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
2026-07-09 12:29:15 +03:00
committed by GitHub
co-authored by Georgi Gerganov
parent ccb0c34223
commit 92b187c97e
9 changed files with 381 additions and 21 deletions
+22 -11
View File
@@ -7299,6 +7299,13 @@ struct ggml_conv_2d_dw_params {
int dilation_y;
};
static inline float ggml_conv_2d_dw_knl_f32(const char * data, int64_t i, ggml_type type) {
if (type == GGML_TYPE_F16) {
return GGML_FP16_TO_FP32(((const ggml_fp16_t *)data)[i]);
}
return ((const float *)data)[i];
}
static void ggml_compute_forward_conv_2d_dw_cwhn(
const ggml_compute_params * params,
const ggml_tensor * src,
@@ -7307,7 +7314,8 @@ static void ggml_compute_forward_conv_2d_dw_cwhn(
const ggml_conv_2d_dw_params & p) {
const int64_t c = p.channels;
const float * knl_data = (const float *)kernel->data;
const char * knl_data = (const char *)kernel->data;
const ggml_type knl_type = kernel->type;
const int64_t rows_total = p.dst_h * p.batch;
const int64_t rows_per_thread = (rows_total + params->nth - 1) / params->nth;
@@ -7315,13 +7323,15 @@ static void ggml_compute_forward_conv_2d_dw_cwhn(
const int64_t row_end = MIN(row_start + rows_per_thread, rows_total);
#ifdef GGML_SIMD
int64_t c_pkg_end = 0;
if (knl_type == GGML_TYPE_F32) {
#if defined(__ARM_FEATURE_SVE)
const int64_t pkg_size = svcntw();
#else
const int64_t pkg_size = GGML_F32_EPR;
#endif
const int64_t pkg_count = c / pkg_size;
const int64_t c_pkg_end = pkg_count * pkg_size;
c_pkg_end = (c / pkg_size) * pkg_size;
}
#else
const int64_t c_pkg_end = 0;
#endif
@@ -7335,8 +7345,7 @@ static void ggml_compute_forward_conv_2d_dw_cwhn(
const int64_t src_x_base = dst_x * p.stride_x - p.pad_x;
#ifdef GGML_SIMD
// Vectorized loop
for (int64_t c_i = 0; c_i < c_pkg_end; c_i += pkg_size) {
for (int64_t c_i = 0; c_i < c_pkg_end; c_i += GGML_F32_EPR) {
GGML_F32_VEC sum = GGML_F32_VEC_ZERO;
for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) {
const int64_t src_y = src_y_base + knl_y * p.dilation_y;
@@ -7348,7 +7357,8 @@ static void ggml_compute_forward_conv_2d_dw_cwhn(
if (src_x < 0 || src_x >= p.src_w) {
continue;
}
GGML_F32_VEC k = GGML_F32_VEC_LOAD(knl_data + (knl_y * p.knl_w + knl_x) * c + c_i);
const float * kp = (const float *)knl_data + (knl_y * p.knl_w + knl_x) * c + c_i;
GGML_F32_VEC k = GGML_F32_VEC_LOAD(kp);
GGML_F32_VEC s = GGML_F32_VEC_LOAD(src_data + (src_y * p.src_w + src_x) * c + c_i);
sum = GGML_F32_VEC_FMA(sum, k, s);
}
@@ -7356,7 +7366,6 @@ static void ggml_compute_forward_conv_2d_dw_cwhn(
GGML_F32_VEC_STORE(dst_data + c_i, sum);
}
#endif
// Scalar loop
for (int64_t c_i = c_pkg_end; c_i < c; ++c_i) {
float sum = 0.0f;
for (int64_t knl_y = 0; knl_y < p.knl_h; ++knl_y) {
@@ -7369,7 +7378,7 @@ static void ggml_compute_forward_conv_2d_dw_cwhn(
if (src_x < 0 || src_x >= p.src_w) {
continue;
}
sum += knl_data[(knl_y * p.knl_w + knl_x) * c + c_i]
sum += ggml_conv_2d_dw_knl_f32(knl_data, (knl_y * p.knl_w + knl_x) * c + c_i, knl_type)
* src_data[(src_y * p.src_w + src_x) * c + c_i];
}
}
@@ -7390,9 +7399,11 @@ static void ggml_compute_forward_conv_2d_dw_whcn(
const int64_t per_thread = (n + params->nth - 1) / params->nth;
const int64_t start = params->ith * per_thread;
const int64_t end = MIN(start + per_thread, n);
const char * knl_base = (const char *)kernel->data;
const ggml_type knl_type = kernel->type;
for (int64_t i = start; i < end; ++i) {
const float * knl_data = (const float *)kernel->data + (i % p.channels) * p.knl_w * p.knl_h;
const int64_t knl_offset = (i % p.channels) * p.knl_w * p.knl_h;
const float * src_data = (const float *)src->data + i * p.src_w * p.src_h;
float * dst_data = (float *)dst->data + i * p.dst_w * p.dst_h;
@@ -7410,7 +7421,7 @@ static void ggml_compute_forward_conv_2d_dw_whcn(
if (src_x < 0 || src_x >= p.src_w) {
continue;
}
sum += knl_data[knl_y * p.knl_w + knl_x]
sum += ggml_conv_2d_dw_knl_f32(knl_base, knl_offset + knl_y * p.knl_w + knl_x, knl_type)
* src_data[src_y * p.src_w + src_x];
}
}
@@ -7442,13 +7453,13 @@ void ggml_compute_forward_conv_2d_dw(
p.dilation_x = dst->op_params[4];
p.dilation_y = dst->op_params[5];
GGML_ASSERT(kernel->type == GGML_TYPE_F32 || kernel->type == GGML_TYPE_F16);
GGML_ASSERT(kernel->ne[3] == p.channels);
GGML_ASSERT(dst->ne[3] == p.batch);
if (ggml_is_contiguous(src)) {
ggml_compute_forward_conv_2d_dw_whcn(params, src, kernel, dst, p);
} else if (ggml_is_contiguous_channels(src)) {
// kernel should also have channels most contiguous in memory
GGML_ASSERT(kernel->nb[0] >= kernel->nb[2] && kernel->nb[1] >= kernel->nb[0]);
ggml_compute_forward_conv_2d_dw_cwhn(params, src, kernel, dst, p);
} else {