ggml : fix conv_transpose_2d for multiple batches (#26132)
* ggml : fix conv_transpose_2d for multiple batches ggml_compute_forward_conv_transpose_2d_impl only computed the first batch (ne[3] of the destination); every batch after the first was left as zero. Both the src1 permutation and the main compute loop now iterate over the batch dimension, and the work buffer size in ggml_graph_plan is scaled by the src1 batch count so the extra permuted batches fit. A multi-batch test case is added to test-backend-ops. Fixes ggml-org/ggml#1448 * metal : fix conv_transpose_2d for multiple batches The kernel only computed batch 0 of the input (src1->ne[3]); every output batch after the first was left as zero, so multi-batch conv_transpose_2d results diverged from the CPU reference. The grid now covers all batches (OW x OH x OC x N), the kernel decodes the batch from the grid z coordinate and offsets both the input and destination indices accordingly. nb3 is passed in the kernel args. Assisted-by: pi:llama.cpp/Qwen3.8-27B --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
co-authored by
Georgi Gerganov
parent
90c26fcd4b
commit
a43c3986b4
@@ -660,6 +660,7 @@ typedef struct {
|
||||
uint64_t nb0;
|
||||
uint64_t nb1;
|
||||
uint64_t nb2;
|
||||
uint64_t nb3;
|
||||
} ggml_metal_kargs_conv_transpose_2d;
|
||||
|
||||
typedef struct {
|
||||
|
||||
@@ -4645,6 +4645,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
|
||||
const int32_t OW = op->ne[0];
|
||||
const int32_t OH = op->ne[1];
|
||||
const int32_t OC = op->ne[2];
|
||||
const int32_t N = op->src[1]->ne[3];
|
||||
|
||||
ggml_metal_kargs_conv_transpose_2d args = {
|
||||
/*.IC =*/ IC,
|
||||
@@ -4657,6 +4658,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
|
||||
/*.nb0 =*/ nb0,
|
||||
/*.nb1 =*/ nb1,
|
||||
/*.nb2 =*/ nb2,
|
||||
/*.nb3 =*/ nb3,
|
||||
};
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op);
|
||||
@@ -4671,7 +4673,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
|
||||
const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16);
|
||||
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1);
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC * N, KW, KH, 1);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
@@ -366,7 +366,8 @@ kernel void kernel_conv_transpose_2d(
|
||||
|
||||
const int64_t out_x = tgpig[0];
|
||||
const int64_t out_y = tgpig[1];
|
||||
const int64_t out_c = tgpig[2];
|
||||
const int64_t batch = tgpig[2] / args.OC;
|
||||
const int64_t out_c = tgpig[2] % args.OC;
|
||||
|
||||
const int64_t kw = tpitg[0];
|
||||
const int64_t kh = tpitg[1];
|
||||
@@ -390,7 +391,7 @@ kernel void kernel_conv_transpose_2d(
|
||||
|
||||
if (in_x >= args.IW) continue;
|
||||
|
||||
const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x;
|
||||
const int64_t input_idx = (args.IW * args.IH) * (args.IC * batch + in_c) + (args.IW) * in_y + in_x;
|
||||
const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw;
|
||||
|
||||
v += (float)src0[kernel_idx] * src1[input_idx];
|
||||
@@ -408,7 +409,7 @@ kernel void kernel_conv_transpose_2d(
|
||||
total += shared_sum[i];
|
||||
}
|
||||
|
||||
device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
|
||||
device float * dst_ptr = (device float *) (dst + batch*args.nb3 + out_c*args.nb2 + out_y * args.nb1 + out_x*args.nb0);
|
||||
dst_ptr[0] = total;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user