From 65dd9133baf3bdebe783af71c67fc4e24a51d55a Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 10 Jul 2026 14:51:31 +0300 Subject: [PATCH] metal : add CONV_2D_DW (depthwise convolution) support (#21565) --- ggml/src/ggml-metal/kernels/conv.metal | 196 +++++++++++++++++++++++++ 1 file changed, 196 insertions(+) diff --git a/ggml/src/ggml-metal/kernels/conv.metal b/ggml/src/ggml-metal/kernels/conv.metal index e4ef3d8f93..5685b5cd49 100644 --- a/ggml/src/ggml-metal/kernels/conv.metal +++ b/ggml/src/ggml-metal/kernels/conv.metal @@ -435,6 +435,202 @@ kernel void kernel_conv_transpose_2d( uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]); +// grid: x = C tile, y = OH, z = OW * N (for channel-contiguous layouts) +template +kernel void kernel_conv_2d_dw_tiled( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]) { + + const int32_t c = (int32_t)(tgpig.x * ntg.x + tpitg.x); + if (c >= args.C) { + return; + } + + const int32_t oh = tgpig.y; + const int32_t own = tgpig.z; + const int32_t ow = own % args.OW; + const int32_t n = own / args.OW; + + const int32_t base_y = oh*args.s1 - args.p1; + + int32_t ky_start = 0; + if (base_y < 0) { + ky_start = (-base_y + args.d1 - 1)/args.d1; + } + int32_t ky_end = args.KH; + const int32_t y_max = args.IH - 1 - base_y; + if (y_max < 0) { + ky_end = ky_start; + } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) { + ky_end = min(ky_end, y_max/args.d1 + 1); + } + + const int32_t base_x = ow*args.s0 - args.p0; + + int32_t kx_start = 0; + if (base_x < 0) { + kx_start = (-base_x + args.d0 - 1)/args.d0; + } + int32_t kx_end = args.KW; + const int32_t x_max = args.IW - 1 - base_x; + if (x_max < 0) { + kx_end = kx_start; + } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) { + kx_end = min(kx_end, x_max/args.d0 + 1); + } + + float acc = 0.0f; + + if (ky_start < ky_end && kx_start < kx_end) { + const uint64_t w_base = (uint64_t) c * args.nb02; + const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12; + + for (int32_t ky = ky_start; ky < ky_end; ++ky) { + const int32_t iy = base_y + ky*args.d1; + const uint64_t src_row = src_base + (uint64_t) iy * args.nb11; + const uint64_t w_row = w_base + (uint64_t) ky * args.nb01; + + for (int32_t kx = kx_start; kx < kx_end; ++kx) { + const int32_t ix = base_x + kx*args.d0; + const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10); + const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00)); + acc += x * w; + } + } + } + + const uint64_t dst_offs = + (uint64_t) n * args.nb3 + + (uint64_t) c * args.nb2 + + (uint64_t) oh * args.nb1 + + (uint64_t) ow * args.nb0; + + *(device float *)(dst + dst_offs) = acc; +} + +// grid: x = OW tile, y = OH, z = C * N (for spatially-contiguous layouts) +template +kernel void kernel_conv_2d_dw( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]) { + + const int32_t oh = tgpig.y; + const int32_t cn = tgpig.z; + const int32_t c = cn % args.C; + const int32_t n = cn / args.C; + + const int32_t base_y = oh*args.s1 - args.p1; + + int32_t ky_start = 0; + if (base_y < 0) { + ky_start = (-base_y + args.d1 - 1)/args.d1; + } + int32_t ky_end = args.KH; + const int32_t y_max = args.IH - 1 - base_y; + if (y_max < 0) { + ky_end = ky_start; + } else if (base_y + (args.KH - 1)*args.d1 >= args.IH) { + ky_end = min(ky_end, y_max/args.d1 + 1); + } + + const uint64_t w_base = (uint64_t) c * args.nb02; + const uint64_t src_base = (uint64_t) n * args.nb13 + (uint64_t) c * args.nb12; + + const int32_t ow = (int32_t)(tgpig.x * ntg.x + tpitg.x); + if (ow >= args.OW) { + return; + } + + float acc = 0.0f; + + const int32_t base_x = ow*args.s0 - args.p0; + + int32_t kx_start = 0; + if (base_x < 0) { + kx_start = (-base_x + args.d0 - 1)/args.d0; + } + int32_t kx_end = args.KW; + const int32_t x_max = args.IW - 1 - base_x; + if (x_max < 0) { + kx_end = kx_start; + } else if (base_x + (args.KW - 1)*args.d0 >= args.IW) { + kx_end = min(kx_end, x_max/args.d0 + 1); + } + + if (ky_start < ky_end && kx_start < kx_end) { + for (int32_t ky = ky_start; ky < ky_end; ++ky) { + const int32_t iy = base_y + ky*args.d1; + const uint64_t src_row = src_base + (uint64_t) iy * args.nb11; + const uint64_t w_row = w_base + (uint64_t) ky * args.nb01; + + for (int32_t kx = kx_start; kx < kx_end; ++kx) { + const int32_t ix = base_x + kx*args.d0; + const float x = *(device const float *)(src + src_row + (uint64_t) ix * args.nb10); + const float w = (float)(*(device const TK *)(weights + w_row + (uint64_t) kx * args.nb00)); + acc += x * w; + } + } + } + + const uint64_t dst_offs = + (uint64_t) n * args.nb3 + + (uint64_t) c * args.nb2 + + (uint64_t) oh * args.nb1 + + (uint64_t) ow * args.nb0; + + *(device float *)(dst + dst_offs) = acc; +} + +template [[host_name("kernel_conv_2d_dw_f32_f32")]] +kernel void kernel_conv_2d_dw( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + +template [[host_name("kernel_conv_2d_dw_f16_f32")]] +kernel void kernel_conv_2d_dw( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + +template [[host_name("kernel_conv_2d_dw_tiled_f32_f32")]] +kernel void kernel_conv_2d_dw_tiled( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + +template [[host_name("kernel_conv_2d_dw_tiled_f16_f32")]] +kernel void kernel_conv_2d_dw_tiled( + constant ggml_metal_kargs_conv_2d_dw & args, + device const char * weights, + device const char * src, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 ntg[[threads_per_threadgroup]]); + template kernel void kernel_conv_3d( constant ggml_metal_kargs_conv_3d & args,