From 60a9937de8662f858e457f1e4be96d0015a8321a Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 11:56:57 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/vulkan/custom_ops_lib.py | 36 +++ .../webgpu/runtime/WebGPUShaderRegistry.cpp | 42 +++- .../runtime/ops/binary_op/binary_op.wgsl | 10 +- .../runtime/ops/binary_op/binary_op.yaml | 9 + .../ops/binary_op/binary_sub_int32_wgsl.h | 78 ++++++ .../runtime/ops/boolean_op/BooleanOp.cpp | 18 +- .../runtime/ops/boolean_op/boolean_op.wgsl | 6 +- .../runtime/ops/boolean_op/compare_eq_wgsl.h | 8 +- .../runtime/ops/boolean_op/compare_ge_wgsl.h | 8 +- .../runtime/ops/boolean_op/compare_gt_wgsl.h | 8 +- .../runtime/ops/boolean_op/compare_le_wgsl.h | 8 +- .../runtime/ops/boolean_op/compare_lt_wgsl.h | 8 +- .../runtime/ops/boolean_op/compare_ne_wgsl.h | 8 +- .../runtime/ops/boolean_op/logical_not_wgsl.h | 8 +- backends/webgpu/runtime/ops/gather/Gather.cpp | 9 +- .../webgpu/runtime/ops/gather/gather.wgsl | 6 +- .../webgpu/runtime/ops/gather/gather_wgsl.h | 8 +- backends/webgpu/runtime/ops/index/Index.cpp | 11 +- backends/webgpu/runtime/ops/index/index.wgsl | 6 +- .../webgpu/runtime/ops/index/index_wgsl.h | 8 +- .../ops/quantized_linear/QuantizedLinear.cpp | 117 +++++++-- .../q4gsw_linear_m3_shared_bicol.wgsl | 159 +++++++++++++ .../q4gsw_linear_m3_shared_bicol_wgsl.h | 183 +++++++++++++++ .../webgpu/runtime/ops/scatter/Scatter.cpp | 124 ++++++++++ .../webgpu/runtime/ops/scatter/scatter.wgsl | 19 ++ .../ops/scatter/scatter_unique_indices.wgsl | 35 +++ .../ops/scatter/scatter_unique_indices_wgsl.h | 59 +++++ .../webgpu/runtime/ops/scatter/scatter_wgsl.h | 43 ++++ .../ops/split_with_sizes/SplitWithSizes.cpp | 9 +- backends/webgpu/runtime/ops/sub/BinaryOp.cpp | 141 ++++++----- backends/webgpu/runtime/ops/sub/SubStorage.h | 34 +++ .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 48 ++-- .../ops/to_copy/to_copy_bool_to_float.wgsl | 6 +- .../ops/to_copy/to_copy_bool_to_float_wgsl.h | 8 +- .../runtime/ops/to_copy/to_copy_convert.wgsl | 6 +- .../ops/to_copy/to_copy_float_to_int_wgsl.h | 8 +- .../ops/to_copy/to_copy_int_to_float_wgsl.h | 8 +- backends/webgpu/runtime/ops/topk/TopK.cpp | 106 +++++++++ backends/webgpu/runtime/ops/topk/topk.wgsl | 198 ++++++++++++++++ backends/webgpu/runtime/ops/topk/topk_wgsl.h | 222 ++++++++++++++++++ backends/webgpu/runtime/ops/where/Where.cpp | 17 +- backends/webgpu/runtime/ops/where/where.wgsl | 6 +- .../webgpu/runtime/ops/where/where_wgsl.h | 8 +- backends/webgpu/test/test_wgsl_codegen.py | 16 +- 44 files changed, 1704 insertions(+), 179 deletions(-) create mode 100644 backends/webgpu/runtime/ops/binary_op/binary_sub_int32_wgsl.h create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol.wgsl create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol_wgsl.h create mode 100644 backends/webgpu/runtime/ops/scatter/Scatter.cpp create mode 100644 backends/webgpu/runtime/ops/scatter/scatter.wgsl create mode 100644 backends/webgpu/runtime/ops/scatter/scatter_unique_indices.wgsl create mode 100644 backends/webgpu/runtime/ops/scatter/scatter_unique_indices_wgsl.h create mode 100644 backends/webgpu/runtime/ops/scatter/scatter_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sub/SubStorage.h create mode 100644 backends/webgpu/runtime/ops/topk/TopK.cpp create mode 100644 backends/webgpu/runtime/ops/topk/topk.wgsl create mode 100644 backends/webgpu/runtime/ops/topk/topk_wgsl.h diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index ffbbc796c62..aea2ecd3677 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -350,6 +350,42 @@ def linear_q4gsw_backward(ctx, grad_out): setup_context=linear_q4gsw_setup_context, ) +######################## +## scatter_src_unique ## +######################## + + +def scatter_src_unique_impl( + self: torch.Tensor, + dim: int, + index: torch.Tensor, + src: torch.Tensor, +) -> torch.Tensor: + normalized_dim = dim if dim >= 0 else dim + self.dim() + if normalized_dim != self.dim() - 1: + raise ValueError("scatter_src_unique requires the final dimension") + if not isinstance(index, FakeTensor): + flattened = index.detach().reshape(-1) + if torch.unique(flattened).numel() != flattened.numel(): + raise ValueError("scatter_src_unique requires unique destinations") + return torch.scatter(self, dim, index, src) + + +def scatter_src_unique_meta( + self: torch.Tensor, + dim: int, + index: torch.Tensor, + src: torch.Tensor, +) -> torch.Tensor: + return torch.empty_like(self) + + +name = "scatter_src_unique" +lib.define(f"{name}(Tensor self, int dim, Tensor index, Tensor src) -> Tensor") +lib.impl(name, scatter_src_unique_impl, "CompositeExplicitAutograd") +lib.impl(name, scatter_src_unique_meta, "Meta") +scatter_src_unique_op = getattr(getattr(torch.ops, namespace), name) + name = "linear_dq8ca_q4gsw" lib.define( f""" diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index fb40eeeb760..943c41acd99 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -101,6 +102,7 @@ #include #include #include +#include #include #include #include @@ -116,6 +118,8 @@ #include #include #include +#include +#include #include #include #include @@ -136,6 +140,7 @@ #include #include #include +#include #include #include #include @@ -161,7 +166,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -288,6 +293,13 @@ constexpr std::array kShaderRegistry = {{ kBinarySubWorkgroupSizeY, kBinarySubWorkgroupSizeZ, }, + { + "binary_sub_int32", + kBinarySubInt32WGSL, + kBinarySubInt32WorkgroupSizeX, + kBinarySubInt32WorkgroupSizeY, + kBinarySubInt32WorkgroupSizeZ, + }, { "bitwise_not", kBitwiseNotWGSL, @@ -792,6 +804,13 @@ constexpr std::array kShaderRegistry = {{ kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY, kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ, }, + { + "q4gsw_linear_m3_shared_bicol", + kQ4gswLinearM3SharedBicolWGSL, + kQ4gswLinearM3SharedBicolWorkgroupSizeX, + kQ4gswLinearM3SharedBicolWorkgroupSizeY, + kQ4gswLinearM3SharedBicolWorkgroupSizeZ, + }, { "q4gsw_qkv_bk64", kQ4gswQkvBk64WGSL, @@ -967,6 +986,20 @@ constexpr std::array kShaderRegistry = {{ kRsqrtWorkgroupSizeY, kRsqrtWorkgroupSizeZ, }, + { + "scatter", + kScatterWGSL, + kScatterWorkgroupSizeX, + kScatterWorkgroupSizeY, + kScatterWorkgroupSizeZ, + }, + { + "scatter_unique_indices", + kScatterUniqueIndicesWGSL, + kScatterUniqueIndicesWorkgroupSizeX, + kScatterUniqueIndicesWorkgroupSizeY, + kScatterUniqueIndicesWorkgroupSizeZ, + }, { "sdpa_compute_attn_weights", kSdpaComputeAttnWeightsWGSL, @@ -1128,6 +1161,13 @@ constexpr std::array kShaderRegistry = {{ kToCopyIntToFloatWorkgroupSizeY, kToCopyIntToFloatWorkgroupSizeZ, }, + { + "topk", + kTopkWGSL, + kTopkWorkgroupSizeX, + kTopkWorkgroupSizeY, + kTopkWorkgroupSizeZ, + }, { "update_cache", kUpdateCacheWGSL, diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl index 4b42665013f..65e56ff0ba1 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl @@ -1,6 +1,6 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; +@group(0) @binding(0) var input1: array<${SCALAR_TYPE}>; +@group(0) @binding(1) var input2: array<${SCALAR_TYPE}>; +@group(0) @binding(2) var output: array<${SCALAR_TYPE}>; struct TensorMeta { ndim: u32, @@ -14,12 +14,12 @@ struct TensorMeta { override wg_size: u32 = 64u; $if USE_ALPHA: - override alpha: f32 = 1.0; + override alpha: ${ALPHA_TYPE} = ${ALPHA_DEFAULT}; $if INLINE: @compute @workgroup_size(wg_size, 1, 1) $else: - fn op(a: f32, b: f32) -> f32 { + fn op(a: ${SCALAR_TYPE}, b: ${SCALAR_TYPE}) -> ${SCALAR_TYPE} { return ${OP_EXPR}; } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml index bd4cd452d35..1bbc41a55cf 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml @@ -5,6 +5,9 @@ binary_op: INLINE: 0 SAME_EXPR: input1[idx] + input2[idx] BROADCAST_EXPR: input1[l1] + input2[l2] + SCALAR_TYPE: f32 + ALPHA_TYPE: f32 + ALPHA_DEFAULT: 1.0 shader_variants: - NAME: binary_div OP_EXPR: a / b @@ -12,6 +15,12 @@ binary_op: - NAME: binary_sub OP_EXPR: a - alpha * b USE_ALPHA: 1 + - NAME: binary_sub_int32 + OP_EXPR: bitcast(bitcast(a) - bitcast(alpha) * bitcast(b)) + USE_ALPHA: 1 + SCALAR_TYPE: i32 + ALPHA_TYPE: i32 + ALPHA_DEFAULT: 1i - NAME: binary_minimum USE_ALPHA: 0 INLINE: 1 diff --git a/backends/webgpu/runtime/ops/binary_op/binary_sub_int32_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_sub_int32_wgsl.h new file mode 100644 index 00000000000..4bb84419e16 --- /dev/null +++ b/backends/webgpu/runtime/ops/binary_op/binary_sub_int32_wgsl.h @@ -0,0 +1,78 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from binary_op.wgsl - DO NOT EDIT. +// wgsl-sha256: 134151da070a891e539f6ede5974310c623a9a7d379c90e69f91bec56ddc9b29 +inline constexpr const char* kBinarySubInt32WGSL = R"( +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: array, 2>, + strides: array, 2>, +} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; +override alpha: i32 = 1i; + +fn op(a: i32, b: i32) -> i32 { + return bitcast(bitcast(a) - bitcast(alpha) * bitcast(b)); +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = op(input1[idx], input2[idx]); + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = op(input1[l1], input2[l2]); +} +)"; + +inline constexpr uint32_t kBinarySubInt32WorkgroupSizeX = 64; +inline constexpr uint32_t kBinarySubInt32WorkgroupSizeY = 1; +inline constexpr uint32_t kBinarySubInt32WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp index e62f387d258..8ecc3825fa2 100644 --- a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp +++ b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp @@ -72,8 +72,8 @@ void dispatch_bool_op( const uint32_t n_words = (numel + 3u) / 4u; uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); - uint32_t workgroup_count = - utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name); + const utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, n_words, wg_size, op_name); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -98,8 +98,12 @@ void dispatch_bool_op( &wg_size_constant, 1); - const size_t dispatch_idx = - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + op_name, + workgroup_count.y}); WGPUBuffer p_buf = params_buf; auto resize = @@ -114,8 +118,10 @@ void dispatch_bool_op( BoolOpParams p = {n, scalar, 0u, 0u}; wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p)); const uint32_t nw = (n + 3u) / 4u; - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name); + const utils::WgCount workgroups = + utils::compute_2d_workgroup_count(g.device(), nw, wg_size, op_name); + g.dispatch_at(dispatch_idx).workgroup_count_x = workgroups.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = workgroups.y; }; graph.add_tensor_resize_hook(self_id, resize); diff --git a/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl b/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl index 4f32792d7dc..944935b8986 100644 --- a/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl +++ b/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl @@ -18,8 +18,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h index 8b7b44e7430..14d79cdcb30 100644 --- a/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 2a3c203d0255086a6e67a9e7cb08538858e8e6cb6a5542f6acebf04b8ccca6b7 +// wgsl-sha256: 558bc966cc511d239ed4901644f515cae9b639bbc4150689c5b2954599417ac9 inline constexpr const char* kCompareEqWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h index bcf7239bd7e..86a0eabcbb2 100644 --- a/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 5fb319a54ed666644f118119639dd8cd33256ea5e5163c4c1392ce6e84b369cc +// wgsl-sha256: d080502cc35fe57711b986bf0abca83bf365de59d9180601717836c35cab37da inline constexpr const char* kCompareGeWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h index 8f911180c8a..7cd8214dd4a 100644 --- a/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: b4a35be8a34774a47be457130125336b5accf8d24615514f3fe1559159644491 +// wgsl-sha256: 8ab67c8415a30f2258ef64c109e94ef608a9b4c309f2812c8362ed626d0289c4 inline constexpr const char* kCompareGtWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h index 87b4ff91c2a..5b7af685cd9 100644 --- a/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 64b8988db9611c1aaff3ba373d32ecb7b276340df337b38a49e12a92727c00cb +// wgsl-sha256: 4ca729891e532e1140bc65cadfc6f1ff6a1529ec374098e6c8f3c5f6ea7ddc71 inline constexpr const char* kCompareLeWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h index 930601d229c..091954ae7c3 100644 --- a/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: df350a7d0b099d38f9e77a27ffe9e56afa93b2c6a3ff84006227e1c1c0b96521 +// wgsl-sha256: fbdc103126daa66d64ebcb0f2e097c755664b8b638142040e56e386b237b4635 inline constexpr const char* kCompareLtWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h index 9020fb77671..6c00bcc9ae0 100644 --- a/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 4f5492c393e494e6a9a734ea797cbfa85f0f1579451fc46f4291f57e1b5daaa2 +// wgsl-sha256: fd93c128bb61e98774b8da14b6fd00d795ab8577cfafd02d769d6b3dcc9e9708 inline constexpr const char* kCompareNeWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h index 65239060454..d85ac0e08b5 100644 --- a/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h +++ b/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from boolean_op.wgsl - DO NOT EDIT. -// wgsl-sha256: 06dc55c61e55c35eb12abf548fafb495097fa8d22daea930af6a32ed724a5b6a +// wgsl-sha256: 648b266b0dc53b413429c292547297686b9e6806684791ed8d9b561c118ce3cb inline constexpr const char* kLogicalNotWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -35,8 +35,10 @@ fn elem_bool(i: u32) -> bool { // One thread per output u32 word packs 4 bool bytes -> no inter-thread race. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let word_idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let word_idx = gid.x + gid.y * (num_workgroups.x * wg_size); let n_words = (params.num_elements + 3u) / 4u; if (word_idx >= n_words) { return; diff --git a/backends/webgpu/runtime/ops/gather/Gather.cpp b/backends/webgpu/runtime/ops/gather/Gather.cpp index 9236a6714a2..5361977c99b 100644 --- a/backends/webgpu/runtime/ops/gather/Gather.cpp +++ b/backends/webgpu/runtime/ops/gather/Gather.cpp @@ -75,7 +75,7 @@ void gather_impl(WebGPUGraph& graph, const std::vector& args) { } uint32_t wg_size = utils::clamp_workgroup_size(device, kGatherWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "gather"); WGPUConstantEntry wg_size_constant = {}; @@ -116,7 +116,12 @@ void gather_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "gather", + workgroup_count.y}); wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(self_meta_buf); diff --git a/backends/webgpu/runtime/ops/gather/gather.wgsl b/backends/webgpu/runtime/ops/gather/gather.wgsl index 5fa428952b3..3f12ed395fe 100644 --- a/backends/webgpu/runtime/ops/gather/gather.wgsl +++ b/backends/webgpu/runtime/ops/gather/gather.wgsl @@ -19,8 +19,10 @@ struct GatherParams { override wg_size: u32 = 256; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let o = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let o = gid.x + gid.y * (num_workgroups.x * wg_size); if (o >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/gather/gather_wgsl.h b/backends/webgpu/runtime/ops/gather/gather_wgsl.h index 8e679d9f9c2..9ba8553c0c5 100644 --- a/backends/webgpu/runtime/ops/gather/gather_wgsl.h +++ b/backends/webgpu/runtime/ops/gather/gather_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from gather.wgsl - DO NOT EDIT. -// wgsl-sha256: c0811711f4603f840241a9cdd458f7b941dcf32336ff84f2b6f1333b8d6f4cfb +// wgsl-sha256: 52a0eb48a4d7f128c60da9e26c0e35c830608317be553c22b536627756bc2ca0 inline constexpr const char* kGatherWGSL = R"( @group(0) @binding(0) var self_: array; @group(0) @binding(1) var indices: array; @@ -36,8 +36,10 @@ struct GatherParams { override wg_size: u32 = 256; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let o = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let o = gid.x + gid.y * (num_workgroups.x * wg_size); if (o >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/index/Index.cpp b/backends/webgpu/runtime/ops/index/Index.cpp index b122cc3af50..8ad3a2010a6 100644 --- a/backends/webgpu/runtime/ops/index/Index.cpp +++ b/backends/webgpu/runtime/ops/index/Index.cpp @@ -95,8 +95,8 @@ void index_impl(WebGPUGraph& graph, const std::vector& args) { uint32_t num_elements = static_cast(out_numel); uint32_t wg_size = utils::clamp_workgroup_size(device, kIndexWorkgroupSizeX); - uint32_t workgroup_count = - utils::compute_1d_workgroup_count(device, num_elements, wg_size, "index"); + const utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_elements, wg_size, "index"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -134,7 +134,12 @@ void index_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "index", + workgroup_count.y}); // The bind group keeps the uniform buffer alive until release. wgpuBufferRelease(uniform_buffer); diff --git a/backends/webgpu/runtime/ops/index/index.wgsl b/backends/webgpu/runtime/ops/index/index.wgsl index b0fd6df81bf..11ac7ad7f5e 100644 --- a/backends/webgpu/runtime/ops/index/index.wgsl +++ b/backends/webgpu/runtime/ops/index/index.wgsl @@ -10,8 +10,10 @@ struct Params { override wg_size: u32 = 64; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let out_bufi = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); if (out_bufi >= params.numel) { return; } diff --git a/backends/webgpu/runtime/ops/index/index_wgsl.h b/backends/webgpu/runtime/ops/index/index_wgsl.h index 839a3b164bb..2561f6ac8d7 100644 --- a/backends/webgpu/runtime/ops/index/index_wgsl.h +++ b/backends/webgpu/runtime/ops/index/index_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from index.wgsl - DO NOT EDIT. -// wgsl-sha256: daed48e60bfcf2b7420d277576d794137d3bff383aef4f68464c98c8a7235c8e +// wgsl-sha256: 56798b45b5ee8cdf295d305f08528698ca3e6847223c846ceabcfb9524ab0e3e inline constexpr const char* kIndexWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -27,8 +27,10 @@ struct Params { override wg_size: u32 = 64; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let out_bufi = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); if (out_bufi >= params.numel) { return; } diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index dcca24b357d..dcddab5e2f6 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include @@ -56,6 +57,9 @@ constexpr uint32_t kQ4gswSteelTile = 64u; constexpr uint32_t kQ4gswSteelBK = 16u; constexpr uint32_t kQ4gswSteelBK64 = 64u; constexpr uint32_t kQ4gswSteelInvocations = 256u; +constexpr uint32_t kQ4gswM3Invocations = 64u; +constexpr uint32_t kQ4gswM3StorageBytes = + 6u * kQ4gswM3Invocations * sizeof(float); constexpr const char* kQ4gswLinearShader = "q4gsw_linear"; constexpr const char* kQ4gswBicolShader = "q4gsw_linear_coop4_bicol"; @@ -96,6 +100,14 @@ bool steel_supported(WGPUDevice device) { limits.maxComputeInvocationsPerWorkgroup >= kQ4gswSteelInvocations; } +bool m3_shared_supported(WGPUDevice device) { + WGPULimits limits = {}; + return wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= kQ4gswM3Invocations && + limits.maxComputeWorkgroupSizeX >= kQ4gswM3Invocations && + limits.maxComputeWorkgroupStorageSize >= kQ4gswM3StorageBytes; +} + bool steel_bk64_eligible( WGPUDevice device, uint32_t K, @@ -148,6 +160,7 @@ uint32_t steel_bk64_workgroup_count( uint32_t compute_q4gsw_workgroup_count( WGPUDevice device, bool use_gemv, + bool use_m3_shared, bool use_bk64, bool use_steel, bool use_shmem_gemm, @@ -172,6 +185,20 @@ uint32_t compute_q4gsw_workgroup_count( } return wgc; } + if (use_m3_shared) { + const uint64_t pairs = (static_cast(n) + 1u) / 2u; + if (pairs == 0u || pairs > UINT32_MAX) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": M=3 N/2 out of range"); + } + const uint32_t wgc = + utils::clamp_workgroup_count(device, static_cast(pairs)); + if (wgc == 0u) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": zero M=3 dispatch"); + } + return wgc; + } if (use_bk64) { const uint32_t count = steel_bk64_workgroup_count(device, m, n, K); if (count == 0u) { @@ -218,11 +245,6 @@ struct Q4gswExecutionState { }; constexpr size_t kQ4gswBicolRoute = 0; -constexpr size_t kQ4gswBk64Route = 1; -constexpr size_t kQ4gswPrefillRoute = 2; -// 2-route (bicol + prefill) layout, used when the BK64 route is not recorded: -// the prefill dispatch sits at index 1, not kQ4gswPrefillRoute (the 3-route 2). -constexpr size_t kQ4gswPrefillRoute2Way = 1; Q4gswExecutionState make_q4gsw_execution_state( WGPUDevice device, @@ -237,6 +259,8 @@ Q4gswExecutionState make_q4gsw_execution_state( uint32_t wg_size, bool use_single_gemv, bool use_dual_route, + bool m3_eligible, + bool record_m3_route, bool record_bk64_route, bool bk64_eligible, bool prefill_use_steel, @@ -259,15 +283,17 @@ Q4gswExecutionState make_q4gsw_execution_state( } const uint32_t m = static_cast(live_m); const bool use_gemv = use_single_gemv || (use_dual_route && m == 1u); - const bool use_bk64 = !use_gemv && bk64_eligible && + const bool use_m3_shared = !use_gemv && m3_eligible && m == 3u; + const bool use_bk64 = !use_gemv && !use_m3_shared && bk64_eligible && utils::is_q4gsw_bk64_live_m(m) && steel_bk64_workgroup_count(device, m, N, K) > 0u; const uint32_t workgroup_count = compute_q4gsw_workgroup_count( device, use_gemv, + use_m3_shared, use_bk64, - !use_gemv && !use_bk64 && prefill_use_steel, - !use_gemv && !use_bk64 && prefill_use_shmem_gemm, + !use_gemv && !use_m3_shared && !use_bk64 && prefill_use_steel, + !use_gemv && !use_m3_shared && !use_bk64 && prefill_use_shmem_gemm, m, N, K, @@ -284,12 +310,16 @@ Q4gswExecutionState make_q4gsw_execution_state( state.params.has_bias = has_bias; state.output_dims = input_dims; state.output_dims.back() = static_cast(N); - state.active_route = use_dual_route - ? (use_gemv ? kQ4gswBicolRoute - : (record_bk64_route - ? (use_bk64 ? kQ4gswBk64Route : kQ4gswPrefillRoute) - : kQ4gswPrefillRoute2Way)) - : 0u; + if (!use_dual_route || use_gemv) { + state.active_route = kQ4gswBicolRoute; + } else if (record_m3_route && use_m3_shared) { + state.active_route = 1u; + } else { + const size_t bk64_route = record_m3_route ? 2u : 1u; + state.active_route = record_bk64_route && use_bk64 + ? bk64_route + : bk64_route + (record_bk64_route ? 1u : 0u); + } state.active_grid = {workgroup_count, 1u}; return state; } @@ -307,6 +337,8 @@ struct Q4gswResizeContext { uint32_t wg_size; bool use_single_gemv; bool use_dual_route; + bool m3_eligible; + bool record_m3_route; bool record_bk64_route; bool bk64_eligible; bool prefill_use_steel; @@ -330,6 +362,8 @@ void resize_q4gsw(WebGPUGraph& graph, const Q4gswResizeContext& context) { context.wg_size, context.use_single_gemv, context.use_dual_route, + context.m3_eligible, + context.record_m3_route, context.record_bk64_route, context.bk64_eligible, context.prefill_use_steel, @@ -456,16 +490,19 @@ void q4gsw_linear_impl_with_input_buffer_internal( const uint32_t wg_size = utils::clamp_workgroup_size( device, get_webgpu_shader_info(kQ4gswLinearShader).workgroup_size_x); const bool bicol_eligible = K % 8u == 0u && gs % 8u == 0u; + const bool m3_eligible = bicol_eligible && m3_shared_supported(device); const bool use_gemv = M == 1u && bicol_eligible; + const bool use_m3_shared = !use_gemv && M == 3u && m3_eligible; const bool use_dual_route = utils::should_record_q4gsw_dual_route( M, bicol_eligible, graph.has_dynamic_shapes(), graph.config().record_q4gsw_decode_route); + const bool record_m3_route = use_dual_route && M >= 3u && m3_eligible; const bool bk64_eligible = steel_bk64_eligible(device, K, N, gs, has_bias != 0u); const bool record_bk64_route = use_dual_route && bk64_eligible && M >= 128u; - const bool use_bk64 = !use_gemv && bk64_eligible && + const bool use_bk64 = !use_gemv && !use_m3_shared && bk64_eligible && utils::is_q4gsw_bk64_live_m(M) && steel_bk64_workgroup_count(device, M, N, K) > 0u; // GEMV (bicol) is a pow2 tree reduction; compute its size only when used. @@ -474,13 +511,13 @@ void q4gsw_linear_impl_with_input_buffer_internal( device, get_webgpu_shader_info(kQ4gswBicolShader).workgroup_size_x) : 0u; // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible. - const bool use_steel = !use_gemv && steel_supported(device) && - steel_workgroup_count(device, M, N, K) > 0u; + const bool use_steel = !use_gemv && !use_m3_shared && + steel_supported(device) && steel_workgroup_count(device, M, N, K) > 0u; // shmem GEMM is now a FALLBACK, not dead: steel shadows it whenever eligible, // so shmem only wins when steel is ineligible (K % 16 != 0, or a // <256-invocation device such as SwiftShader) and the shape still hits the // large K/N thresholds; otherwise the register-tiled path handles it. - const bool use_shmem_gemm = !use_gemv && !use_steel && + const bool use_shmem_gemm = !use_gemv && !use_m3_shared && !use_steel && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); const char* prefill_shader_name = use_steel ? kQ4gswSteelShader : use_shmem_gemm ? kQ4gswShmemShader @@ -535,6 +572,8 @@ void q4gsw_linear_impl_with_input_buffer_internal( wg_size, use_gemv, use_dual_route, + m3_eligible, + record_m3_route, record_bk64_route, bk64_eligible, use_steel, @@ -590,6 +629,17 @@ void q4gsw_linear_impl_with_input_buffer_internal( initial_state.active_grid.x, "linear_q4gsw_coop4_bicol", initial_state.active_grid.y}); + size_t m3_idx = 0; + if (record_m3_route) { + utils::ComputePipelineBundle m3_bundle = make_shared_bundle( + kQ4gswLinearM3SharedBicolWGSL, bicol_bundle, true, 0u); + m3_idx = graph.add_dispatch( + {m3_bundle.pipeline, + m3_bundle.bind_group, + initial_state.active_grid.x, + "linear_q4gsw_m3_shared_bicol", + initial_state.active_grid.y}); + } size_t bk64_idx = 0; if (record_bk64_route) { utils::ComputePipelineBundle bk64_bundle = make_shared_bundle( @@ -615,7 +665,18 @@ void q4gsw_linear_impl_with_input_buffer_internal( initial_state.active_grid.x, prefill_label, initial_state.active_grid.y}); - if (record_bk64_route) { + if (record_m3_route && record_bk64_route) { + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, + {m3_idx, m3_idx + 1}, + {bk64_idx, bk64_idx + 1}, + {prefill_idx, prefill_idx + 1}}); + } else if (record_m3_route) { + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, + {m3_idx, m3_idx + 1}, + {prefill_idx, prefill_idx + 1}}); + } else if (record_bk64_route) { route_group = graph.register_dispatch_route_group( {{bicol_idx, bicol_idx + 1}, {bk64_idx, bk64_idx + 1}, @@ -627,17 +688,21 @@ void q4gsw_linear_impl_with_input_buffer_internal( graph.select_dispatch_route( route_group, initial_state.active_route, {initial_state.active_grid}); } else { - const bool fixed_wg = use_gemv ? false : (use_bk64 || fixed_prefill_wg); - utils::ComputePipelineBundle bundle = make_bundle( - get_webgpu_shader_info(shader_name).source, - fixed_wg, - use_gemv ? gemv_wg_size : wg_size); + const bool fixed_wg = + use_m3_shared || use_bk64 || (!use_gemv && fixed_prefill_wg); + const char* shader_source = use_m3_shared + ? kQ4gswLinearM3SharedBicolWGSL + : get_webgpu_shader_info(shader_name).source; + utils::ComputePipelineBundle bundle = + make_bundle(shader_source, fixed_wg, use_gemv ? gemv_wg_size : wg_size); dispatch_idx = graph.add_dispatch( {bundle.pipeline, bundle.bind_group, initial_state.active_grid.x, use_gemv ? "linear_q4gsw_coop4_bicol" - : (use_bk64 ? "linear_q4gsw_bk64" : prefill_label), + : (use_m3_shared + ? "linear_q4gsw_m3_shared_bicol" + : (use_bk64 ? "linear_q4gsw_bk64" : prefill_label)), initial_state.active_grid.y}); } @@ -656,6 +721,8 @@ void q4gsw_linear_impl_with_input_buffer_internal( wg_size, use_gemv, use_dual_route, + m3_eligible, + record_m3_route, record_bk64_route, bk64_eligible, use_steel, diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol.wgsl new file mode 100644 index 00000000000..81894b975fb --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol.wgsl @@ -0,0 +1,159 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(5) var params: Params; + +const WG: u32 = 64u; +var partial0: array; +var partial1: array; +var partial2: array; +var partial3: array; +var partial4: array; +var partial5: array; + +@compute @workgroup_size(WG, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(num_workgroups) ngrp: vec3, + @builtin(local_invocation_id) lid: vec3) { + if (params.M != 3u) { + return; + } + let num_pairs = (params.N + 1u) >> 1u; + let num_words = params.K >> 3u; + let row_words = params.K_packed >> 2u; + var pair = wid.x; + loop { + if (pair >= num_pairs) { + break; + } + let col0 = pair << 1u; + let col1 = col0 + 1u; + let has1 = col1 < params.N; + let wbase0 = col0 * row_words; + let wbase1 = col1 * row_words; + var acc0: f32 = 0.0; + var acc1: f32 = 0.0; + var acc2: f32 = 0.0; + var acc3: f32 = 0.0; + var acc4: f32 = 0.0; + var acc5: f32 = 0.0; + var w: u32 = lid.x; + loop { + if (w >= num_words) { + break; + } + let k0 = w << 3u; + let scale_row = (k0 / params.group_size) * params.padded_N; + let word0 = t_weight[wbase0 + w]; + let scale0 = t_scales[scale_row + col0]; + var word1: u32 = 0u; + var scale1: f32 = 0.0; + if (has1) { + word1 = t_weight[wbase1 + w]; + scale1 = t_scales[scale_row + col1]; + } + for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) { + let kk = bi << 1u; + let input00 = t_input[k0 + kk]; + let input01 = t_input[k0 + kk + 1u]; + let input10 = t_input[params.K + k0 + kk]; + let input11 = t_input[params.K + k0 + kk + 1u]; + let input20 = t_input[2u * params.K + k0 + kk]; + let input21 = t_input[2u * params.K + k0 + kk + 1u]; + let byte0 = (word0 >> (bi * 8u)) & 0xFFu; + let lo0 = f32(i32(byte0 & 0x0Fu) - 8); + let hi0 = f32(i32((byte0 >> 4u) & 0x0Fu) - 8); + acc0 = acc0 + input00 * lo0 * scale0; + acc0 = acc0 + input01 * hi0 * scale0; + acc2 = acc2 + input10 * lo0 * scale0; + acc2 = acc2 + input11 * hi0 * scale0; + acc4 = acc4 + input20 * lo0 * scale0; + acc4 = acc4 + input21 * hi0 * scale0; + let byte1 = (word1 >> (bi * 8u)) & 0xFFu; + let lo1 = f32(i32(byte1 & 0x0Fu) - 8); + let hi1 = f32(i32((byte1 >> 4u) & 0x0Fu) - 8); + acc1 = acc1 + input00 * lo1 * scale1; + acc1 = acc1 + input01 * hi1 * scale1; + acc3 = acc3 + input10 * lo1 * scale1; + acc3 = acc3 + input11 * hi1 * scale1; + acc5 = acc5 + input20 * lo1 * scale1; + acc5 = acc5 + input21 * hi1 * scale1; + } + w = w + WG; + } + + partial0[lid.x] = acc0; + partial1[lid.x] = acc1; + partial2[lid.x] = acc2; + partial3[lid.x] = acc3; + partial4[lid.x] = acc4; + partial5[lid.x] = acc5; + workgroupBarrier(); + var stride: u32 = WG >> 1u; + loop { + if (stride == 0u) { + break; + } + if (lid.x < stride) { + partial0[lid.x] = partial0[lid.x] + partial0[lid.x + stride]; + partial1[lid.x] = partial1[lid.x] + partial1[lid.x + stride]; + partial2[lid.x] = partial2[lid.x] + partial2[lid.x + stride]; + partial3[lid.x] = partial3[lid.x] + partial3[lid.x + stride]; + partial4[lid.x] = partial4[lid.x] + partial4[lid.x + stride]; + partial5[lid.x] = partial5[lid.x] + partial5[lid.x + stride]; + } + workgroupBarrier(); + stride = stride >> 1u; + } + if (lid.x == 0u) { + var out0 = partial0[0]; + var out1 = partial1[0]; + var out2 = partial2[0]; + var out3 = partial3[0]; + var out4 = partial4[0]; + var out5 = partial5[0]; + if (params.has_bias != 0u) { + let bias0 = t_bias[col0]; + out0 = out0 + bias0; + out2 = out2 + bias0; + out4 = out4 + bias0; + if (has1) { + let bias1 = t_bias[col1]; + out1 = out1 + bias1; + out3 = out3 + bias1; + out5 = out5 + bias1; + } + } + t_out[col0] = out0; + t_out[params.N + col0] = out2; + t_out[2u * params.N + col0] = out4; + if (has1) { + t_out[col1] = out1; + t_out[params.N + col1] = out3; + t_out[2u * params.N + col1] = out5; + } + } + workgroupBarrier(); + pair = pair + ngrp.x; + } +} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol_wgsl.h new file mode 100644 index 00000000000..0f49852b0b7 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_m3_shared_bicol_wgsl.h @@ -0,0 +1,183 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_linear_m3_shared_bicol.wgsl - DO NOT EDIT. +// wgsl-sha256: 7a670ca4277c4a181dbf4590353f29e0478130b3018a5f25c65b153735f2a3c0 +inline constexpr const char* kQ4gswLinearM3SharedBicolWGSL = R"( +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(5) var params: Params; + +const WG: u32 = 64u; +var partial0: array; +var partial1: array; +var partial2: array; +var partial3: array; +var partial4: array; +var partial5: array; + +@compute @workgroup_size(WG, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(num_workgroups) ngrp: vec3, + @builtin(local_invocation_id) lid: vec3) { + if (params.M != 3u) { + return; + } + let num_pairs = (params.N + 1u) >> 1u; + let num_words = params.K >> 3u; + let row_words = params.K_packed >> 2u; + var pair = wid.x; + loop { + if (pair >= num_pairs) { + break; + } + let col0 = pair << 1u; + let col1 = col0 + 1u; + let has1 = col1 < params.N; + let wbase0 = col0 * row_words; + let wbase1 = col1 * row_words; + var acc0: f32 = 0.0; + var acc1: f32 = 0.0; + var acc2: f32 = 0.0; + var acc3: f32 = 0.0; + var acc4: f32 = 0.0; + var acc5: f32 = 0.0; + var w: u32 = lid.x; + loop { + if (w >= num_words) { + break; + } + let k0 = w << 3u; + let scale_row = (k0 / params.group_size) * params.padded_N; + let word0 = t_weight[wbase0 + w]; + let scale0 = t_scales[scale_row + col0]; + var word1: u32 = 0u; + var scale1: f32 = 0.0; + if (has1) { + word1 = t_weight[wbase1 + w]; + scale1 = t_scales[scale_row + col1]; + } + for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) { + let kk = bi << 1u; + let input00 = t_input[k0 + kk]; + let input01 = t_input[k0 + kk + 1u]; + let input10 = t_input[params.K + k0 + kk]; + let input11 = t_input[params.K + k0 + kk + 1u]; + let input20 = t_input[2u * params.K + k0 + kk]; + let input21 = t_input[2u * params.K + k0 + kk + 1u]; + let byte0 = (word0 >> (bi * 8u)) & 0xFFu; + let lo0 = f32(i32(byte0 & 0x0Fu) - 8); + let hi0 = f32(i32((byte0 >> 4u) & 0x0Fu) - 8); + acc0 = acc0 + input00 * lo0 * scale0; + acc0 = acc0 + input01 * hi0 * scale0; + acc2 = acc2 + input10 * lo0 * scale0; + acc2 = acc2 + input11 * hi0 * scale0; + acc4 = acc4 + input20 * lo0 * scale0; + acc4 = acc4 + input21 * hi0 * scale0; + let byte1 = (word1 >> (bi * 8u)) & 0xFFu; + let lo1 = f32(i32(byte1 & 0x0Fu) - 8); + let hi1 = f32(i32((byte1 >> 4u) & 0x0Fu) - 8); + acc1 = acc1 + input00 * lo1 * scale1; + acc1 = acc1 + input01 * hi1 * scale1; + acc3 = acc3 + input10 * lo1 * scale1; + acc3 = acc3 + input11 * hi1 * scale1; + acc5 = acc5 + input20 * lo1 * scale1; + acc5 = acc5 + input21 * hi1 * scale1; + } + w = w + WG; + } + + partial0[lid.x] = acc0; + partial1[lid.x] = acc1; + partial2[lid.x] = acc2; + partial3[lid.x] = acc3; + partial4[lid.x] = acc4; + partial5[lid.x] = acc5; + workgroupBarrier(); + var stride: u32 = WG >> 1u; + loop { + if (stride == 0u) { + break; + } + if (lid.x < stride) { + partial0[lid.x] = partial0[lid.x] + partial0[lid.x + stride]; + partial1[lid.x] = partial1[lid.x] + partial1[lid.x + stride]; + partial2[lid.x] = partial2[lid.x] + partial2[lid.x + stride]; + partial3[lid.x] = partial3[lid.x] + partial3[lid.x + stride]; + partial4[lid.x] = partial4[lid.x] + partial4[lid.x + stride]; + partial5[lid.x] = partial5[lid.x] + partial5[lid.x + stride]; + } + workgroupBarrier(); + stride = stride >> 1u; + } + if (lid.x == 0u) { + var out0 = partial0[0]; + var out1 = partial1[0]; + var out2 = partial2[0]; + var out3 = partial3[0]; + var out4 = partial4[0]; + var out5 = partial5[0]; + if (params.has_bias != 0u) { + let bias0 = t_bias[col0]; + out0 = out0 + bias0; + out2 = out2 + bias0; + out4 = out4 + bias0; + if (has1) { + let bias1 = t_bias[col1]; + out1 = out1 + bias1; + out3 = out3 + bias1; + out5 = out5 + bias1; + } + } + t_out[col0] = out0; + t_out[params.N + col0] = out2; + t_out[2u * params.N + col0] = out4; + if (has1) { + t_out[col1] = out1; + t_out[params.N + col1] = out3; + t_out[2u * params.N + col1] = out5; + } + } + workgroupBarrier(); + pair = pair + ngrp.x; + } +} +)"; + +inline constexpr uint32_t kQ4gswLinearM3SharedBicolWorkgroupSizeX = 64; +inline constexpr uint32_t kQ4gswLinearM3SharedBicolWorkgroupSizeY = 1; +inline constexpr uint32_t kQ4gswLinearM3SharedBicolWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/scatter/Scatter.cpp b/backends/webgpu/runtime/ops/scatter/Scatter.cpp new file mode 100644 index 00000000000..8b068ba8280 --- /dev/null +++ b/backends/webgpu/runtime/ops/scatter/Scatter.cpp @@ -0,0 +1,124 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +constexpr int64_t kVocabSize = 262144; +constexpr int64_t kSelectedCount = 4096; + +void scatter_impl( + WebGPUGraph& graph, + const std::vector& args, + bool unique_indices) { + if (args.size() != 5u) { + throw std::runtime_error("WebGPU scatter: malformed argument list"); + } + const int input_id = args[0]; + const int dim_id = args[1]; + const int index_id = args[2]; + const int source_id = args[3]; + const int output_id = args[4]; + + if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int || + graph.get_int(dim_id) != -1) { + throw std::runtime_error("WebGPU scatter: requires dim=-1"); + } + const auto& input = graph.get_tensor(input_id); + const auto& index = graph.get_tensor(index_id); + const auto& source = graph.get_tensor(source_id); + const auto& output = graph.get_tensor(output_id); + const std::vector vocab_dims = {1, 1, kVocabSize}; + const std::vector selected_dims = {1, 1, kSelectedCount}; + if (input.buffer == nullptr || input.is_int || + input.elem_size != sizeof(float) || input.dims != vocab_dims || + input.nbytes != kVocabSize * sizeof(float)) { + throw std::runtime_error("WebGPU scatter: input must be fp32 [1,1,262144]"); + } + if (index.buffer == nullptr || !index.is_int || + index.elem_size != sizeof(int32_t) || index.dims != selected_dims || + index.nbytes != kSelectedCount * sizeof(int32_t)) { + throw std::runtime_error( + "WebGPU scatter: index must use effective i32 [1,1,4096] storage"); + } + if (source.buffer == nullptr || source.is_int || + source.elem_size != sizeof(float) || source.dims != selected_dims || + source.nbytes != kSelectedCount * sizeof(float)) { + throw std::runtime_error("WebGPU scatter: source must be fp32 [1,1,4096]"); + } + if (output.buffer == nullptr || output.is_int || + output.elem_size != sizeof(float) || output.dims != vocab_dims || + output.nbytes != kVocabSize * sizeof(float)) { + throw std::runtime_error( + "WebGPU scatter: output must be fp32 [1,1,262144]"); + } + + add_flat_copy(graph, input_id, output_id); + const char* shader = + unique_indices ? kScatterUniqueIndicesWGSL : kScatterWGSL; + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + graph.device(), + shader, + { + {0, WGPUBufferBindingType_Storage, output.buffer, output.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + index.buffer, + index.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + source.buffer, + source.nbytes}, + }); + const uint32_t workgroups = unique_indices + ? static_cast(utils::div_up( + kSelectedCount, kScatterUniqueIndicesWorkgroupSizeX)) + : 1u; + try { + graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroups, + unique_indices ? "scatter_unique_indices" : "scatter"}); + } catch (...) { + wgpuComputePipelineRelease(bundle.pipeline); + wgpuBindGroupRelease(bundle.bind_group); + throw; + } +} + +void scatter_generic_impl(WebGPUGraph& graph, const std::vector& args) { + scatter_impl(graph, args, false); +} + +void scatter_unique_impl(WebGPUGraph& graph, const std::vector& args) { + scatter_impl(graph, args, true); +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.scatter.src, scatter_generic_impl); + WEBGPU_REGISTER_OP(et_vk.scatter_src_unique.default, scatter_unique_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/scatter/scatter.wgsl b/backends/webgpu/runtime/ops/scatter/scatter.wgsl new file mode 100644 index 00000000000..a860d96f749 --- /dev/null +++ b/backends/webgpu/runtime/ops/scatter/scatter.wgsl @@ -0,0 +1,19 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var output: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var source: array; + +@compute @workgroup_size(1) +fn main() { + for (var i = 0u; i < 4096u; i += 1u) { + let destination = indices[i]; + if (destination >= 0 && destination < 262144) { + output[u32(destination)] = source[i]; + } + } +} diff --git a/backends/webgpu/runtime/ops/scatter/scatter_unique_indices.wgsl b/backends/webgpu/runtime/ops/scatter/scatter_unique_indices.wgsl new file mode 100644 index 00000000000..26b8b7983f4 --- /dev/null +++ b/backends/webgpu/runtime/ops/scatter/scatter_unique_indices.wgsl @@ -0,0 +1,35 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var output: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var source: array; + +// One invocation per SOURCE element, replacing the shipping scatter.wgsl's +// single one-thread workgroup that walked all 4096 writes serially +// (measured 393.216 us = 6 timestamp quanta per dispatch). +// +// The shipping kernel is last-write-wins over ascending i, mirroring portable +// ExecuTorch op_scatter.cpp. This kernel is arbitrary-write-wins. The two are +// bit-identical iff the destinations are pairwise DISTINCT, which they are by +// construction here: `indices` is `token_ordering[topk_indices]`, a gather of +// 32 DISTINCT rows out of a [2048,128] permutation of [0,262144), so no two +// source elements can name the same destination. The out-of-range guard is +// kept verbatim, and is per-invocation, so negative / >= vocab entries are +// dropped exactly as before. +const WG: u32 = 64u; + +@compute @workgroup_size(WG, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= 4096u) { + return; + } + let destination = indices[i]; + if (destination >= 0 && destination < 262144) { + output[u32(destination)] = source[i]; + } +} diff --git a/backends/webgpu/runtime/ops/scatter/scatter_unique_indices_wgsl.h b/backends/webgpu/runtime/ops/scatter/scatter_unique_indices_wgsl.h new file mode 100644 index 00000000000..17c81e39f8f --- /dev/null +++ b/backends/webgpu/runtime/ops/scatter/scatter_unique_indices_wgsl.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from scatter_unique_indices.wgsl - DO NOT EDIT. +// wgsl-sha256: 1db5230fb370bd30c52e857ca7134a5dfb6e04c381e3a2670f4fdcfe1d625d39 +inline constexpr const char* kScatterUniqueIndicesWGSL = R"( +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var output: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var source: array; + +// One invocation per SOURCE element, replacing the shipping scatter.wgsl's +// single one-thread workgroup that walked all 4096 writes serially +// (measured 393.216 us = 6 timestamp quanta per dispatch). +// +// The shipping kernel is last-write-wins over ascending i, mirroring portable +// ExecuTorch op_scatter.cpp. This kernel is arbitrary-write-wins. The two are +// bit-identical iff the destinations are pairwise DISTINCT, which they are by +// construction here: `indices` is `token_ordering[topk_indices]`, a gather of +// 32 DISTINCT rows out of a [2048,128] permutation of [0,262144), so no two +// source elements can name the same destination. The out-of-range guard is +// kept verbatim, and is per-invocation, so negative / >= vocab entries are +// dropped exactly as before. +const WG: u32 = 64u; + +@compute @workgroup_size(WG, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= 4096u) { + return; + } + let destination = indices[i]; + if (destination >= 0 && destination < 262144) { + output[u32(destination)] = source[i]; + } +} +)"; + +inline constexpr uint32_t kScatterUniqueIndicesWorkgroupSizeX = 64; +inline constexpr uint32_t kScatterUniqueIndicesWorkgroupSizeY = 1; +inline constexpr uint32_t kScatterUniqueIndicesWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/scatter/scatter_wgsl.h b/backends/webgpu/runtime/ops/scatter/scatter_wgsl.h new file mode 100644 index 00000000000..00c3fbc2084 --- /dev/null +++ b/backends/webgpu/runtime/ops/scatter/scatter_wgsl.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from scatter.wgsl - DO NOT EDIT. +// wgsl-sha256: ac43cd32373a4c93649c8174169b87d0c2c4eaf7be9f318a3c97273f5a1682bb +inline constexpr const char* kScatterWGSL = R"( +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var output: array; +@group(0) @binding(1) var indices: array; +@group(0) @binding(2) var source: array; + +@compute @workgroup_size(1) +fn main() { + for (var i = 0u; i < 4096u; i += 1u) { + let destination = indices[i]; + if (destination >= 0 && destination < 262144) { + output[u32(destination)] = source[i]; + } + } +} +)"; + +inline constexpr uint32_t kScatterWorkgroupSizeX = 1; +inline constexpr uint32_t kScatterWorkgroupSizeY = 1; +inline constexpr uint32_t kScatterWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp b/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp index 75542377f69..9a7a6d54f94 100644 --- a/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp +++ b/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp @@ -92,7 +92,7 @@ void split_with_sizes_impl(WebGPUGraph& graph, const std::vector& args) { uint32_t wg_size = utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "split_with_sizes"); WGPUConstantEntry wg_size_constant = {}; @@ -130,7 +130,12 @@ void split_with_sizes_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "split_with_sizes_copy", + workgroup_count.y}); graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp index 7b96bf0c9ac..f7685c0c7b9 100644 --- a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp @@ -10,11 +10,14 @@ #include #include #include +#include #include +#include #include #include +#include #include #include @@ -30,17 +33,40 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { WGPUDevice device = graph.device(); - float alpha = 1.0f; - if (graph.get_value_type(alpha_id) == WebGPUGraph::ValueType::Int) { - alpha = static_cast(graph.get_int(alpha_id)); - } else if (graph.get_value_type(alpha_id) == WebGPUGraph::ValueType::Double) { - alpha = static_cast(graph.get_double(alpha_id)); + const auto alpha_type = graph.get_value_type(alpha_id); + int64_t alpha_int64 = 1; + float alpha_fp32 = 1.0f; + if (alpha_type == WebGPUGraph::ValueType::Int) { + alpha_int64 = graph.get_int(alpha_id); + alpha_fp32 = static_cast(alpha_int64); + } else if (alpha_type == WebGPUGraph::ValueType::Double) { + alpha_fp32 = static_cast(graph.get_double(alpha_id)); + } else { + throw std::runtime_error("sub: alpha must be an int or double scalar"); } const auto& in1_tensor = graph.get_tensor(in1_id); const auto& in2_tensor = graph.get_tensor(in2_id); const auto& out_tensor = graph.get_tensor(out_id); + const bool all_int32 = + classify_sub_storage(in1_tensor, in2_tensor, out_tensor) == + SubStorage::Int32; + if (all_int32 && + (alpha_type != WebGPUGraph::ValueType::Int || + alpha_int64 < std::numeric_limits::min() || + alpha_int64 > std::numeric_limits::max())) { + throw std::runtime_error("sub: int32 operands require an int32 alpha"); + } + + int32_t rhs = 0; + if (all_int32 && alpha_int64 == 1 && out_tensor.dims == in1_tensor.dims && + in2_tensor.is_int && in2_tensor.elem_size == sizeof(int32_t) && + in2_tensor.nbytes == sizeof(int32_t) && + graph.try_read_prepacked_int32_scalar(in2_id, rhs)) { + graph.add_input_bias_projection(out_id, in1_id, -static_cast(rhs)); + } + // Rank guard (NCHW backend is <= 4 dims; 1D dispatch only). if (out_tensor.dims.size() > kTensorMetaMaxNdim || in1_tensor.dims.size() > kTensorMetaMaxNdim || @@ -58,26 +84,27 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { fill_tensor_meta_broadcast(in1_tensor, out_ndim, &in1_meta); fill_tensor_meta_broadcast(in2_tensor, out_ndim, &in2_meta); - // fp32-only: nbytes must equal numel * 4 for every operand. - if (out_tensor.nbytes != - static_cast(out_meta.numel) * sizeof(float) || - in1_tensor.nbytes != - static_cast(in1_meta.numel) * sizeof(float) || - in2_tensor.nbytes != - static_cast(in2_meta.numel) * sizeof(float)) { - throw std::runtime_error("sub: non-fp32 operand (nbytes != numel * 4)"); + const size_t element_size = all_int32 ? sizeof(int32_t) : sizeof(float); + if (out_tensor.nbytes != static_cast(out_meta.numel) * element_size || + in1_tensor.nbytes != static_cast(in1_meta.numel) * element_size || + in2_tensor.nbytes != static_cast(in2_meta.numel) * element_size) { + throw std::runtime_error("sub: operand byte size does not match dtype"); } - uint32_t wg_size = - utils::clamp_workgroup_size(device, kBinarySubWorkgroupSizeX); - utils::WgCount workgroup_count = - utils::compute_2d_workgroup_count(device, out_meta.numel, wg_size, "sub"); + const uint32_t default_wg_size = + all_int32 ? kBinarySubInt32WorkgroupSizeX : kBinarySubWorkgroupSizeX; + const char* shader = all_int32 ? kBinarySubInt32WGSL : kBinarySubWGSL; + const char* kernel_name = all_int32 ? "sub_int32" : "sub"; + uint32_t wg_size = utils::clamp_workgroup_size(device, default_wg_size); + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, out_meta.numel, wg_size, kernel_name); WGPUConstantEntry constants[2] = {}; constants[0].key = {"wg_size", WGPU_STRLEN}; constants[0].value = static_cast(wg_size); constants[1].key = {"alpha", WGPU_STRLEN}; - constants[1].value = static_cast(alpha); + constants[1].value = all_int32 ? static_cast(alpha_int64) + : static_cast(alpha_fp32); WGPUBuffer out_meta_buf = utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); @@ -89,7 +116,7 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, - kBinarySubWGSL, + shader, { {0, WGPUBufferBindingType_ReadOnlyStorage, @@ -114,45 +141,51 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { {bundle.pipeline, bundle.bind_group, workgroup_count.x, - "sub", + kernel_name, workgroup_count.y}); // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; - auto sub_resize = - [in1_id, in2_id, out_id, wg_size, dispatch_idx, o_buf, a_buf, b_buf]( - WebGPUGraph& g) { - const auto& a = g.cur_dims(in1_id); - const auto& b = g.cur_dims(in2_id); - const size_t r = std::max(a.size(), b.size()); - std::vector out_d(r, 1); - for (size_t i = 0; i < r; i++) { - const int64_t av = (i + a.size() < r) ? 1 : a[i - (r - a.size())]; - const int64_t bv = (i + b.size() < r) ? 1 : b[i - (r - b.size())]; - if (av != bv && av != 1 && bv != 1) { - throw std::runtime_error( - "sub(resize): operands are not broadcast-compatible"); - } - out_d[i] = av > bv ? av : bv; - } - g.set_cur_dims(out_id, out_d); - const uint32_t out_ndim = static_cast(r); - WebGPUTensor ta, tb, to; - ta.dims = a; - tb.dims = b; - to.dims = out_d; - TensorMeta om, am, bm; - fill_tensor_meta_broadcast(to, out_ndim, &om); - fill_tensor_meta_broadcast(ta, out_ndim, &am); - fill_tensor_meta_broadcast(tb, out_ndim, &bm); - wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om)); - wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); - wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm)); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), om.numel, wg_size, "sub(resize)"); - g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; - g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; - }; + auto sub_resize = [in1_id, + in2_id, + out_id, + wg_size, + dispatch_idx, + o_buf, + a_buf, + b_buf, + kernel_name](WebGPUGraph& g) { + const auto& a = g.cur_dims(in1_id); + const auto& b = g.cur_dims(in2_id); + const size_t r = std::max(a.size(), b.size()); + std::vector out_d(r, 1); + for (size_t i = 0; i < r; i++) { + const int64_t av = (i + a.size() < r) ? 1 : a[i - (r - a.size())]; + const int64_t bv = (i + b.size() < r) ? 1 : b[i - (r - b.size())]; + if (av != bv && av != 1 && bv != 1) { + throw std::runtime_error( + "sub(resize): operands are not broadcast-compatible"); + } + out_d[i] = av > bv ? av : bv; + } + g.set_cur_dims(out_id, out_d); + const uint32_t out_ndim = static_cast(r); + WebGPUTensor ta, tb, to; + ta.dims = a; + tb.dims = b; + to.dims = out_d; + TensorMeta om, am, bm; + fill_tensor_meta_broadcast(to, out_ndim, &om); + fill_tensor_meta_broadcast(ta, out_ndim, &am); + fill_tensor_meta_broadcast(tb, out_ndim, &bm); + wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om)); + wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); + wgpuQueueWriteBuffer(g.queue(), b_buf, 0, &bm, sizeof(bm)); + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), om.numel, wg_size, kernel_name); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }; graph.add_tensor_resize_hook(in1_id, sub_resize); graph.add_tensor_resize_hook(in2_id, sub_resize); diff --git a/backends/webgpu/runtime/ops/sub/SubStorage.h b/backends/webgpu/runtime/ops/sub/SubStorage.h new file mode 100644 index 00000000000..adac8107b17 --- /dev/null +++ b/backends/webgpu/runtime/ops/sub/SubStorage.h @@ -0,0 +1,34 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace executorch::backends::webgpu { + +enum class SubStorage { Float32, Int32 }; + +template +SubStorage +classify_sub_storage(const Tensor& in1, const Tensor& in2, const Tensor& out) { + const bool all_int32 = in1.is_int && !in1.is_bool && + in1.elem_size == sizeof(int32_t) && in2.is_int && !in2.is_bool && + in2.elem_size == sizeof(int32_t) && out.is_int && !out.is_bool && + out.elem_size == sizeof(int32_t); + const bool any_integer = in1.is_int || in1.is_bool || in2.is_int || + in2.is_bool || out.is_int || out.is_bool; + if (any_integer && !all_int32) { + throw std::runtime_error( + "sub: integer operands must all be signed int32 tensors"); + } + return all_int32 ? SubStorage::Int32 : SubStorage::Float32; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp index 54d80016296..0c6a7914c76 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -60,8 +60,8 @@ void add_convert_op( uint32_t num_elements = static_cast(out_tensor.nbytes / 4); uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); - uint32_t workgroup_count = - utils::compute_1d_workgroup_count(device, num_elements, wg_size, op_name); + const utils::WgCount workgroup_count = + utils::compute_2d_workgroup_count(device, num_elements, wg_size, op_name); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; @@ -92,8 +92,12 @@ void add_convert_op( &wg_size_constant, 1); - const size_t dispatch_idx = - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + op_name, + workgroup_count.y}); // Dynamic shapes: recompute num_elements/dispatch for the live shape. WGPUBuffer params_buf = uniform_buffer; @@ -106,12 +110,13 @@ void add_convert_op( ConvertParams p = {}; p.num_elements = static_cast(numel); wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - g.device(), - static_cast(numel), - wg_size, - "to_copy(resize)"); + const utils::WgCount workgroups = utils::compute_2d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = workgroups.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = workgroups.y; }); } @@ -146,7 +151,7 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { const uint32_t wg_size = utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); - const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, num_elements, wg_size, "to_copy_bool_to_float"); WGPUConstantEntry wg_size_constant = {}; @@ -177,8 +182,12 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { &wg_size_constant, 1); - const size_t dispatch_idx = - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "to_copy_bool_to_float", + workgroup_count.y}); WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( @@ -194,12 +203,13 @@ void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { ConvertParams p = {}; p.num_elements = static_cast(numel); wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - g.device(), - static_cast(numel), - wg_size, - "to_copy_bool_to_float(resize)"); + const utils::WgCount workgroups = utils::compute_2d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy_bool_to_float(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = workgroups.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = workgroups.y; }); } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl index 239730de65d..d039e1d513f 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl @@ -12,8 +12,10 @@ struct Params { @group(0) @binding(2) var params: Params; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h index ef7e40976f3..36d343e5638 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from to_copy_bool_to_float.wgsl - DO NOT EDIT. -// wgsl-sha256: 29fd43b2f638489e9b8d72b2cc9140d07174c750cbdc291e63793319a2fa5961 +// wgsl-sha256: e8ad49a183e13cdef81e5ee374a12ca6f15b265f16eff1e33506b47f0e49522b inline constexpr const char* kToCopyBoolToFloatWGSL = R"( override wg_size: u32 = 256u; @@ -29,8 +29,10 @@ struct Params { @group(0) @binding(2) var params: Params; @compute @workgroup_size(wg_size) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl index f1113f0e14c..5ffb2507f27 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_convert.wgsl @@ -9,8 +9,10 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h index 1a384c747d8..11a266661aa 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_float_to_int_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from to_copy_convert.wgsl - DO NOT EDIT. -// wgsl-sha256: c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f +// wgsl-sha256: 241d51293095623126da4f106092bd2d7327c26e00c9cad39bdb8dc546f48149 inline constexpr const char* kToCopyFloatToIntWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -26,8 +26,10 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h index 6fdf37ec2b7..841b25a4c31 100644 --- a/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_int_to_float_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from to_copy_convert.wgsl - DO NOT EDIT. -// wgsl-sha256: e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195 +// wgsl-sha256: 9506570f98b9888a65603157f75d380f7784539205610ccbcc459d6afb6629c5 inline constexpr const char* kToCopyIntToFloatWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -26,8 +26,10 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/topk/TopK.cpp b/backends/webgpu/runtime/ops/topk/TopK.cpp new file mode 100644 index 00000000000..386d0c86bf2 --- /dev/null +++ b/backends/webgpu/runtime/ops/topk/TopK.cpp @@ -0,0 +1,106 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +constexpr int64_t kInputWidth = 2048; +constexpr int64_t kOutputWidth = 32; + +void topk_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() != 6u) { + throw std::runtime_error("WebGPU topk: malformed argument list"); + } + const int input_id = args[0]; + const int k_id = args[1]; + const int dim_id = args[2]; + const int largest_id = args[3]; + const int sorted_id = args[4]; + const int output_list_id = args[5]; + + if (graph.get_value_type(k_id) != WebGPUGraph::ValueType::Int || + graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int || + graph.get_value_type(largest_id) != WebGPUGraph::ValueType::Bool || + graph.get_value_type(sorted_id) != WebGPUGraph::ValueType::Bool || + graph.get_value_type(output_list_id) != + WebGPUGraph::ValueType::ValueList) { + throw std::runtime_error("WebGPU topk: malformed scalar arguments"); + } + if (graph.get_int(k_id) != kOutputWidth || graph.get_int(dim_id) != -1 || + !graph.get_bool(largest_id) || !graph.get_bool(sorted_id)) { + throw std::runtime_error( + "WebGPU topk: requires k=32, dim=-1, largest=true, sorted=true"); + } + + const auto& output_ids = graph.get_value_list(output_list_id); + if (output_ids.size() != 2u) { + throw std::runtime_error( + "WebGPU topk: expected values and indices outputs"); + } + const auto& input = graph.get_tensor(input_id); + const auto& values = graph.get_tensor(output_ids[0]); + const auto& indices = graph.get_tensor(output_ids[1]); + const std::vector input_dims = {1, 1, kInputWidth}; + const std::vector output_dims = {1, 1, kOutputWidth}; + if (input.buffer == nullptr || input.is_int || + input.elem_size != sizeof(float) || input.dims != input_dims || + input.nbytes != kInputWidth * sizeof(float)) { + throw std::runtime_error("WebGPU topk: input must be fp32 [1,1,2048]"); + } + if (values.buffer == nullptr || values.is_int || + values.elem_size != sizeof(float) || values.dims != output_dims || + values.nbytes != kOutputWidth * sizeof(float)) { + throw std::runtime_error("WebGPU topk: values must be fp32 [1,1,32]"); + } + if (indices.buffer == nullptr || !indices.is_int || + indices.elem_size != sizeof(int32_t) || indices.dims != output_dims || + indices.nbytes != kOutputWidth * sizeof(int32_t)) { + throw std::runtime_error( + "WebGPU topk: indices must use effective i32 [1,1,32] storage"); + } + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + graph.device(), + kTopkWGSL, + { + {0, WGPUBufferBindingType_Storage, values.buffer, values.nbytes}, + {1, WGPUBufferBindingType_Storage, indices.buffer, indices.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + input.buffer, + input.nbytes}, + }); + try { + graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, 1u, "topk_staged_serial"}); + } catch (...) { + wgpuComputePipelineRelease(bundle.pipeline); + wgpuBindGroupRelease(bundle.bind_group); + throw; + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.topk.default, topk_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/topk/topk.wgsl b/backends/webgpu/runtime/ops/topk/topk.wgsl new file mode 100644 index 00000000000..94f18172144 --- /dev/null +++ b/backends/webgpu/runtime/ops/topk/topk.wgsl @@ -0,0 +1,198 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var values_out: array; +@group(0) @binding(1) var indices_out: array; +@group(0) @binding(2) var scores: array; + +// The accepted topk.wgsl selects k=32 of 2048 with a 32-entry binary heap on a +// SINGLE invocation (@workgroup_size(1), dispatched 1x1x1). Measured on the W3 +// profile at 524.288 us per dispatch = 8 timestamp quanta, i.e. ~256 ns per +// element: one lane, one outstanding global load at a time, latency fully +// exposed because the very next instruction branches on the loaded value. +// +// This kernel changes NOTHING about the algorithm. The heap, the comparator +// and the emission order are transcribed character for character from the +// accepted kernel. The only differences are: +// +// 1. the row is staged into workgroup memory by 64 lanes first, so the +// serial scan reads `vals[i]` (threadgroup) instead of `scores[i]` +// (device) -- the loads are hoisted out of the dependent chain and +// issued 64-wide; +// 2. `heap_values[0]` -- the only heap slot the hot loop reads -- is +// mirrored into a register `heap_root`, refreshed after every +// adjust_heap call that can touch slot 0. `heap_values` is dynamically +// indexed through a pointer, so it is thread-local (device-backed) +// memory, and without this the hot loop pays a second dependent load per +// element. +// +// Both are behaviour-preserving by construction, so this kernel is bit-exact +// against the accepted one on EVERY input, including forced ties, NaN +// payloads and signed zeros. In particular it does NOT adopt a +// lowest-index-wins tie rule; that rule is an explicitly killed mutation +// (`tie_by_low_index`) of the accepted CPU authority bundle. + +const WG: u32 = 64u; +const N: u32 = 2048u; +const PER_LANE: u32 = N / WG; + +var vals: array; + +fn is_nan_bits(bits: u32) -> bool { + return (bits & 0x7f800000u) == 0x7f800000u && + (bits & 0x007fffffu) != 0u; +} + +fn float_less_than_bits(lhs: u32, rhs: u32) -> bool { + let lhs_nan = is_nan_bits(lhs); + let rhs_nan = is_nan_bits(rhs); + if (lhs_nan || rhs_nan) { + return !lhs_nan && rhs_nan; + } + + let lhs_magnitude = lhs & 0x7fffffffu; + let rhs_magnitude = rhs & 0x7fffffffu; + if (lhs_magnitude == 0u && rhs_magnitude == 0u) { + return false; + } + + let lhs_negative = (lhs & 0x80000000u) != 0u; + let rhs_negative = (rhs & 0x80000000u) != 0u; + if (lhs_negative != rhs_negative) { + return lhs_negative; + } + if (lhs_negative) { + return lhs > rhs; + } + return lhs < rhs; +} + +fn greater(lhs: u32, rhs: u32) -> bool { + return float_less_than_bits(rhs, lhs); +} + +fn push_heap( + heap_values: ptr>, + heap_indices: ptr>, + initial_hole: u32, + top: u32, + value_bits: u32, + value_index: u32) { + var hole = initial_hole; + while (hole > top) { + let parent = (hole - 1u) / 2u; + if (!greater((*heap_values)[parent], value_bits)) { + break; + } + (*heap_values)[hole] = (*heap_values)[parent]; + (*heap_indices)[hole] = (*heap_indices)[parent]; + hole = parent; + } + (*heap_values)[hole] = value_bits; + (*heap_indices)[hole] = value_index; +} + +fn adjust_heap( + heap_values: ptr>, + heap_indices: ptr>, + initial_hole: u32, + length: u32, + value_bits: u32, + value_index: u32) { + let top = initial_hole; + var hole = initial_hole; + var second_child = initial_hole; + while (second_child < (length - 1u) / 2u) { + second_child = 2u * (second_child + 1u); + if (greater( + (*heap_values)[second_child], + (*heap_values)[second_child - 1u])) { + second_child -= 1u; + } + (*heap_values)[hole] = (*heap_values)[second_child]; + (*heap_indices)[hole] = (*heap_indices)[second_child]; + hole = second_child; + } + if ((length & 1u) == 0u && second_child == (length - 2u) / 2u) { + second_child = 2u * (second_child + 1u); + (*heap_values)[hole] = (*heap_values)[second_child - 1u]; + (*heap_indices)[hole] = (*heap_indices)[second_child - 1u]; + hole = second_child - 1u; + } + push_heap( + heap_values, + heap_indices, + hole, + top, + value_bits, + value_index); +} + +@compute @workgroup_size(WG, 1, 1) +fn main(@builtin(local_invocation_id) lid: vec3) { + for (var j = 0u; j < PER_LANE; j += 1u) { + let idx = j * WG + lid.x; + vals[idx] = scores[idx]; + } + workgroupBarrier(); + if (lid.x != 0u) { + return; + } + + var heap_values: array; + var heap_indices: array; + for (var i = 0u; i < 32u; i += 1u) { + heap_values[i] = vals[i]; + heap_indices[i] = i; + } + + var parent = 15u; + loop { + let value_bits = heap_values[parent]; + let value_index = heap_indices[parent]; + adjust_heap( + &heap_values, + &heap_indices, + parent, + 32u, + value_bits, + value_index); + if (parent == 0u) { + break; + } + parent -= 1u; + } + + var heap_root = heap_values[0]; + for (var i = 32u; i < 2048u; i += 1u) { + let value_bits = vals[i]; + if (greater(value_bits, heap_root)) { + adjust_heap(&heap_values, &heap_indices, 0u, 32u, value_bits, i); + heap_root = heap_values[0]; + } + } + + var last = 32u; + while (last > 1u) { + last -= 1u; + let value_bits = heap_values[last]; + let value_index = heap_indices[last]; + heap_values[last] = heap_values[0]; + heap_indices[last] = heap_indices[0]; + adjust_heap( + &heap_values, + &heap_indices, + 0u, + last, + value_bits, + value_index); + } + + for (var i = 0u; i < 32u; i += 1u) { + values_out[i] = heap_values[i]; + indices_out[i] = heap_indices[i]; + } +} diff --git a/backends/webgpu/runtime/ops/topk/topk_wgsl.h b/backends/webgpu/runtime/ops/topk/topk_wgsl.h new file mode 100644 index 00000000000..0543b69ba14 --- /dev/null +++ b/backends/webgpu/runtime/ops/topk/topk_wgsl.h @@ -0,0 +1,222 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from topk.wgsl - DO NOT EDIT. +// wgsl-sha256: 5ee60ff98938deb3a7f531e6bcc9168dce6e356eebfc3f4dc063abfe9caf4437 +inline constexpr const char* kTopkWGSL = R"( +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +@group(0) @binding(0) var values_out: array; +@group(0) @binding(1) var indices_out: array; +@group(0) @binding(2) var scores: array; + +// The accepted topk.wgsl selects k=32 of 2048 with a 32-entry binary heap on a +// SINGLE invocation (@workgroup_size(1), dispatched 1x1x1). Measured on the W3 +// profile at 524.288 us per dispatch = 8 timestamp quanta, i.e. ~256 ns per +// element: one lane, one outstanding global load at a time, latency fully +// exposed because the very next instruction branches on the loaded value. +// +// This kernel changes NOTHING about the algorithm. The heap, the comparator +// and the emission order are transcribed character for character from the +// accepted kernel. The only differences are: +// +// 1. the row is staged into workgroup memory by 64 lanes first, so the +// serial scan reads `vals[i]` (threadgroup) instead of `scores[i]` +// (device) -- the loads are hoisted out of the dependent chain and +// issued 64-wide; +// 2. `heap_values[0]` -- the only heap slot the hot loop reads -- is +// mirrored into a register `heap_root`, refreshed after every +// adjust_heap call that can touch slot 0. `heap_values` is dynamically +// indexed through a pointer, so it is thread-local (device-backed) +// memory, and without this the hot loop pays a second dependent load per +// element. +// +// Both are behaviour-preserving by construction, so this kernel is bit-exact +// against the accepted one on EVERY input, including forced ties, NaN +// payloads and signed zeros. In particular it does NOT adopt a +// lowest-index-wins tie rule; that rule is an explicitly killed mutation +// (`tie_by_low_index`) of the accepted CPU authority bundle. + +const WG: u32 = 64u; +const N: u32 = 2048u; +const PER_LANE: u32 = N / WG; + +var vals: array; + +fn is_nan_bits(bits: u32) -> bool { + return (bits & 0x7f800000u) == 0x7f800000u && + (bits & 0x007fffffu) != 0u; +} + +fn float_less_than_bits(lhs: u32, rhs: u32) -> bool { + let lhs_nan = is_nan_bits(lhs); + let rhs_nan = is_nan_bits(rhs); + if (lhs_nan || rhs_nan) { + return !lhs_nan && rhs_nan; + } + + let lhs_magnitude = lhs & 0x7fffffffu; + let rhs_magnitude = rhs & 0x7fffffffu; + if (lhs_magnitude == 0u && rhs_magnitude == 0u) { + return false; + } + + let lhs_negative = (lhs & 0x80000000u) != 0u; + let rhs_negative = (rhs & 0x80000000u) != 0u; + if (lhs_negative != rhs_negative) { + return lhs_negative; + } + if (lhs_negative) { + return lhs > rhs; + } + return lhs < rhs; +} + +fn greater(lhs: u32, rhs: u32) -> bool { + return float_less_than_bits(rhs, lhs); +} + +fn push_heap( + heap_values: ptr>, + heap_indices: ptr>, + initial_hole: u32, + top: u32, + value_bits: u32, + value_index: u32) { + var hole = initial_hole; + while (hole > top) { + let parent = (hole - 1u) / 2u; + if (!greater((*heap_values)[parent], value_bits)) { + break; + } + (*heap_values)[hole] = (*heap_values)[parent]; + (*heap_indices)[hole] = (*heap_indices)[parent]; + hole = parent; + } + (*heap_values)[hole] = value_bits; + (*heap_indices)[hole] = value_index; +} + +fn adjust_heap( + heap_values: ptr>, + heap_indices: ptr>, + initial_hole: u32, + length: u32, + value_bits: u32, + value_index: u32) { + let top = initial_hole; + var hole = initial_hole; + var second_child = initial_hole; + while (second_child < (length - 1u) / 2u) { + second_child = 2u * (second_child + 1u); + if (greater( + (*heap_values)[second_child], + (*heap_values)[second_child - 1u])) { + second_child -= 1u; + } + (*heap_values)[hole] = (*heap_values)[second_child]; + (*heap_indices)[hole] = (*heap_indices)[second_child]; + hole = second_child; + } + if ((length & 1u) == 0u && second_child == (length - 2u) / 2u) { + second_child = 2u * (second_child + 1u); + (*heap_values)[hole] = (*heap_values)[second_child - 1u]; + (*heap_indices)[hole] = (*heap_indices)[second_child - 1u]; + hole = second_child - 1u; + } + push_heap( + heap_values, + heap_indices, + hole, + top, + value_bits, + value_index); +} + +@compute @workgroup_size(WG, 1, 1) +fn main(@builtin(local_invocation_id) lid: vec3) { + for (var j = 0u; j < PER_LANE; j += 1u) { + let idx = j * WG + lid.x; + vals[idx] = scores[idx]; + } + workgroupBarrier(); + if (lid.x != 0u) { + return; + } + + var heap_values: array; + var heap_indices: array; + for (var i = 0u; i < 32u; i += 1u) { + heap_values[i] = vals[i]; + heap_indices[i] = i; + } + + var parent = 15u; + loop { + let value_bits = heap_values[parent]; + let value_index = heap_indices[parent]; + adjust_heap( + &heap_values, + &heap_indices, + parent, + 32u, + value_bits, + value_index); + if (parent == 0u) { + break; + } + parent -= 1u; + } + + var heap_root = heap_values[0]; + for (var i = 32u; i < 2048u; i += 1u) { + let value_bits = vals[i]; + if (greater(value_bits, heap_root)) { + adjust_heap(&heap_values, &heap_indices, 0u, 32u, value_bits, i); + heap_root = heap_values[0]; + } + } + + var last = 32u; + while (last > 1u) { + last -= 1u; + let value_bits = heap_values[last]; + let value_index = heap_indices[last]; + heap_values[last] = heap_values[0]; + heap_indices[last] = heap_indices[0]; + adjust_heap( + &heap_values, + &heap_indices, + 0u, + last, + value_bits, + value_index); + } + + for (var i = 0u; i < 32u; i += 1u) { + values_out[i] = heap_values[i]; + indices_out[i] = heap_indices[i]; + } +} +)"; + +inline constexpr uint32_t kTopkWorkgroupSizeX = 64; +inline constexpr uint32_t kTopkWorkgroupSizeY = 1; +inline constexpr uint32_t kTopkWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/where/Where.cpp b/backends/webgpu/runtime/ops/where/Where.cpp index e70888308bf..86efc2d3d58 100644 --- a/backends/webgpu/runtime/ops/where/Where.cpp +++ b/backends/webgpu/runtime/ops/where/Where.cpp @@ -75,7 +75,7 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { const size_t cond_bind_size = (cond_tensor.nbytes + 3) & ~size_t(3); uint32_t wg_size = utils::clamp_workgroup_size(device, kWhereWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "where"); WGPUConstantEntry wg_size_constant = {}; @@ -120,8 +120,12 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - const size_t dispatch_idx = - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "where", + workgroup_count.y}); // Dynamic shapes: rebuild the 4 broadcast TensorMeta UBOs + dispatch count. WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf, @@ -170,9 +174,10 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { wgpuQueueWriteBuffer(g.queue(), c_buf, 0, &cm, sizeof(cm)); wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am)); wgpuQueueWriteBuffer(g.queue(), bb_buf, 0, &bm, sizeof(bm)); - g.dispatch_at(dispatch_idx).workgroup_count_x = - utils::compute_1d_workgroup_count( - g.device(), om.numel, wg_size, "where(resize)"); + const utils::WgCount workgroups = utils::compute_2d_workgroup_count( + g.device(), om.numel, wg_size, "where(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = workgroups.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = workgroups.y; }; graph.add_tensor_resize_hook(cond_id, where_resize); graph.add_tensor_resize_hook(a_id, where_resize); diff --git a/backends/webgpu/runtime/ops/where/where.wgsl b/backends/webgpu/runtime/ops/where/where.wgsl index 8c11db74200..4eb7e0562a0 100644 --- a/backends/webgpu/runtime/ops/where/where.wgsl +++ b/backends/webgpu/runtime/ops/where/where.wgsl @@ -23,8 +23,10 @@ fn cond_is_true(i: u32) -> bool { } @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/where/where_wgsl.h b/backends/webgpu/runtime/ops/where/where_wgsl.h index 01177210207..12b8710d1ac 100644 --- a/backends/webgpu/runtime/ops/where/where_wgsl.h +++ b/backends/webgpu/runtime/ops/where/where_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from where.wgsl - DO NOT EDIT. -// wgsl-sha256: 2c5c6491e95822f767920ad82b81b005808ad5550567b85fa98c1521f188dbfc +// wgsl-sha256: f254e27661bfaa523237e63f9ccd90665daa436e45f769903ee581b006b06576 inline constexpr const char* kWhereWGSL = R"( @group(0) @binding(0) var cond: array; @group(0) @binding(1) var input_a: array; @@ -40,8 +40,10 @@ fn cond_is_true(i: u32) -> bool { } @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index f2c750a50ac..244ab6eb67a 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -165,6 +165,12 @@ def test_parse_workgroup_ignores_commented_constants(self) -> None: ) self.assertEqual(g.parse_workgroup_size(src), (64, 1, 1)) + def test_scatter_unique_indices_workgroup_size(self) -> None: + source = ( + g.BACKEND_ROOT / "runtime/ops/scatter/scatter_unique_indices.wgsl" + ).read_text() + self.assertEqual(g.parse_workgroup_size(source), (64, 1, 1)) + def test_render_header_shape(self) -> None: wgsl = "@compute @workgroup_size(64, 1, 1)\nfn main(){}\n" h = g.render_header(Path("runtime/ops/update_cache/update_cache.wgsl"), wgsl) @@ -240,14 +246,14 @@ def test_generated_output_manifest_digest(self) -> None: digest.update(b"\0") digest.update(output.read_bytes()) digest.update(b"\0") - self.assertEqual(len(outputs), 144) + self.assertEqual(len(outputs), 149) self.assertEqual( digest.hexdigest(), - "b4a5a79ea7cdd1f18867106365da79b6faaa5f1065409cadec8bab2d64f3139e", + "8b2879a6ba11b57fd67aa961793ef9ff5142fdefaa7d9dcf41ab26276331f546", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "2717f916578362e00727dabd1d7e91a28c8cada0c238fa6dc323511cd672369a", + "1c26ac3f0671aeec5c648f78c9f2cbeb02b65f79478712a4f4de5c7e75446e8c", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -950,12 +956,12 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: "to_copy_float_to_int": ( "f32", "i32", - "c331e00e3171eecbe6317ac9df0a5f9cd6d25da26a9a587250f1cc6086dc3c8f", + "241d51293095623126da4f106092bd2d7327c26e00c9cad39bdb8dc546f48149", ), "to_copy_int_to_float": ( "i32", "f32", - "e18dd733a3838f83eded4977a2a2b21119099c8409b234f12474fae5acc9b195", + "9506570f98b9888a65603157f75d380f7784539205610ccbcc459d6afb6629c5", ), } self.assertEqual(set(variants), set(expected))