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:
Chen Yuan
2026-04-30 22:22:18 -07:00
committed by GitHub
parent a95a11e5b8
commit aab68217b7
3 changed files with 383 additions and 0 deletions
@@ -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,