diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index d7692363a..bed9265b8 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -355,6 +355,30 @@ struct ggml_webgpu_conv2d_pipeline_key_hash { } }; +// Same type fields as conv2d plus the input layout (WHCN vs CWHN). +struct ggml_webgpu_conv2d_dw_pipeline_key { + ggml_type weight_type; + ggml_type input_type; + ggml_type output_type; + bool whcn; + + bool operator==(const ggml_webgpu_conv2d_dw_pipeline_key & other) const { + return weight_type == other.weight_type && input_type == other.input_type && output_type == other.output_type && + whcn == other.whcn; + } +}; + +struct ggml_webgpu_conv2d_dw_pipeline_key_hash { + size_t operator()(const ggml_webgpu_conv2d_dw_pipeline_key & key) const { + size_t seed = 0; + ggml_webgpu_hash_combine(seed, key.weight_type); + ggml_webgpu_hash_combine(seed, key.input_type); + ggml_webgpu_hash_combine(seed, key.output_type); + ggml_webgpu_hash_combine(seed, key.whcn); + return seed; + } +}; + /** Im2Col **/ struct ggml_webgpu_im2col_pipeline_key { ggml_type input_type; @@ -1210,6 +1234,8 @@ class ggml_webgpu_shader_lib { soft_max_pipelines; std::unordered_map conv2d_pipelines; + std::unordered_map + conv2d_dw_pipelines; std::unordered_map im2col_pipelines; @@ -3172,6 +3198,50 @@ class ggml_webgpu_shader_lib { return conv2d_pipelines[key]; } + // whcn selects the input layout: contiguous WHCN vs contiguous-channels CWHN + webgpu_pipeline get_conv2d_dw_pipeline(const ggml_webgpu_shader_lib_context & context, bool whcn) { + ggml_webgpu_conv2d_dw_pipeline_key key = {}; + key.weight_type = context.src0->type; + key.input_type = context.src1->type; + key.output_type = context.dst->type; + key.whcn = whcn; + + auto it = conv2d_dw_pipelines.find(key); + if (it != conv2d_dw_pipelines.end()) { + return it->second; + } + + std::vector defines; + std::string variant = whcn ? "conv_2d_dw_whcn" : "conv_2d_dw_cwhn"; + + 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 CONV_2D_DW shader"); + } + }; + + push_type_defines("WEIGHT", key.weight_type); + push_type_defines("INPUT", key.input_type); + push_type_defines("OUTPUT", key.output_type); + if (whcn) { + defines.push_back("WHCN"); + } + defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size)); + + auto processed = preprocessor.preprocess(wgsl_conv2d_dw, defines); + auto decisions = std::make_shared(); + decisions->wg_size = context.max_wg_size; + webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant); + pipeline.context = decisions; + conv2d_dw_pipelines[key] = pipeline; + return conv2d_dw_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; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 29025e9ba..75286ec73 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -978,6 +978,67 @@ static webgpu_encoded_op ggml_webgpu_conv_2d(webgpu_context & ctx, return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y); } +// Same param/binding layout as conv_2d; the shader differs +static webgpu_encoded_op ggml_webgpu_conv_2d_dw(webgpu_context & ctx, + ggml_tensor * src0, + ggml_tensor * src1, + ggml_tensor * dst) { + const int32_t s0 = ggml_get_op_params_i32(dst, 0); + const int32_t s1 = ggml_get_op_params_i32(dst, 1); + const int32_t p0 = ggml_get_op_params_i32(dst, 2); + const int32_t p1 = ggml_get_op_params_i32(dst, 3); + const int32_t d0 = ggml_get_op_params_i32(dst, 4); + const int32_t d1 = ggml_get_op_params_i32(dst, 5); + + // Scalar params matching conv2d_dw.wgsl (weight src0 [KW,KH,1,C], input src1, output dst). + std::vector params = { + (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)), + (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)), + (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), + + (uint32_t) ggml_nelements(dst), + (uint32_t) dst->ne[2], + (uint32_t) dst->ne[3], + (uint32_t) dst->ne[0], + (uint32_t) dst->ne[1], + (uint32_t) src1->ne[0], + (uint32_t) src1->ne[1], + (uint32_t) src0->ne[0], + (uint32_t) src0->ne[1], + + (uint32_t) s0, + (uint32_t) s1, + (uint32_t) p0, + (uint32_t) p1, + (uint32_t) d0, + (uint32_t) d1, + }; + + std::vector entries = { + ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0), + ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1), + ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst), + }; + + ggml_webgpu_shader_lib_context shader_lib_ctx = {}; + shader_lib_ctx.src0 = src0; + shader_lib_ctx.src1 = src1; + shader_lib_ctx.dst = dst; + shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; + + // Input layout: contiguous -> WHCN, contiguous-channels -> CWHN + const bool whcn = ggml_is_contiguous(src1); + webgpu_pipeline pipeline = ctx->shader_lib->get_conv2d_dw_pipeline(shader_lib_ctx, whcn); + auto * decisions = static_cast(pipeline.context.get()); + + uint32_t wg_x; + uint32_t wg_y; + uint32_t total_wg = CEIL_DIV((uint32_t) ggml_nelements(dst), decisions->wg_size); + compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y); + + return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y); +} + static webgpu_encoded_op ggml_webgpu_im2col(webgpu_context & ctx, ggml_tensor * src0, ggml_tensor * src1, @@ -3164,6 +3225,8 @@ static std::optional ggml_webgpu_encode(webgpu_context ctx, return ggml_webgpu_sum_rows(ctx, src0, node); case GGML_OP_CONV_2D: return ggml_webgpu_conv_2d(ctx, src0, src1, node); + case GGML_OP_CONV_2D_DW: + return ggml_webgpu_conv_2d_dw(ctx, src0, src1, node); case GGML_OP_IM2COL: return ggml_webgpu_im2col(ctx, src0, src1, node); case GGML_OP_UPSCALE: @@ -4349,6 +4412,12 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); break; + case GGML_OP_CONV_2D_DW: + supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && + (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) && + (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16) && + (ggml_is_contiguous(src1) || ggml_is_contiguous_channels(src1)); + break; case GGML_OP_IM2COL: supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/conv2d_dw.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/conv2d_dw.wgsl new file mode 100644 index 000000000..42d6f027c --- /dev/null +++ b/ggml/src/ggml-webgpu/wgsl-shaders/conv2d_dw.wgsl @@ -0,0 +1,137 @@ +#include "common_decls.tmpl" +enable f16; + +// Ported from the Vulkan backend's conv2d_dw.comp. Two variants (based on WHCN) +// selected by the input (src1) layout: contiguous -> WHCN, else CWHN. +// weight (src0) is [KW,KH,1,C]; output matches the input layout. + +@group(0) @binding(0) +#if defined(WEIGHT_F32) +var weights: array; +#elif defined(WEIGHT_F16) +var weights: array; +#endif + +@group(0) @binding(1) +#if defined(INPUT_F32) +var input: array; +#elif defined(INPUT_F16) +var input: array; +#endif + +@group(0) @binding(2) +#if defined(OUTPUT_F32) +var output: array; +#elif defined(OUTPUT_F16) +var output: array; +#endif + +struct Params { + offset_w: u32, + offset_i: u32, + offset_o: u32, + + ne: u32, + channels: u32, + batches: u32, + dst_w: u32, dst_h: u32, + src_w: u32, src_h: u32, + knl_w: u32, knl_h: u32, + + stride_x: i32, stride_y: i32, + pad_x: i32, pad_y: i32, + dilation_x: i32, dilation_y: i32, +}; + +@group(0) @binding(3) +var params: Params; + +fn load_weight(idx: u32) -> f32 { + #if defined(WEIGHT_F32) + return weights[idx]; + #elif defined(WEIGHT_F16) + return f32(weights[idx]); + #endif +} +fn load_input(idx: u32) -> f32 { + #if defined(INPUT_F32) + return input[idx]; + #elif defined(INPUT_F16) + return f32(input[idx]); + #endif +} +fn store_output(idx: u32, val: f32) { + #if defined(OUTPUT_F32) + output[idx] = val; + #elif defined(OUTPUT_F16) + output[idx] = f16(val); + #endif +} + +#if defined(WHCN) +// Input/output/kernel contiguous in [W, H, C, N] order (kernel [KW,KH,C]). +fn conv_2d_dw(idx: u32) -> f32 { + let i0 = idx / params.dst_w; + let dst_x = idx - i0 * params.dst_w; + let i1 = i0 / params.dst_h; + let dst_y = i0 - i1 * params.dst_h; + let n = i1 / params.channels; + let c = i1 - n * params.channels; + + let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w + + c * params.src_h * params.src_w; + let knl_i = params.offset_w + c * params.knl_h * params.knl_w; + + var sum: f32 = 0.0; + for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) { + let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y; + if (src_y < 0 || src_y >= i32(params.src_h)) { continue; } + for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) { + let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x; + if (src_x < 0 || src_x >= i32(params.src_w)) { continue; } + let v = load_input(src_i + u32(src_y) * params.src_w + u32(src_x)); + let k = load_weight(knl_i + ky * params.knl_w + kx); + sum += v * k; + } + } + return sum; +} +#else +// Channels contiguous (CWHN): channel is the innermost axis. +fn conv_2d_dw(idx: u32) -> f32 { + let i0 = idx / params.channels; + let c = idx - i0 * params.channels; + let i1 = i0 / params.dst_w; + let dst_x = i0 - i1 * params.dst_w; + let n = i1 / params.dst_h; + let dst_y = i1 - n * params.dst_h; + + let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w; + let src_row = params.src_w * params.channels; + let knl_row = params.knl_w * params.channels; + + var sum: f32 = 0.0; + for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) { + let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y; + if (src_y < 0 || src_y >= i32(params.src_h)) { continue; } + for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) { + let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x; + if (src_x < 0 || src_x >= i32(params.src_w)) { continue; } + let v = load_input(src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c); + let k = load_weight(params.offset_w + ky * knl_row + kx * params.channels + c); + sum += v * k; + } + } + return sum; +} +#endif + +@compute @workgroup_size(WG_SIZE) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_wg: vec3 +) { + let idx = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y; + if (idx >= params.ne) { return; } + store_output(params.offset_o + idx, conv_2d_dw(idx)); +}