diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 585e10d45..356f9abd7 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -962,6 +962,7 @@ struct vk_device_struct { vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32; + vk_pipeline pipeline_cpy_transpose_02_16, pipeline_cpy_transpose_02_32; // [src0 0=fp32,1=fp16][dst] vk_pipeline pipeline_set_rows_i32[2][GGML_TYPE_COUNT]; vk_pipeline pipeline_set_rows_i64[2][GGML_TYPE_COUNT]; @@ -5525,6 +5526,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_32, "cpy_transpose_32", cpy_transpose_32_len, cpy_transpose_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_16, "cpy_transpose_16", cpy_transpose_16_len, cpy_transpose_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_02_32, "cpy_transpose_02_32", cpy_transpose_02_32_len, cpy_transpose_02_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_02_16, "cpy_transpose_02_16", cpy_transpose_02_16_len, cpy_transpose_02_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q1_0], "cpy_f32_q1_0", cpy_f32_q1_0_len, cpy_f32_q1_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q2_0], "cpy_f32_q2_0", cpy_f32_q2_0_len, cpy_f32_q2_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); @@ -8931,6 +8934,18 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const } } + // Same, for a 0<->2 swap: src dim2 is the innermost dimension. + bool transpose02 = dst && !contig && src->nb[2] == ggml_type_size(to) && + ggml_is_contiguous(dst) && ggml_are_same_shape(dst, src); + + if (transpose02 && src->type == to) { + if (ggml_type_size(to) == 4) { + return ctx->device->pipeline_cpy_transpose_02_32; + } else if (ggml_type_size(to) == 2) { + return ctx->device->pipeline_cpy_transpose_02_16; + } + } + if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { if (contig) { return ctx->device->pipeline_contig_cpy_f32_f32; @@ -12192,7 +12207,16 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co elements = { ne, 1, 1 }; } - if (pipeline == ctx->device->pipeline_cpy_transpose_32 || + if (pipeline == ctx->device->pipeline_cpy_transpose_02_32 || + pipeline == ctx->device->pipeline_cpy_transpose_02_16) { + // 32x32 tiles over dims 0 and 2; dim1 and dim3 are the batch + elements[0] = (uint32_t)CEIL_DIV(dst->ne[0], 32); + elements[1] = (uint32_t)CEIL_DIV(dst->ne[2], 32); + elements[2] = (uint32_t)(dst->ne[1]*dst->ne[3]); + elements[0] = std::min(elements[0], ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + elements[1] = std::min(elements[1], ctx->device->properties.limits.maxComputeWorkGroupCount[1]); + elements[2] = std::min(elements[2], ctx->device->properties.limits.maxComputeWorkGroupCount[2]); + } else if (pipeline == ctx->device->pipeline_cpy_transpose_32 || pipeline == ctx->device->pipeline_cpy_transpose_16) { // 32x32 tiles elements[0] = (uint32_t)CEIL_DIV(dst->ne[0], 32); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp b/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp new file mode 100644 index 000000000..5a3d66dab --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp @@ -0,0 +1,61 @@ +#version 450 + +#include "types.glsl" +#include "generic_unary_head.glsl" + +// workgroup does 32x32 tile, but uses 32x8 threads +#define TILE_DIM 32 +layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in; + +// +1 padding avoids shared-memory bank conflicts on the transposed read +shared uint sh[TILE_DIM][TILE_DIM + 1]; + +void iter(uvec3 wg_id) { + const uint tile_i0 = wg_id.x; // tiles dst ne10 (== src ne00) + const uint tile_i2 = wg_id.y; // tiles dst ne12 (== src ne02) + + const uint tid_col = gl_LocalInvocationID.x; + const uint tid_row = gl_LocalInvocationID.y; + + const uint i1 = wg_id.z % p.ne11; + const uint i3 = wg_id.z / p.ne11; + const uint i01 = i1; + const uint i03 = i3; + + [[unroll]] for (uint y = 0; y < 4; ++y) { + const uint i00 = tile_i0 * TILE_DIM + tid_row + 8 * y; + const uint i02 = tile_i2 * TILE_DIM + tid_col; + if (i00 < p.ne00 && i01 < p.ne01 && i02 < p.ne02 && i03 < p.ne03) { + const uint src_idx = i00 * p.nb00 + i01 * p.nb01 + i02 * p.nb02 + i03 * p.nb03; + sh[tid_row + 8 * y][tid_col] = uint(data_a[get_aoffset() + src_idx]); + } + } + + barrier(); + + [[unroll]] for (uint y = 0; y < 4; ++y) { + const uint i0 = tile_i0 * TILE_DIM + tid_col; + const uint i2 = tile_i2 * TILE_DIM + tid_row + 8 * y; + if (i0 < p.ne10 && i1 < p.ne11 && i2 < p.ne12 && i3 < p.ne13) { + const uint dst_idx = i0 * p.nb10 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13; + data_d[get_doffset() + dst_idx] = D_TYPE(sh[tid_col][tid_row + 8 * y]); + } + } +} + +#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) + +void main() { + bool need_barrier = false; + for (uint z = gl_WorkGroupID.z; z < p.ne11 * p.ne13; z += gl_NumWorkGroups.z) { + for (uint y = gl_WorkGroupID.y; y < CEIL_DIV(p.ne12, TILE_DIM); y += gl_NumWorkGroups.y) { + for (uint x = gl_WorkGroupID.x; x < CEIL_DIV(p.ne10, TILE_DIM); x += gl_NumWorkGroups.x) { + if (need_barrier) { + barrier(); + } + need_barrier = true; + iter(uvec3(x, y, z)); + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 6c9f76af1..fbc2ea3ca 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -826,6 +826,8 @@ void process_shaders() { string_to_spv("cpy_transpose_16", "copy_transpose.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); + string_to_spv("cpy_transpose_02_16", "copy_transpose_02.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); + string_to_spv("cpy_transpose_02_32", "copy_transpose_02.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) { string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 219235e2d..52cd747d5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3061,28 +3061,36 @@ struct test_cpy : public test_case { }; // GGML_OP_CONT +// permute = {0, 0, 0, 0} means no permutation: the source is transposed (or +// view-sliced). A non-identity permute applies ggml_permute before ggml_cont. struct test_cont : public test_case { const ggml_type type; const std::array ne; bool use_view_slice; + const std::array permute; std::string vars() override { - return VARS_TO_STR3(type, ne, use_view_slice); + return VARS_TO_STR4(type, ne, use_view_slice, permute); } test_cont(ggml_type type = GGML_TYPE_F32, std::array ne = {10, 10, 10, 1}, - bool use_view_slice = false) - : type(type), ne(ne), use_view_slice(use_view_slice) {} + bool use_view_slice = false, + std::array permute = {0, 0, 0, 0}) + : type(type), ne(ne), use_view_slice(use_view_slice), permute(permute) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * src = ggml_new_tensor(ctx, type, 4, ne.data()); ggml_set_param(src); ggml_set_name(src, "src"); + const bool permuted = permute[0] != 0 || permute[1] != 0 || permute[2] != 0 || permute[3] != 0; ggml_tensor * dst; - if (use_view_slice) { + if (permuted) { + dst = ggml_permute(ctx, src, permute[0], permute[1], permute[2], permute[3]); + ggml_set_name(dst, "src_permuted"); + } else if (use_view_slice) { dst = ggml_view_4d(ctx, src, src->ne[0], 1, src->ne[2], src->ne[3], src->nb[1], src->nb[2], src->nb[3], src->nb[0] * (src->ne[1] - 1)); ggml_set_name(dst, "src_view_slice"); @@ -8905,6 +8913,20 @@ static std::vector> make_test_cases_eval() { } } + for (ggml_type type_dst : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + for (std::array ne : std::initializer_list>{ + {10, 10, 10, 1}, {33, 5, 7, 1}, {64, 3, 65, 1}, {2, 3, 5, 7}, + // large, tile-aligned and tile-unaligned, matching the perf cases + {1024, 64, 64, 1}, {2304, 64, 64, 1}, {1000, 33, 65, 1} }) { + for (std::array perm : std::initializer_list>{ + {2, 1, 0, 3}, // 0<->2 swap + {1, 2, 0, 3}, // 3-cycle + {0, 2, 1, 3} }) { + test_cases.emplace_back(new test_cont(type_dst, ne, false, perm)); + } + } + } + auto add_test_bin_bcast = [&](ggml_type type, std::array ne, std::array nr, bool perm1 = false, bool src_overlap = false) { for (auto op : {ggml_add, ggml_sub, ggml_mul, ggml_div}) { test_cases.emplace_back(new test_bin_bcast(op, type, ne, nr, 1, perm1, src_overlap)); @@ -10056,6 +10078,17 @@ static std::vector> make_test_cases_perf() { } } + // CONT of a 0<->2 permute at DeepSeek-V4 lightning-indexer shapes: + // indexer_kq is [n_kv, n_tokens, n_head=64] and gets ggml_cont(ggml_permute(.., 2,1,0,3)). + for (int64_t n_kv : { 1024, 1280, 2048, 2304 }) { + test_cases.emplace_back(new test_cont( + GGML_TYPE_F32, {n_kv, 64, 64, 1}, false, {2, 1, 0, 3})); + } + for (int64_t n_kv : { 2048, 2304 }) { + test_cases.emplace_back(new test_cont( + GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3})); + } + // Conv2d: K=CRS=NPQ=4096 matmul performance uint32_t iwh_idx = 0; uint32_t kwh_idx = 1;