ggml-webgpu: add the upscale shader (#22419)
* shader(upscale): add the upscale shader with nearest, bilinear and bicubic implementations * shader(upscale): use macro
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
#ifndef GGML_WEBGPU_SHADER_LIB_HPP
|
||||
#define GGML_WEBGPU_SHADER_LIB_HPP
|
||||
|
||||
#include "ggml-impl.h"
|
||||
#include "ggml-wgsl-shaders.hpp"
|
||||
#include "ggml.h"
|
||||
#include "pre_wgsl.hpp"
|
||||
@@ -405,6 +406,31 @@ struct ggml_webgpu_scale_pipeline_key_hash {
|
||||
}
|
||||
};
|
||||
|
||||
/** Upscale **/
|
||||
|
||||
struct ggml_webgpu_upscale_pipeline_key {
|
||||
ggml_type input_type;
|
||||
ggml_type output_type;
|
||||
uint32_t base_mode;
|
||||
bool antialias;
|
||||
|
||||
bool operator==(const ggml_webgpu_upscale_pipeline_key & other) const {
|
||||
return input_type == other.input_type && output_type == other.output_type && base_mode == other.base_mode &&
|
||||
antialias == other.antialias;
|
||||
}
|
||||
};
|
||||
|
||||
struct ggml_webgpu_upscale_pipeline_key_hash {
|
||||
size_t operator()(const ggml_webgpu_upscale_pipeline_key & key) const {
|
||||
size_t seed = 0;
|
||||
ggml_webgpu_hash_combine(seed, key.input_type);
|
||||
ggml_webgpu_hash_combine(seed, key.output_type);
|
||||
ggml_webgpu_hash_combine(seed, key.base_mode);
|
||||
ggml_webgpu_hash_combine(seed, key.antialias);
|
||||
return seed;
|
||||
}
|
||||
};
|
||||
|
||||
/** Concat **/
|
||||
|
||||
struct ggml_webgpu_concat_pipeline_key {
|
||||
@@ -1049,6 +1075,8 @@ class ggml_webgpu_shader_lib {
|
||||
webgpu_pipeline,
|
||||
ggml_webgpu_rms_norm_mul_pipeline_key_hash>
|
||||
rms_norm_mul_pipelines;
|
||||
std::unordered_map<ggml_webgpu_upscale_pipeline_key, webgpu_pipeline, ggml_webgpu_upscale_pipeline_key_hash>
|
||||
upscale_pipelines;
|
||||
|
||||
public:
|
||||
ggml_webgpu_shader_lib(wgpu::Device device) { this->device = device; }
|
||||
@@ -2947,6 +2975,72 @@ class ggml_webgpu_shader_lib {
|
||||
return im2col_pipelines[key];
|
||||
}
|
||||
|
||||
webgpu_pipeline get_upscale_pipeline(const ggml_webgpu_shader_lib_context & context) {
|
||||
const uint32_t mode_flags = (uint32_t) ggml_get_op_params_i32(context.dst, 0);
|
||||
const uint32_t base_mode = mode_flags & 0xFFu;
|
||||
const bool antialias = (mode_flags & GGML_SCALE_FLAG_ANTIALIAS) != 0u;
|
||||
|
||||
ggml_webgpu_upscale_pipeline_key key = {};
|
||||
key.input_type = context.src0->type;
|
||||
key.output_type = context.dst->type;
|
||||
key.base_mode = base_mode;
|
||||
key.antialias = antialias;
|
||||
|
||||
auto it = upscale_pipelines.find(key);
|
||||
if (it != upscale_pipelines.end()) {
|
||||
return it->second;
|
||||
}
|
||||
|
||||
std::vector<std::string> defines;
|
||||
std::string variant = "upscale";
|
||||
|
||||
if (key.input_type == GGML_TYPE_F16) {
|
||||
defines.push_back("SRC_F16");
|
||||
variant += "_src_f16";
|
||||
} else {
|
||||
variant += "_src_f32";
|
||||
}
|
||||
|
||||
if (key.output_type == GGML_TYPE_F16) {
|
||||
defines.push_back("DST_F16");
|
||||
variant += "_dst_f16";
|
||||
} else {
|
||||
variant += "_dst_f32";
|
||||
}
|
||||
|
||||
switch (base_mode) {
|
||||
case GGML_SCALE_MODE_NEAREST:
|
||||
defines.push_back("NEAREST");
|
||||
variant += "_nearest";
|
||||
break;
|
||||
case GGML_SCALE_MODE_BILINEAR:
|
||||
defines.push_back("BILINEAR");
|
||||
variant += "_bilinear";
|
||||
break;
|
||||
case GGML_SCALE_MODE_BICUBIC:
|
||||
defines.push_back("BICUBIC");
|
||||
variant += "_bicubic";
|
||||
break;
|
||||
default:
|
||||
GGML_ABORT("Unsupported upscale mode");
|
||||
}
|
||||
|
||||
if (antialias) {
|
||||
defines.push_back("ANTIALIAS");
|
||||
variant += "_aa";
|
||||
}
|
||||
|
||||
defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
|
||||
|
||||
auto processed = preprocessor.preprocess(wgsl_upscale, 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;
|
||||
upscale_pipelines[key] = pipeline;
|
||||
return upscale_pipelines[key];
|
||||
}
|
||||
|
||||
private:
|
||||
static webgpu_pipeline ggml_webgpu_create_pipeline(wgpu::Device & device,
|
||||
std::string shader_code,
|
||||
|
||||
Reference in New Issue
Block a user