ggml-webgpu: add support for im2col (#22259)

* shader(im2col): implement the im2col shader

* shader(im2col): clean the formatting issues

* shader(im2col): clean the editorconfig checker warning

* fix(shader): address the workgroup issues of im2col and conv2d
This commit is contained in:
Chen Yuan
2026-04-22 23:17:41 -04:00
committed by GitHub
parent 86db42e97f
commit b76429a69c
3 changed files with 268 additions and 19 deletions
@@ -281,6 +281,25 @@ struct ggml_webgpu_conv2d_pipeline_key_hash {
}
};
/** Im2Col **/
struct ggml_webgpu_im2col_pipeline_key {
ggml_type input_type;
ggml_type output_type;
bool operator==(const ggml_webgpu_im2col_pipeline_key & other) const {
return input_type == other.input_type && output_type == other.output_type;
}
};
struct ggml_webgpu_im2col_pipeline_key_hash {
size_t operator()(const ggml_webgpu_im2col_pipeline_key & key) const {
size_t seed = 0;
ggml_webgpu_hash_combine(seed, key.input_type);
ggml_webgpu_hash_combine(seed, key.output_type);
return seed;
}
};
/** Gated Delta Net **/
struct ggml_webgpu_gated_delta_net_pipeline_key {
int type;
@@ -833,6 +852,8 @@ class ggml_webgpu_shader_lib {
soft_max_pipelines;
std::unordered_map<ggml_webgpu_conv2d_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_pipeline_key_hash>
conv2d_pipelines;
std::unordered_map<ggml_webgpu_im2col_pipeline_key, webgpu_pipeline, ggml_webgpu_im2col_pipeline_key_hash>
im2col_pipelines;
std::unordered_map<ggml_webgpu_rms_norm_mul_pipeline_key,
webgpu_pipeline,
@@ -2504,6 +2525,44 @@ class ggml_webgpu_shader_lib {
return conv2d_pipelines[key];
}
webgpu_pipeline get_im2col_pipeline(const ggml_webgpu_shader_lib_context & context) {
ggml_webgpu_im2col_pipeline_key key = {};
key.input_type = context.src1->type;
key.output_type = context.dst->type;
auto it = im2col_pipelines.find(key);
if (it != im2col_pipelines.end()) {
return it->second;
}
std::vector<std::string> defines;
std::string variant = "im2col";
auto push_type_defines = [&](const char * prefix, ggml_type type) {
std::string s_prefix = prefix;
if (type == GGML_TYPE_F32) {
defines.push_back(s_prefix + "_F32");
} else if (type == GGML_TYPE_F16) {
defines.push_back(s_prefix + "_F16");
} else {
GGML_ABORT("Unsupported type for IM2COL shader");
}
};
push_type_defines("INPUT", key.input_type);
push_type_defines("OUTPUT", key.output_type);
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
auto processed = preprocessor.preprocess(wgsl_im2col, defines);
auto decisions = std::make_shared<ggml_webgpu_generic_shader_decisions>();
decisions->wg_size = context.max_wg_size;
webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
pipeline.context = decisions;
im2col_pipelines[key] = pipeline;
return im2col_pipelines[key];
}
private:
static webgpu_pipeline ggml_webgpu_create_pipeline(wgpu::Device & device,
std::string shader_code,