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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user