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:
co-authored by
Georgi Gerganov
parent
ccb0c34223
commit
92b187c97e
+22
-11
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user