diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 4cb246154..766158b04 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -12520,9 +12520,108 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc ggml_vk_op_f32(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_OPT_STEP_SGD, { (uint32_t)n, 0, 0.0f, 0.0f, 0.0f, 0.0f }); } +// Fast path for concat along dim 0 where one source is stored "transposed" +// (nb[1] == type_size, dim1 innermost) and the other source + dst are +// contiguous along dim 0. The generic concat shader reads the transposed +// source with a catastrophic uncoalesced stride; here we instead copy the +// contiguous source with copy.comp and transpose the other source into the +// matching dst sub-region with the tiled copy_transpose shader (shared-memory +// transpose, coalesced read+write). Mirrors the precedent in +// ggml_vk_cpy_to_contiguous: direct dispatch with custom push constants. +static void ggml_vk_concat_transpose_fastpath(ggml_backend_vk_context * ctx, vk_context& subctx, + const ggml_tensor * ctg, const ggml_tensor * trp, + ggml_tensor * dst, uint32_t off_ctg, uint32_t off_trp) { + const uint32_t ts = ggml_type_size(dst->type); + + vk_pipeline pipeline_cpy = (ts == 4) ? ctx->device->pipeline_cpy_f32_f32 + : ctx->device->pipeline_cpy_f16_f16; + vk_pipeline pipeline_trp = (ts == 4) ? ctx->device->pipeline_cpy_transpose_32 + : ctx->device->pipeline_cpy_transpose_16; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline_cpy, 1); + ggml_pipeline_request_descriptor_sets(ctx, pipeline_trp, 1); + + vk_subbuffer ctg_buf = ggml_vk_tensor_subbuffer(ctx, ctg, true); + vk_subbuffer trp_buf = ggml_vk_tensor_subbuffer(ctx, trp, true); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + const uint32_t a_misalign_ctg = get_misalign_bytes(ctx, ctg) / ts; + const uint32_t a_misalign_trp = get_misalign_bytes(ctx, trp) / ts; + const uint32_t d_misalign = get_misalign_bytes(ctx, dst) / ts; + + // Dispatch A: contiguous copy of `ctg` into dst[off_ctg : off_ctg + ctg->ne[0], :] + if (ctg->ne[0] > 0) { + const uint32_t ne_ctg = (uint32_t) ggml_nelements(ctg); + vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(ctg, dst, ne_ctg); + pc.ne10 = (uint32_t) ctg->ne[0]; // only ctg's columns in dst + pc.misalign_offsets = (a_misalign_ctg << 16) | (d_misalign + off_ctg); + init_pushconst_fastdiv(pc); + std::array el = ne_ctg > 262144 ? std::array{512, 512, CEIL_DIV(ne_ctg, 262144)} + : ne_ctg > 512 ? std::array{512, CEIL_DIV(ne_ctg, 512), 1} + : std::array{ne_ctg, 1, 1}; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_cpy, { ctg_buf, dst_buf }, pc, el); + } + + // Dispatch B: tiled transpose of `trp` into dst[off_trp : off_trp + trp->ne[0], :] + if (trp->ne[0] > 0) { + vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(trp, dst, ggml_nelements(trp)); + pc.ne10 = (uint32_t) trp->ne[0]; // dst bound = trp columns (NOT dst->ne[0]) + pc.misalign_offsets = (a_misalign_trp << 16) | (d_misalign + off_trp); + init_pushconst_fastdiv(pc); + std::array el = { + (uint32_t) CEIL_DIV(trp->ne[0], 32), + (uint32_t) CEIL_DIV(trp->ne[1], 32), + (uint32_t) (trp->ne[2] * trp->ne[3]), + }; + el[0] = std::min(el[0], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + el[1] = std::min(el[1], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[1]); + el[2] = std::min(el[2], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[2]); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_trp, { trp_buf, dst_buf }, pc, el); + } + + ggml_vk_sync_buffers(ctx, subctx); +} + static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { int * op_params = (int *)dst->op_params; + // Fast path: concat along dim 0 with one source "transposed" (nb[1]==type_size, + // dim1 innermost) and the other source + dst contiguous along dim 0. The generic + // shader reads the transposed source uncoalesced (~5.6ms on RX 580 vs ~130us for + // a coalesced copy); here we instead use a tiled shared-memory transpose. + auto src_transposed_2d = [&](const ggml_tensor * s) { + const uint32_t ts = ggml_type_size(s->type); + return s->nb[1] == ts // dim1 innermost + && s->nb[0] == s->ne[1] * ts // consistent 2D transpose + && s->ne[2] == 1 && s->ne[3] == 1; + }; + const uint32_t dst_ts = ggml_type_size(dst->type); + const bool dim0_ok = (op_params[0] == 0) && (dst->nb[0] == dst_ts); + const bool types_ok = (dst_ts == 4 || dst_ts == 2) + && (src0->type == src1->type && src0->type == dst->type) + && (ggml_blck_size(dst->type) == 1); + const bool shapes_ok = (src0->ne[1] == src1->ne[1] && src1->ne[1] == dst->ne[1]) + && (src0->ne[2] == src1->ne[2] && src0->ne[2] == dst->ne[2]) + && (src0->ne[3] == src1->ne[3] && src0->ne[3] == dst->ne[3]) + && (src0->ne[0] + src1->ne[0] == dst->ne[0]); + // doffset is 16-bit for unary shaders; guard against overflow + const bool off_fits = ((get_misalign_bytes(ctx, dst)/dst_ts + dst->ne[0])) < 0xFFFFu; + + const bool s1_trans = dim0_ok && types_ok && shapes_ok && off_fits + && src_transposed_2d(src1) && ggml_is_contiguous(src0); + const bool s0_trans = dim0_ok && types_ok && shapes_ok && off_fits + && src_transposed_2d(src0) && ggml_is_contiguous(src1); + + if (s1_trans || s0_trans) { + const ggml_tensor * ctg = s1_trans ? src0 : src1; // contiguous source + const ggml_tensor * trp = s1_trans ? src1 : src0; // transposed source + // dst offsets (in elements) where each source region begins + const uint32_t off_ctg = s1_trans ? 0u : (uint32_t) src1->ne[0]; + const uint32_t off_trp = s1_trans ? (uint32_t) src0->ne[0] : 0u; + ggml_vk_concat_transpose_fastpath(ctx, subctx, ctg, trp, dst, off_ctg, off_trp); + return; + } + const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); const uint32_t dst_type_size = ggml_type_size(dst->type);