metal : add CONV_2D_DW (depthwise convolution) support (#21565)
* metal : add CONV_2D_DW (depthwise 2D convolution) support * test : add perf cases for CONV_2D_DW * metal : use 3D dispatch for CONV_2D_DW kernel * metal : add channel-tiled CONV_2D_DW kernel for non-contiguous layouts * metal : simplify CONV_2D_DW dispatch and trim comments * metal : merge duplicate CONV_2D_DW pipeline getters * tests : add F16 CONV2D_DW tests * cpu : fix F16 kernel support for CONV_2D_DW * tests : remove commented-out CONV_2D_DW test block --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
Sou-ly
co-authored by
Georgi Gerganov
Georgi Gerganov
parent
ccb0c34223
commit
92b187c97e
@@ -4908,6 +4908,202 @@ kernel void kernel_conv_2d<half>(
|
||||
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 <typename TK>
|
||||
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 <typename TK>
|
||||
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<float>(
|
||||
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<half>(
|
||||
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<float>(
|
||||
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<half>(
|
||||
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]]);
|
||||
|
||||
typedef void (conv_transpose_1d_t)(
|
||||
constant ggml_metal_kargs_conv_transpose_1d & args,
|
||||
device const float * src0,
|
||||
|
||||
Reference in New Issue
Block a user