Skip to content

Commit 0cfb191

Browse files
bghirabghiraggerganov
authored
metal: accelerated conv2d (ggml-org#17175)
* metal: accelerated conv2d * cont : cleanup --------- Co-authored-by: bghira <[email protected]> Co-authored-by: Georgi Gerganov <[email protected]>
1 parent 2776db6 commit 0cfb191

File tree

7 files changed

+258
-5
lines changed

7 files changed

+258
-5
lines changed

‎ggml/src/ggml-metal/ggml-metal-device.cpp‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1438,6 +1438,30 @@ ggml_metal_pipeline_t ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_met
14381438
return res;
14391439
}
14401440

1441+
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_conv_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
1442+
assert(op->op == GGML_OP_CONV_2D);
1443+
1444+
GGML_ASSERT(ggml_is_contiguous(op->src[0]));
1445+
GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
1446+
GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
1447+
GGML_ASSERT(op->type == GGML_TYPE_F32);
1448+
1449+
char base[256];
1450+
char name[256];
1451+
1452+
snprintf(base, 256, "kernel_conv_2d_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type));
1453+
snprintf(name, 256, "%s", base);
1454+
1455+
ggml_metal_pipeline_t res = ggml_metal_library_get_pipeline(lib, name);
1456+
if (res) {
1457+
return res;
1458+
}
1459+
1460+
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1461+
1462+
return res;
1463+
}
1464+
14411465
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_upscale(ggml_metal_library_t lib, const ggml_tensor * op) {
14421466
assert(op->op == GGML_OP_UPSCALE);
14431467

‎ggml/src/ggml-metal/ggml-metal-device.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ ggml_metal_pipeline_t ggml_metal_library_get_pipeline_rope (ggml_me
133133
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_im2col (ggml_metal_library_t lib, const struct ggml_tensor * op);
134134
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
135135
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
136+
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_conv_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
136137
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_upscale (ggml_metal_library_t lib, const struct ggml_tensor * op);
137138
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_pad (ggml_metal_library_t lib, const struct ggml_tensor * op);
138139
ggml_metal_pipeline_t ggml_metal_library_get_pipeline_pad_reflect_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);

‎ggml/src/ggml-metal/ggml-metal-device.m‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -885,6 +885,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
885885
return true;
886886
case GGML_OP_IM2COL:
887887
return ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32 && (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32);
888+
case GGML_OP_CONV_2D:
889+
return ggml_is_contiguous(op->src[0]) &&
890+
op->src[1]->type == GGML_TYPE_F32 &&
891+
op->type == GGML_TYPE_F32 &&
892+
(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
888893
case GGML_OP_POOL_1D:
889894
return false;
890895
case GGML_OP_UPSCALE:

‎ggml/src/ggml-metal/ggml-metal-impl.h‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -528,6 +528,36 @@ typedef struct {
528528
uint64_t nb2;
529529
} ggml_metal_kargs_conv_transpose_2d;
530530

531+
typedef struct {
532+
uint64_t nb00;
533+
uint64_t nb01;
534+
uint64_t nb02;
535+
uint64_t nb03;
536+
uint64_t nb10;
537+
uint64_t nb11;
538+
uint64_t nb12;
539+
uint64_t nb13;
540+
uint64_t nb0;
541+
uint64_t nb1;
542+
uint64_t nb2;
543+
uint64_t nb3;
544+
int32_t IW;
545+
int32_t IH;
546+
int32_t KW;
547+
int32_t KH;
548+
int32_t IC;
549+
int32_t OC;
550+
int32_t OW;
551+
int32_t OH;
552+
int32_t N;
553+
int32_t s0;
554+
int32_t s1;
555+
int32_t p0;
556+
int32_t p1;
557+
int32_t d0;
558+
int32_t d1;
559+
} ggml_metal_kargs_conv_2d;
560+
531561
typedef struct {
532562
uint64_t ofs0;
533563
uint64_t ofs1;

‎ggml/src/ggml-metal/ggml-metal-ops.cpp‎

Lines changed: 83 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
#include <cassert>
1212
#include <algorithm>
13+
#include <limits>
1314

1415
static ggml_metal_buffer_id ggml_metal_get_buffer_id(const ggml_tensor * t) {
1516
if (!t) {
@@ -364,6 +365,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
364365
{
365366
n_fuse = ggml_metal_op_im2col(ctx, idx);
366367
} break;
368+
case GGML_OP_CONV_2D:
369+
{
370+
n_fuse = ggml_metal_op_conv_2d(ctx, idx);
371+
} break;
367372
case GGML_OP_CONV_TRANSPOSE_1D:
368373
{
369374
n_fuse = ggml_metal_op_conv_transpose_1d(ctx, idx);
@@ -1036,11 +1041,6 @@ int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) {
10361041

10371042
nth = std::min(nth, nk0);
10381043

1039-
if (nth*nrptg > ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) {
1040-
nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
1041-
nrptg = 1;
1042-
}
1043-
10441044
ggml_metal_kargs_set_rows args = {
10451045
/*.nk0 =*/ nk0,
10461046
/*.ne01 =*/ ne01,
@@ -3082,6 +3082,84 @@ int ggml_metal_op_im2col(ggml_metal_op_t ctx, int idx) {
30823082
return 1;
30833083
}
30843084

3085+
int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) {
3086+
ggml_tensor * op = ctx->node(idx);
3087+
3088+
ggml_metal_library_t lib = ctx->lib;
3089+
ggml_metal_encoder_t enc = ctx->enc;
3090+
3091+
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
3092+
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
3093+
GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
3094+
GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb);
3095+
GGML_TENSOR_LOCALS( int32_t, ne, op, ne);
3096+
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
3097+
3098+
GGML_ASSERT(ggml_is_contiguous(op->src[0]));
3099+
GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
3100+
GGML_ASSERT(op->type == GGML_TYPE_F32);
3101+
GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
3102+
3103+
const int32_t s0 = ((const int32_t *) op->op_params)[0];
3104+
const int32_t s1 = ((const int32_t *) op->op_params)[1];
3105+
const int32_t p0 = ((const int32_t *) op->op_params)[2];
3106+
const int32_t p1 = ((const int32_t *) op->op_params)[3];
3107+
const int32_t d0 = ((const int32_t *) op->op_params)[4];
3108+
const int32_t d1 = ((const int32_t *) op->op_params)[5];
3109+
3110+
ggml_metal_kargs_conv_2d args = {
3111+
/*.nb00 =*/ nb00,
3112+
/*.nb01 =*/ nb01,
3113+
/*.nb02 =*/ nb02,
3114+
/*.nb03 =*/ nb03,
3115+
/*.nb10 =*/ nb10,
3116+
/*.nb11 =*/ nb11,
3117+
/*.nb12 =*/ nb12,
3118+
/*.nb13 =*/ nb13,
3119+
/*.nb0 =*/ nb0,
3120+
/*.nb1 =*/ nb1,
3121+
/*.nb2 =*/ nb2,
3122+
/*.nb3 =*/ nb3,
3123+
/*.IW =*/ ne10,
3124+
/*.IH =*/ ne11,
3125+
/*.KW =*/ ne00,
3126+
/*.KH =*/ ne01,
3127+
/*.IC =*/ ne02,
3128+
/*.OC =*/ ne03,
3129+
/*.OW =*/ ne0,
3130+
/*.OH =*/ ne1,
3131+
/*.N =*/ ne3,
3132+
/*.s0 =*/ s0,
3133+
/*.s1 =*/ s1,
3134+
/*.p0 =*/ p0,
3135+
/*.p1 =*/ p1,
3136+
/*.d0 =*/ d0,
3137+
/*.d1 =*/ d1,
3138+
};
3139+
3140+
ggml_metal_pipeline_t pipeline = ggml_metal_library_get_pipeline_conv_2d(lib, op);
3141+
3142+
int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline);
3143+
nth = std::min(nth, 256);
3144+
nth = std::max(nth, 1);
3145+
3146+
const uint64_t n_out = ggml_nelements(op);
3147+
3148+
uint64_t tg = (n_out + nth - 1)/nth;
3149+
tg = std::max<uint64_t>(tg, 1);
3150+
tg = std::min<uint64_t>(tg, (uint64_t) std::numeric_limits<int>::max());
3151+
3152+
ggml_metal_encoder_set_pipeline(enc, pipeline);
3153+
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
3154+
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
3155+
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2);
3156+
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3);
3157+
3158+
ggml_metal_encoder_dispatch_threadgroups(enc, tg, 1, 1, nth, 1, 1);
3159+
3160+
return 1;
3161+
}
3162+
30853163
int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) {
30863164
ggml_tensor * op = ctx->node(idx);
30873165

‎ggml/src/ggml-metal/ggml-metal-ops.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ int ggml_metal_op_group_norm (ggml_metal_op_t ctx, int idx);
7070
int ggml_metal_op_norm (ggml_metal_op_t ctx, int idx);
7171
int ggml_metal_op_rope (ggml_metal_op_t ctx, int idx);
7272
int ggml_metal_op_im2col (ggml_metal_op_t ctx, int idx);
73+
int ggml_metal_op_conv_2d (ggml_metal_op_t ctx, int idx);
7374
int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx);
7475
int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx);
7576
int ggml_metal_op_upscale (ggml_metal_op_t ctx, int idx);

‎ggml/src/ggml-metal/ggml-metal.metal‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4146,6 +4146,120 @@ template [[host_name("kernel_im2col_f16")]] kernel im2col_t kernel_im2col<half>;
41464146
//template [[host_name("kernel_im2col_ext_f32")]] kernel im2col_ext_t kernel_im2col_ext<float>;
41474147
//template [[host_name("kernel_im2col_ext_f16")]] kernel im2col_ext_t kernel_im2col_ext<half>;
41484148

4149+
template <typename TK>
4150+
kernel void kernel_conv_2d(
4151+
constant ggml_metal_kargs_conv_2d & args,
4152+
device const char * weights,
4153+
device const char * src,
4154+
device char * dst,
4155+
uint3 tgpig[[threadgroup_position_in_grid]],
4156+
uint3 tgpg[[threadgroups_per_grid]],
4157+
uint3 tpitg[[thread_position_in_threadgroup]],
4158+
uint3 ntg[[threads_per_threadgroup]]) {
4159+
4160+
const uint threads_per_tg = ntg.x * ntg.y * ntg.z;
4161+
const uint tg_index = (tgpig.z * tgpg.y + tgpig.y) * tgpg.x + tgpig.x;
4162+
const uint local_thread = tpitg.z * (ntg.x * ntg.y) + tpitg.y * ntg.x + tpitg.x;
4163+
const uint thread_index = tg_index * threads_per_tg + local_thread;
4164+
const uint64_t total_threads = (uint64_t) threads_per_tg * tgpg.x * tgpg.y * tgpg.z;
4165+
const uint64_t total_outputs = (uint64_t) args.N * args.OC * args.OH * args.OW;
4166+
4167+
for (uint64_t index = thread_index; index < total_outputs; index += total_threads) {
4168+
uint64_t tmp = index;
4169+
4170+
const int32_t ow = tmp % args.OW; tmp /= args.OW;
4171+
const int32_t oh = tmp % args.OH; tmp /= args.OH;
4172+
const int32_t oc = tmp % args.OC; tmp /= args.OC;
4173+
const int32_t n = tmp;
4174+
4175+
float acc = 0.0f;
4176+
4177+
const int32_t base_x = ow*args.s0 - args.p0;
4178+
const int32_t base_y = oh*args.s1 - args.p1;
4179+
4180+
int32_t ky_start = 0;
4181+
if (base_y < 0) {
4182+
ky_start = (-base_y + args.d1 - 1)/args.d1;
4183+
}
4184+
int32_t ky_end = args.KH;
4185+
const int32_t y_max = args.IH - 1 - base_y;
4186+
if (y_max < 0) {
4187+
ky_end = ky_start;
4188+
} else if (base_y + (args.KH - 1)*args.d1 >= args.IH) {
4189+
ky_end = min(ky_end, y_max/args.d1 + 1);
4190+
}
4191+
4192+
int32_t kx_start = 0;
4193+
if (base_x < 0) {
4194+
kx_start = (-base_x + args.d0 - 1)/args.d0;
4195+
}
4196+
int32_t kx_end = args.KW;
4197+
const int32_t x_max = args.IW - 1 - base_x;
4198+
if (x_max < 0) {
4199+
kx_end = kx_start;
4200+
} else if (base_x + (args.KW - 1)*args.d0 >= args.IW) {
4201+
kx_end = min(kx_end, x_max/args.d0 + 1);
4202+
}
4203+
4204+
if (ky_start < ky_end && kx_start < kx_end) {
4205+
const uint64_t src_base_n = (uint64_t) n * args.nb13;
4206+
const uint64_t w_base_oc = (uint64_t) oc * args.nb03;
4207+
4208+
for (int32_t ic = 0; ic < args.IC; ++ic) {
4209+
const uint64_t src_base_nc = src_base_n + (uint64_t) ic * args.nb12;
4210+
const uint64_t w_base_ocic = w_base_oc + (uint64_t) ic * args.nb02;
4211+
4212+
for (int32_t ky = ky_start; ky < ky_end; ++ky) {
4213+
const int32_t iy = base_y + ky*args.d1;
4214+
const uint64_t src_base_row = src_base_nc + (uint64_t) iy * args.nb11;
4215+
const uint64_t w_base_row = w_base_ocic + (uint64_t) ky * args.nb01;
4216+
4217+
for (int32_t kx = kx_start; kx < kx_end; ++kx) {
4218+
const int32_t ix = base_x + kx*args.d0;
4219+
const uint64_t src_offs = src_base_row + (uint64_t) ix * args.nb10;
4220+
const uint64_t w_offs = w_base_row + (uint64_t) kx * args.nb00;
4221+
4222+
const float x = *(device const float *)(src + src_offs);
4223+
const float w = (float) (*(device const TK *)(weights + w_offs));
4224+
4225+
acc += x * w;
4226+
}
4227+
}
4228+
}
4229+
}
4230+
4231+
const uint64_t dst_offs =
4232+
(uint64_t) n * args.nb3 +
4233+
(uint64_t) oc * args.nb2 +
4234+
(uint64_t) oh * args.nb1 +
4235+
(uint64_t) ow * args.nb0;
4236+
4237+
*(device float *)(dst + dst_offs) = acc;
4238+
}
4239+
}
4240+
4241+
template [[host_name("kernel_conv_2d_f32_f32")]]
4242+
kernel void kernel_conv_2d<float>(
4243+
constant ggml_metal_kargs_conv_2d & args,
4244+
device const char * weights,
4245+
device const char * src,
4246+
device char * dst,
4247+
uint3 tgpig[[threadgroup_position_in_grid]],
4248+
uint3 tgpg[[threadgroups_per_grid]],
4249+
uint3 tpitg[[thread_position_in_threadgroup]],
4250+
uint3 ntg[[threads_per_threadgroup]]);
4251+
4252+
template [[host_name("kernel_conv_2d_f16_f32")]]
4253+
kernel void kernel_conv_2d<half>(
4254+
constant ggml_metal_kargs_conv_2d & args,
4255+
device const char * weights,
4256+
device const char * src,
4257+
device char * dst,
4258+
uint3 tgpig[[threadgroup_position_in_grid]],
4259+
uint3 tgpg[[threadgroups_per_grid]],
4260+
uint3 tpitg[[thread_position_in_threadgroup]],
4261+
uint3 ntg[[threads_per_threadgroup]]);
4262+
41494263
typedef void (conv_transpose_1d_t)(
41504264
constant ggml_metal_kargs_conv_transpose_1d & args,
41514265
device const float * src0,

0 commit comments

Comments
 (0)