metal: add CONV_3D (#19927)
* Apply suggestions from code review Co-authored-by: Georgi Gerganov <ggerganov@gmail.com> * metal:add conv_3d backend Rebased with master and resolved conflicts. * Resolved issues related to changes in variable names * kernel void kernel_upscale_bilinear_f32 was missing in my branch, added back, should pass all tests now --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
co-authored by
Georgi Gerganov
parent
7a0b6a635e
commit
177c75852a
@@ -394,6 +394,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_CONV_3D:
|
||||
{
|
||||
n_fuse = ggml_metal_op_conv_3d(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_UPSCALE:
|
||||
{
|
||||
n_fuse = ggml_metal_op_upscale(ctx, idx);
|
||||
@@ -3697,6 +3701,77 @@ int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_conv_3d(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;
|
||||
|
||||
// 1. Extract standard dimensions and byte strides
|
||||
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
|
||||
GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb);
|
||||
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
|
||||
|
||||
// 2. Extract hyperparams from op_params
|
||||
const int32_t s0 = ((const int32_t *)(op->op_params))[0];
|
||||
const int32_t s1 = ((const int32_t *)(op->op_params))[1];
|
||||
const int32_t s2 = ((const int32_t *)(op->op_params))[2];
|
||||
const int32_t p0 = ((const int32_t *)(op->op_params))[3];
|
||||
const int32_t p1 = ((const int32_t *)(op->op_params))[4];
|
||||
const int32_t p2 = ((const int32_t *)(op->op_params))[5];
|
||||
const int32_t d0 = ((const int32_t *)(op->op_params))[6];
|
||||
const int32_t d1 = ((const int32_t *)(op->op_params))[7];
|
||||
const int32_t d2 = ((const int32_t *)(op->op_params))[8];
|
||||
const int32_t IC = ((const int32_t *)(op->op_params))[9];
|
||||
const int32_t N = ((const int32_t *)(op->op_params))[10];
|
||||
const int32_t OC = ((const int32_t *)(op->op_params))[11];
|
||||
|
||||
// 3. Build the parameter struct using the macro-generated variables
|
||||
ggml_metal_kargs_conv_3d args = {
|
||||
/*.IW =*/ (int32_t)op->src[1]->ne[0],
|
||||
/*.IH =*/ (int32_t)op->src[1]->ne[1],
|
||||
/*.ID =*/ (int32_t)op->src[1]->ne[2],
|
||||
/*.OW =*/ (int32_t)op->ne[0],
|
||||
/*.OH =*/ (int32_t)op->ne[1],
|
||||
/*.OD =*/ (int32_t)op->ne[2],
|
||||
/*.KW =*/ (int32_t)op->src[0]->ne[0],
|
||||
/*.KH =*/ (int32_t)op->src[0]->ne[1],
|
||||
/*.KD =*/ (int32_t)op->src[0]->ne[2],
|
||||
s0, s1, s2,
|
||||
p0, p1, p2,
|
||||
d0, d1, d2,
|
||||
IC, N, OC,
|
||||
nb00, nb01, nb02, nb03, // Weight strides
|
||||
nb10, nb11, nb12, nb13, // Input strides
|
||||
nb0, nb1, nb2, nb3 // Output strides
|
||||
};
|
||||
|
||||
// 4. Fetch the JIT pipeline
|
||||
auto pipeline = ggml_metal_library_get_pipeline_conv_3d(lib, op);
|
||||
|
||||
// 5. Grid mapping
|
||||
int nth0 = 32; // Standard SIMD width for Apple Silicon
|
||||
int nth1 = 1;
|
||||
int nth2 = 1;
|
||||
|
||||
int64_t spatial_volume = args.OW * args.OH * args.OD;
|
||||
|
||||
int ntg0 = (spatial_volume + nth0 - 1) / nth0;
|
||||
int ntg1 = args.OC;
|
||||
int ntg2 = args.N;
|
||||
|
||||
// 6. Bind and Dispatch via the ggml C wrapper
|
||||
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->src[1]), 2);
|
||||
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, ntg0, ntg1, ntg2, nth0, nth1, nth2);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user