diff --git a/ggml/src/ggml-metal/kernels/conv.metal b/ggml/src/ggml-metal/kernels/conv.metal index 33014e0ed0..e4ef3d8f93 100644 --- a/ggml/src/ggml-metal/kernels/conv.metal +++ b/ggml/src/ggml-metal/kernels/conv.metal @@ -118,6 +118,48 @@ kernel void kernel_im2col_ext( template [[host_name("kernel_im2col_ext_f32")]] kernel im2col_ext_t kernel_im2col_ext; template [[host_name("kernel_im2col_ext_f16")]] kernel im2col_ext_t kernel_im2col_ext; +template +kernel void kernel_col2im_1d( + constant ggml_metal_kargs_col2im_1d & args, + device const T * col, + device T * dst, + uint tgpig [[threadgroup_position_in_grid]], + uint tpitg [[thread_position_in_threadgroup]], + uint ntg [[threads_per_threadgroup]]) { + + const int idx = tgpig * ntg + tpitg; + if (idx >= args.T_out * args.OC) { + return; + } + + const int t_out = idx % args.T_out; + const int oc = idx / args.T_out; + const int t_abs = t_out + args.p0; // absolute position in uncropped signal + + int t_in_min = (t_abs - args.K + args.s0) / args.s0; // ceil((t_abs - K + 1) / s0) + if (t_in_min < 0) { + t_in_min = 0; + } + int t_in_max = t_abs / args.s0; + if (t_in_max >= args.T_in) { + t_in_max = args.T_in - 1; + } + + float sum = 0.0f; + for (int t_in = t_in_min; t_in <= t_in_max; t_in++) { + const int k = t_abs - t_in * args.s0; + sum += float(col[(oc * args.K + k) + t_in * args.K_OC]); + } + + dst[t_out + oc * args.T_out] = T(sum); +} + +template [[host_name("kernel_col2im_1d_f32")]] kernel void kernel_col2im_1d(constant ggml_metal_kargs_col2im_1d &, device const float *, device float *, uint, uint, uint); +template [[host_name("kernel_col2im_1d_f16")]] kernel void kernel_col2im_1d(constant ggml_metal_kargs_col2im_1d &, device const half *, device half *, uint, uint, uint); +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d(constant ggml_metal_kargs_col2im_1d &, device const bfloat *, device bfloat *, uint, uint, uint); +#endif + template kernel void kernel_conv_2d( constant ggml_metal_kargs_conv_2d & args,