metal: add col2im_1d op (f32/f16/bf16) (#25176)
* metal: add col2im_1d op (f32/f16/bf16) Gather kernel mirroring the CPU/CUDA path: each output (t_out, oc) reads its ceil(K/s0) source columns with an F32 accumulator, a single write and no atomics. One thread per output element, 256 per threadgroup. * metal: check dst contiguity and type match in supports_op for COL2IM_1D Align the GGML_OP_COL2IM_1D predicate with the CPU, CUDA, and Vulkan backends: the kernel writes dst with linear indexing and assumes the same type as src0, so supports_op must also require a contiguous dst and op->type == op->src[0]->type. * Update ggml/src/ggml-metal/ggml-metal.metal Co-authored-by: YiChen Lv <63285796+forforever73@users.noreply.github.com> --------- Co-authored-by: YiChen Lv <63285796+forforever73@users.noreply.github.com>
This commit is contained in:
@@ -395,6 +395,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
|
||||
{
|
||||
n_fuse = ggml_metal_op_conv_transpose_2d(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_COL2IM_1D:
|
||||
{
|
||||
n_fuse = ggml_metal_op_col2im_1d(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_CONV_3D:
|
||||
{
|
||||
n_fuse = ggml_metal_op_conv_3d(ctx, idx);
|
||||
@@ -3854,6 +3858,47 @@ int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_col2im_1d(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_library_t lib = ctx->lib;
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
|
||||
const int32_t s0 = ((const int32_t *)(op->op_params))[0];
|
||||
const int32_t OC = ((const int32_t *)(op->op_params))[1];
|
||||
const int32_t p0 = ((const int32_t *)(op->op_params))[2];
|
||||
|
||||
const int32_t K_OC = (int32_t) op->src[0]->ne[0];
|
||||
const int32_t T_in = (int32_t) op->src[0]->ne[1];
|
||||
const int32_t K = K_OC / OC;
|
||||
const int32_t T_out = (int32_t) op->ne[0];
|
||||
|
||||
ggml_metal_kargs_col2im_1d args = {
|
||||
/*.T_in =*/ T_in,
|
||||
/*.T_out =*/ T_out,
|
||||
/*.OC =*/ OC,
|
||||
/*.K =*/ K,
|
||||
/*.K_OC =*/ K_OC,
|
||||
/*.s0 =*/ s0,
|
||||
/*.p0 =*/ p0,
|
||||
};
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_col2im_1d(lib, op);
|
||||
|
||||
const int total = T_out * OC;
|
||||
const int nth = 256;
|
||||
const int ntg = (total + nth - 1) / nth;
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user