From c3cb71f2745894ce3e177bf1bcf4779e27cd810e Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:37 -0700 Subject: [PATCH 1/4] [ExecuTorch][WebGPU] Generate extrema and unary shader variants Pull Request resolved: https://github.com/pytorch/executorch/pull/21450 **Generate extrema and unary shader families** The extrema reductions and ten no-parameter unary kernels duplicated shader skeletons that could drift independently. This consolidates amax/amin behind one extrema template and abs/cos/exp/hardswish/neg/round/rsqrt/sin/sqrt/tanh behind one unary template while preserving the generated runtime payloads. Key changes: - Generate amax/amin from one extrema manifest. - Generate ten unary payloads from one operator-expression manifest. - Lock expanded bytes, registry entries, delegation, and boundary numerics. The attempted Unary lifecycle migration is intentionally not part of the stack: its performance campaign did not produce an authoritative passing result, so the Unary builder, interface, and activation/sigmoid call sites are restored to their pre-migration bytes. Co-authored-with: Claude Code. ghstack-source-id: 411961475 @exported-using-ghexport Differential Revision: [D113979760](https://our.internmc.facebook.com/intern/diff/D113979760/) --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 4 +- backends/webgpu/runtime/ops/amax/Reduce.cpp | 2 +- backends/webgpu/runtime/ops/amin/Reduce.cpp | 2 +- backends/webgpu/runtime/ops/amin/amin.wgsl | 49 -------- .../runtime/ops/{amax => extrema}/amax_wgsl.h | 2 +- .../runtime/ops/{amin => extrema}/amin_wgsl.h | 2 +- .../{amax/amax.wgsl => extrema/extrema.wgsl} | 4 +- .../webgpu/runtime/ops/extrema/extrema.yaml | 13 +++ backends/webgpu/runtime/ops/unary/abs.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/abs_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/cos.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/cos_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/exp_wgsl.h | 2 +- .../webgpu/runtime/ops/unary/hardswish.wgsl | 21 ---- .../webgpu/runtime/ops/unary/hardswish_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/neg.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/neg_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/round.wgsl | 21 ---- .../webgpu/runtime/ops/unary/round_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/rsqrt.wgsl | 21 ---- .../webgpu/runtime/ops/unary/rsqrt_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/sin.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/sin_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/sqrt.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/sqrt_wgsl.h | 2 +- backends/webgpu/runtime/ops/unary/tanh.wgsl | 21 ---- backends/webgpu/runtime/ops/unary/tanh_wgsl.h | 2 +- .../ops/unary/{exp.wgsl => unary.wgsl} | 2 +- backends/webgpu/runtime/ops/unary/unary.yaml | 29 +++++ .../test/native/test_compute_dispatch.cpp | 109 ++++++++++++++++++ backends/webgpu/test/op_tests/cases.py | 55 +++++---- backends/webgpu/test/ops/test_reduce.py | 108 ++++++++++++++++- backends/webgpu/test/test_wgsl_codegen.py | 90 ++++++++++++++- 33 files changed, 397 insertions(+), 283 deletions(-) delete mode 100644 backends/webgpu/runtime/ops/amin/amin.wgsl rename backends/webgpu/runtime/ops/{amax => extrema}/amax_wgsl.h (97%) rename backends/webgpu/runtime/ops/{amin => extrema}/amin_wgsl.h (97%) rename backends/webgpu/runtime/ops/{amax/amax.wgsl => extrema/extrema.wgsl} (94%) create mode 100644 backends/webgpu/runtime/ops/extrema/extrema.yaml delete mode 100644 backends/webgpu/runtime/ops/unary/abs.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/cos.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/hardswish.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/neg.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/round.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/rsqrt.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/sin.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/sqrt.wgsl delete mode 100644 backends/webgpu/runtime/ops/unary/tanh.wgsl rename backends/webgpu/runtime/ops/unary/{exp.wgsl => unary.wgsl} (94%) create mode 100644 backends/webgpu/runtime/ops/unary/unary.yaml diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 1cf4b952a5f..b475a6cd999 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -13,8 +13,6 @@ #include #include #include -#include -#include #include #include #include @@ -48,6 +46,8 @@ #include #include #include +#include +#include #include #include #include diff --git a/backends/webgpu/runtime/ops/amax/Reduce.cpp b/backends/webgpu/runtime/ops/amax/Reduce.cpp index 14675d97e96..c1dff62e035 100644 --- a/backends/webgpu/runtime/ops/amax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amax/Reduce.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/amin/Reduce.cpp b/backends/webgpu/runtime/ops/amin/Reduce.cpp index fbe574fdf0b..24a5ebeb826 100644 --- a/backends/webgpu/runtime/ops/amin/Reduce.cpp +++ b/backends/webgpu/runtime/ops/amin/Reduce.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/amin/amin.wgsl b/backends/webgpu/runtime/ops/amin/amin.wgsl deleted file mode 100644 index 4778800ab3d..00000000000 --- a/backends/webgpu/runtime/ops/amin/amin.wgsl +++ /dev/null @@ -1,49 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_rows: u32, - reduce_size: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -// Cooperative shared-memory reduction; mirrors Vulkan reduce.glsl (a group of -// threads co-operates per reduction row, partials aggregated in shared memory). -// Fixed upper bound (>= any clamped wg_size); only [0, wg_size) is used. -var partials: array; - -@compute @workgroup_size(wg_size) -fn main( - @builtin(workgroup_id) wid: vec3, - @builtin(local_invocation_id) lid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // One workgroup per reduction row; 2D-fold lifts the 65535 grid cap. - let row = wid.x + wid.y * num_workgroups.x; - if (row >= params.num_rows) { - return; - } - let base = row * params.reduce_size; - - // Each thread reduces a strided slice of the row into a partial. Seed with - // the row's first element (always valid; reduce_size >= 1) so threads that - // own no element contribute a real value, not an out-of-range identity. - var acc = input[base]; - var i = lid.x; - while (i < params.reduce_size) { - acc = min(acc, input[base + i]); - i = i + wg_size; - } - partials[lid.x] = acc; - workgroupBarrier(); - - // Thread 0 aggregates the wg_size partials (mirrors Vulkan's group aggregate). - if (lid.x == 0u) { - var m = partials[0]; - for (var t = 1u; t < wg_size; t = t + 1u) { - m = min(m, partials[t]); - } - output[row] = m; - } -} diff --git a/backends/webgpu/runtime/ops/amax/amax_wgsl.h b/backends/webgpu/runtime/ops/extrema/amax_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/amax/amax_wgsl.h rename to backends/webgpu/runtime/ops/extrema/amax_wgsl.h index 48ec8f20e27..0b7a4d2238e 100644 --- a/backends/webgpu/runtime/ops/amax/amax_wgsl.h +++ b/backends/webgpu/runtime/ops/extrema/amax_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from amax.wgsl - DO NOT EDIT. +// @generated from extrema.wgsl - DO NOT EDIT. // wgsl-sha256: 35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0 inline constexpr const char* kAmaxWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/amin/amin_wgsl.h b/backends/webgpu/runtime/ops/extrema/amin_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/amin/amin_wgsl.h rename to backends/webgpu/runtime/ops/extrema/amin_wgsl.h index 40a97c67a63..8b8bf5456b1 100644 --- a/backends/webgpu/runtime/ops/amin/amin_wgsl.h +++ b/backends/webgpu/runtime/ops/extrema/amin_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from amin.wgsl - DO NOT EDIT. +// @generated from extrema.wgsl - DO NOT EDIT. // wgsl-sha256: 8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2 inline constexpr const char* kAminWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/amax/amax.wgsl b/backends/webgpu/runtime/ops/extrema/extrema.wgsl similarity index 94% rename from backends/webgpu/runtime/ops/amax/amax.wgsl rename to backends/webgpu/runtime/ops/extrema/extrema.wgsl index 2f23b9c38fa..6c55d60d9b7 100644 --- a/backends/webgpu/runtime/ops/amax/amax.wgsl +++ b/backends/webgpu/runtime/ops/extrema/extrema.wgsl @@ -32,7 +32,7 @@ fn main( var acc = input[base]; var i = lid.x; while (i < params.reduce_size) { - acc = max(acc, input[base + i]); + acc = ${REDUCE_FN}(acc, input[base + i]); i = i + wg_size; } partials[lid.x] = acc; @@ -42,7 +42,7 @@ fn main( if (lid.x == 0u) { var m = partials[0]; for (var t = 1u; t < wg_size; t = t + 1u) { - m = max(m, partials[t]); + m = ${REDUCE_FN}(m, partials[t]); } output[row] = m; } diff --git a/backends/webgpu/runtime/ops/extrema/extrema.yaml b/backends/webgpu/runtime/ops/extrema/extrema.yaml new file mode 100644 index 00000000000..a85da849c60 --- /dev/null +++ b/backends/webgpu/runtime/ops/extrema/extrema.yaml @@ -0,0 +1,13 @@ +# 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. + +extrema: + parameter_names_with_default_values: + REDUCE_FN: max + shader_variants: + - NAME: amax + - NAME: amin + REDUCE_FN: min diff --git a/backends/webgpu/runtime/ops/unary/abs.wgsl b/backends/webgpu/runtime/ops/unary/abs.wgsl deleted file mode 100644 index e3e10c75dd9..00000000000 --- a/backends/webgpu/runtime/ops/unary/abs.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = abs(x); -} diff --git a/backends/webgpu/runtime/ops/unary/abs_wgsl.h b/backends/webgpu/runtime/ops/unary/abs_wgsl.h index 3d2873e69c4..0efb72d5189 100644 --- a/backends/webgpu/runtime/ops/unary/abs_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/abs_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from abs.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 39d3c163fdf6a92286828f4b3217e00294e3ca5634a878ed5fd34e3b1cdf0a27 inline constexpr const char* kAbsWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/cos.wgsl b/backends/webgpu/runtime/ops/unary/cos.wgsl deleted file mode 100644 index c2bafe7b248..00000000000 --- a/backends/webgpu/runtime/ops/unary/cos.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = cos(x); -} diff --git a/backends/webgpu/runtime/ops/unary/cos_wgsl.h b/backends/webgpu/runtime/ops/unary/cos_wgsl.h index 422b5618146..4ca99df88f3 100644 --- a/backends/webgpu/runtime/ops/unary/cos_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/cos_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from cos.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 9df78873e5fae98d347c26db2a02b047ea3d5d2c93f0761cb9ac6995f9a71ab2 inline constexpr const char* kCosWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/exp_wgsl.h b/backends/webgpu/runtime/ops/unary/exp_wgsl.h index 28b6c67cb18..cbf85fd415a 100644 --- a/backends/webgpu/runtime/ops/unary/exp_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/exp_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from exp.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 3171399bc36acf9c1cb2a03c2a31038318203c4c63ab03c4881df7a660346020 inline constexpr const char* kExpWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/hardswish.wgsl b/backends/webgpu/runtime/ops/unary/hardswish.wgsl deleted file mode 100644 index 2278caf3664..00000000000 --- a/backends/webgpu/runtime/ops/unary/hardswish.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = select(select(x * (x + 3.0) / 6.0, x, x >= 3.0), 0.0, x <= -3.0); -} diff --git a/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h b/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h index 0c991547b9a..43f104c7f8f 100644 --- a/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/hardswish_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from hardswish.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: c874a15ef6cdaec71187296016cc2a1515f5e7c889b97dfa8fd4b278e6e2c3d5 inline constexpr const char* kHardswishWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/neg.wgsl b/backends/webgpu/runtime/ops/unary/neg.wgsl deleted file mode 100644 index c977957957d..00000000000 --- a/backends/webgpu/runtime/ops/unary/neg.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = -x; -} diff --git a/backends/webgpu/runtime/ops/unary/neg_wgsl.h b/backends/webgpu/runtime/ops/unary/neg_wgsl.h index d528c45fea0..c4c7bb989ac 100644 --- a/backends/webgpu/runtime/ops/unary/neg_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/neg_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from neg.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 8851b9f42d14153f6f04484fee2f8bf67bda26dea892ff48768e09e6ad49cee1 inline constexpr const char* kNegWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/round.wgsl b/backends/webgpu/runtime/ops/unary/round.wgsl deleted file mode 100644 index 2269ca59988..00000000000 --- a/backends/webgpu/runtime/ops/unary/round.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = round(x); -} diff --git a/backends/webgpu/runtime/ops/unary/round_wgsl.h b/backends/webgpu/runtime/ops/unary/round_wgsl.h index 209305bc855..8c805ba4f80 100644 --- a/backends/webgpu/runtime/ops/unary/round_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/round_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from round.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 8f3e0edbeb81aa50f35e691c78554e8057fa8d78fe8a86454f4f42e5e8871452 inline constexpr const char* kRoundWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/rsqrt.wgsl b/backends/webgpu/runtime/ops/unary/rsqrt.wgsl deleted file mode 100644 index 1f50c4f66ac..00000000000 --- a/backends/webgpu/runtime/ops/unary/rsqrt.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = 1.0 / sqrt(x); -} diff --git a/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h b/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h index 83e74a1b8cf..58bd05d010e 100644 --- a/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/rsqrt_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from rsqrt.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 108765d5a23b87473f34651875d08abf2a5fa8980bd92fc8cbe3617295097747 inline constexpr const char* kRsqrtWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/sin.wgsl b/backends/webgpu/runtime/ops/unary/sin.wgsl deleted file mode 100644 index ffd2a07ea8a..00000000000 --- a/backends/webgpu/runtime/ops/unary/sin.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = sin(x); -} diff --git a/backends/webgpu/runtime/ops/unary/sin_wgsl.h b/backends/webgpu/runtime/ops/unary/sin_wgsl.h index f22229b2342..54184a1ccb8 100644 --- a/backends/webgpu/runtime/ops/unary/sin_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/sin_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from sin.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: e5762804773659d348fddddcef4935807ae6fe7d92c92eb17a2f44aae8f2c5b9 inline constexpr const char* kSinWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/sqrt.wgsl b/backends/webgpu/runtime/ops/unary/sqrt.wgsl deleted file mode 100644 index e34ff440007..00000000000 --- a/backends/webgpu/runtime/ops/unary/sqrt.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = sqrt(x); -} diff --git a/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h b/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h index 260a4ce4266..42dcdb838c0 100644 --- a/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/sqrt_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from sqrt.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 008534ae365969f5c180b42e8d6d0b131df78f181e5435abbcafc3ffb8be8aac inline constexpr const char* kSqrtWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/tanh.wgsl b/backends/webgpu/runtime/ops/unary/tanh.wgsl deleted file mode 100644 index 5a541e21699..00000000000 --- a/backends/webgpu/runtime/ops/unary/tanh.wgsl +++ /dev/null @@ -1,21 +0,0 @@ -@group(0) @binding(0) var input: array; -@group(0) @binding(1) var output: array; - -struct Params { - num_elements: u32, -} -@group(0) @binding(2) var params: Params; - -override wg_size: u32 = 256u; - -@compute @workgroup_size(wg_size) -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; - } - let x = input[idx]; - output[idx] = tanh(clamp(x, -15.0, 15.0)); -} diff --git a/backends/webgpu/runtime/ops/unary/tanh_wgsl.h b/backends/webgpu/runtime/ops/unary/tanh_wgsl.h index eef1ecd91af..51ba4d3919c 100644 --- a/backends/webgpu/runtime/ops/unary/tanh_wgsl.h +++ b/backends/webgpu/runtime/ops/unary/tanh_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from tanh.wgsl - DO NOT EDIT. +// @generated from unary.wgsl - DO NOT EDIT. // wgsl-sha256: 5bd7eb1c6411940d84a9b311884f35b39f15b82103b14bab02902290ed6b0339 inline constexpr const char* kTanhWGSL = R"( @group(0) @binding(0) var input: array; diff --git a/backends/webgpu/runtime/ops/unary/exp.wgsl b/backends/webgpu/runtime/ops/unary/unary.wgsl similarity index 94% rename from backends/webgpu/runtime/ops/unary/exp.wgsl rename to backends/webgpu/runtime/ops/unary/unary.wgsl index b69aa509d8e..d974a2f3319 100644 --- a/backends/webgpu/runtime/ops/unary/exp.wgsl +++ b/backends/webgpu/runtime/ops/unary/unary.wgsl @@ -17,5 +17,5 @@ fn main( return; } let x = input[idx]; - output[idx] = exp(x); + output[idx] = ${OPERATOR}; } diff --git a/backends/webgpu/runtime/ops/unary/unary.yaml b/backends/webgpu/runtime/ops/unary/unary.yaml new file mode 100644 index 00000000000..84c0dd660d7 --- /dev/null +++ b/backends/webgpu/runtime/ops/unary/unary.yaml @@ -0,0 +1,29 @@ +# 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. + +unary: + parameter_names_with_default_values: + OPERATOR: "abs(x)" + shader_variants: + - NAME: abs + - NAME: cos + OPERATOR: "cos(x)" + - NAME: exp + OPERATOR: "exp(x)" + - NAME: hardswish + OPERATOR: "select(select(x * (x + 3.0) / 6.0, x, x >= 3.0), 0.0, x <= -3.0)" + - NAME: neg + OPERATOR: "-x" + - NAME: round + OPERATOR: "round(x)" + - NAME: rsqrt + OPERATOR: "1.0 / sqrt(x)" + - NAME: sin + OPERATOR: "sin(x)" + - NAME: sqrt + OPERATOR: "sqrt(x)" + - NAME: tanh + OPERATOR: "tanh(clamp(x, -15.0, 15.0))" diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index b2b0d028567..6ed7229604b 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -440,6 +440,18 @@ void record_resize_probe(WebGPUGraph&, const ResizeProbeContext& context) { ++*context.calls; } +struct ThrowingResizeContext { + bool* fail; + int* calls; +}; + +void maybe_throw_resize(WebGPUGraph&, const ThrowingResizeContext& context) { + ++*context.calls; + if (*context.fail) { + throw std::runtime_error("resize hook failure"); + } +} + struct GridPickerContext { int tensor_id; uint32_t x_bias; @@ -487,6 +499,25 @@ TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { int tensor_calls = 0; int symint_observed = 0; int symint_calls = 0; + using ResizeProbeFn = void (*)(WebGPUGraph&, const ResizeProbeContext&); + EXPECT_THROW( + graph.add_tensor_resize_hook( + kResizeQ, + static_cast(nullptr), + ResizeProbeContext{0, &tensor_observed, &tensor_calls}), + std::runtime_error); + EXPECT_THROW( + graph.add_tensor_resize_hook( + kResizeSymInt, + record_resize_probe, + ResizeProbeContext{0, &tensor_observed, &tensor_calls}), + std::runtime_error); + EXPECT_THROW( + graph.add_resize_hook( + kResizeQ, + record_resize_probe, + ResizeProbeContext{0, &symint_observed, &symint_calls}), + std::runtime_error); { ResizeProbeContext tensor_context = {17, &tensor_observed, &tensor_calls}; ResizeProbeContext symint_context = {23, &symint_observed, &symint_calls}; @@ -510,6 +541,36 @@ TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { EXPECT_EQ(symint_calls, 1); } +TEST(WebGPUResizeHooks, RestoresDirtyTriggerWhenHookThrows) { + WebGPUGraph graph; + build_resize_test_graph(graph); + bool hook_fails = true; + int hook_calls = 0; + int picker_calls = 0; + graph.add_tensor_resize_hook( + kResizeQ, + maybe_throw_resize, + ThrowingResizeContext{&hook_fails, &hook_calls}); + const size_t dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_hook_retry"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 1u, 2u, &picker_calls, nullptr}); + picker_calls = 0; + + graph.resize_input(kResizeQ, {4, 3}); + EXPECT_THROW(graph.propagate_resize(), std::runtime_error); + EXPECT_EQ(hook_calls, 1); + EXPECT_EQ(picker_calls, 0); + expect_dispatch_grid(graph, dispatch, 9u, 10u); + + hook_fails = false; + EXPECT_NO_THROW(graph.propagate_resize()); + EXPECT_EQ(hook_calls, 2); + EXPECT_EQ(picker_calls, 1); + expect_dispatch_grid(graph, dispatch, 5u, 5u); +} + TEST(WebGPUDynamicDispatch, InitializesAndIsolatesTriggeredGrids) { WebGPUGraph graph; build_resize_test_graph(graph); @@ -615,12 +676,60 @@ TEST(WebGPUDynamicDispatch, HandlesCascadesAndStagesPickerFailures) { EXPECT_THROW(graph.propagate_resize(), std::runtime_error); expect_dispatch_grid(graph, first, 9u, 10u); expect_dispatch_grid(graph, second, 11u, 12u); + second_fails = false; + EXPECT_NO_THROW(graph.propagate_resize()); + expect_dispatch_grid(graph, first, 5u, 5u); + expect_dispatch_grid(graph, second, 7u, 7u); } TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { WebGPUGraph graph; build_resize_test_graph(graph); int calls = 0; + const size_t dispatches_before_invalid_trigger = graph.num_dispatches(); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_negative_trigger"), + -1, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_oob_trigger"), + graph.num_values(), + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_symint_trigger"), + kResizeSymInt, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}), + std::runtime_error); + EXPECT_EQ(calls, 0); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_invalid_trigger); + + auto pick_zero_grid = [](const WebGPUGraph&, const WebGPUDispatchGrid& grid) { + return grid; + }; + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_zero_x"), + kResizeQ, + +pick_zero_grid, + WebGPUDispatchGrid{0u, 1u}), + std::runtime_error); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_zero_y"), + kResizeQ, + +pick_zero_grid, + WebGPUDispatchGrid{1u, 0u}), + std::runtime_error); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_invalid_trigger); + bool initial_fails = true; const size_t dispatches_before_failure = graph.num_dispatches(); EXPECT_THROW( diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index be46d35b520..aa8f000bceb 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -112,7 +112,15 @@ DequantizeConstModule, QuantizeModule, ) -from executorch.backends.webgpu.test.ops.test_reduce import AmaxModule, AminModule +from executorch.backends.webgpu.test.ops.test_reduce import ( + amax_sign_trap_input, + amax_tie_input, + AmaxModule, + amin_sign_trap_input, + amin_tie_input, + AminModule, + EXTREMA_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_repeat import RepeatModule from executorch.backends.webgpu.test.ops.test_rms_norm import ( _CASES, @@ -489,35 +497,42 @@ def _floor_divide_suite() -> WebGPUTestSuite: ) -def _reduce_suite(module_cls) -> WebGPUTestSuite: - # Last-dim reduction; both keepdim variants over a 2d and a 3d shape. - return WebGPUTestSuite( - module_factory=lambda keepdim: module_cls(keepdim), - cases=[ - Case(name="keepdim_2d", construct={"keepdim": True}, inputs=((M1, M2),)), - Case(name="nodim_2d", construct={"keepdim": False}, inputs=((M1, M2),)), - Case( - name="keepdim_3d", - construct={"keepdim": True}, - inputs=((S, S1, S2),), - ), +def _reduce_suite(module_cls, op: str) -> WebGPUTestSuite: + generators = { + ("amax", "sign_trap"): amax_sign_trap_input, + ("amax", "tie"): amax_tie_input, + ("amin", "sign_trap"): amin_sign_trap_input, + ("amin", "tie"): amin_tie_input, + } + cases = [] + for name, shape, dim, keepdim, input_class in EXTREMA_CONFIGS: + inputs = (shape,) + kwargs = {} + if input_class != "default": + inputs = (InputSpec(shape=shape, gen=generators[(op, input_class)]),) + kwargs = {"atol": 0.0, "rtol": 0.0} + cases.append( Case( - name="nodim_3d", - construct={"keepdim": False}, - inputs=((S, S1, S2),), - ), - ], + name=name, + construct={"keepdim": keepdim, "dim": dim}, + inputs=inputs, + **kwargs, + ) + ) + return WebGPUTestSuite( + module_factory=module_cls, + cases=cases, ) @register_op_test("amax") def _amax_suite() -> WebGPUTestSuite: - return _reduce_suite(AmaxModule) + return _reduce_suite(AmaxModule, "amax") @register_op_test("amin") def _amin_suite() -> WebGPUTestSuite: - return _reduce_suite(AminModule) + return _reduce_suite(AminModule, "amin") @register_op_test("flip") diff --git a/backends/webgpu/test/ops/test_reduce.py b/backends/webgpu/test/ops/test_reduce.py index c7b91d8156c..cdfd7eed96c 100644 --- a/backends/webgpu/test/ops/test_reduce.py +++ b/backends/webgpu/test/ops/test_reduce.py @@ -19,12 +19,14 @@ from __future__ import annotations +import math import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes class ReduceModule(torch.nn.Module): @@ -61,6 +63,70 @@ def _det_input(shape) -> torch.Tensor: return ((flat % 17) - 8).div(16.0).reshape(shape) +# Shared structural and manifest-driven extrema case authority. +EXTREMA_CONFIGS = ( + ("keepdim_2d", (37, 41), -1, True, "default"), + ("nodim_2d", (37, 41), -1, False, "default"), + ("keepdim_3d", (5, 7, 11), -1, True, "default"), + ("nodim_3d", (5, 7, 11), -1, False, "default"), + ("sign_trap_63_drop", (3, 63), -1, False, "sign_trap"), + ("tie_64_keep", (2, 64), -1, True, "tie"), + ("tie_65_posdim_drop", (2, 65), 1, False, "tie"), + ("sign_trap_255_keep", (2, 255), -1, True, "sign_trap"), + ("tie_256_drop", (2, 256), -1, False, "tie"), + ("tie_257_posdim_keep", (2, 257), 1, True, "tie"), +) + + +def _sign_trap_input(shape, *, for_max: bool) -> torch.Tensor: + magnitude = ((torch.arange(math.prod(shape), dtype=torch.float32) % 31) + 1).div( + 8.0 + ) + rows = magnitude.reshape(-1, shape[-1]) + for row_index, row in enumerate(rows): + row.add_(float(row_index)) + values = -magnitude if for_max else magnitude + return values.reshape(shape) + + +def amax_sign_trap_input(shape) -> torch.Tensor: + return _sign_trap_input(shape, for_max=True) + + +def amin_sign_trap_input(shape) -> torch.Tensor: + return _sign_trap_input(shape, for_max=False) + + +def _tie_input(shape, *, for_max: bool) -> torch.Tensor: + values = ((torch.arange(math.prod(shape), dtype=torch.float32) % 29) - 14).div(8.0) + rows = values.reshape(-1, shape[-1]) + for row_index, row in enumerate(rows): + extreme = 16.0 + float(row_index) + if not for_max: + extreme = -extreme + row[-2] = extreme + row[-1] = extreme + return values.reshape(shape) + + +def amax_tie_input(shape) -> torch.Tensor: + return _tie_input(shape, for_max=True) + + +def amin_tie_input(shape) -> torch.Tensor: + return _tie_input(shape, for_max=False) + + +def _extrema_input(op: str, input_class: str, shape) -> torch.Tensor: + if input_class == "default": + return _det_input(shape) + if input_class == "sign_trap": + return ( + amax_sign_trap_input(shape) if op == "amax" else amin_sign_trap_input(shape) + ) + return amax_tie_input(shape) if op == "amax" else amin_tie_input(shape) + + def _export(m: torch.nn.Module, x: torch.Tensor): ep = torch.export.export(m, (x,)) return to_edge_transform_and_lower( @@ -128,21 +194,55 @@ def export_reduce_model( class AmaxModule(torch.nn.Module): - def __init__(self, keepdim: bool) -> None: + def __init__(self, keepdim: bool, dim: int = -1) -> None: super().__init__() self.keepdim = keepdim + self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.amax(x, dim=-1, keepdim=self.keepdim) + return torch.amax(x, dim=self.dim, keepdim=self.keepdim) class AminModule(torch.nn.Module): - def __init__(self, keepdim: bool) -> None: + def __init__(self, keepdim: bool, dim: int = -1) -> None: super().__init__() self.keepdim = keepdim + self.dim = dim def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.amin(x, dim=-1, keepdim=self.keepdim) + return torch.amin(x, dim=self.dim, keepdim=self.keepdim) + + +class TestExtrema(unittest.TestCase): + def test_config_contract(self) -> None: + self.assertEqual( + EXTREMA_CONFIGS, + ( + ("keepdim_2d", (37, 41), -1, True, "default"), + ("nodim_2d", (37, 41), -1, False, "default"), + ("keepdim_3d", (5, 7, 11), -1, True, "default"), + ("nodim_3d", (5, 7, 11), -1, False, "default"), + ("sign_trap_63_drop", (3, 63), -1, False, "sign_trap"), + ("tie_64_keep", (2, 64), -1, True, "tie"), + ("tie_65_posdim_drop", (2, 65), 1, False, "tie"), + ("sign_trap_255_keep", (2, 255), -1, True, "sign_trap"), + ("tie_256_drop", (2, 256), -1, False, "tie"), + ("tie_257_posdim_keep", (2, 257), 1, True, "tie"), + ), + ) + + def test_exports_fully_delegated(self) -> None: + for op, module_cls in (("amax", AmaxModule), ("amin", AminModule)): + for name, shape, dim, keepdim, input_class in EXTREMA_CONFIGS: + with self.subTest(op=op, config=name): + x = _extrema_input(op, input_class, shape) + ep = torch.export.export(module_cls(keepdim, dim).eval(), (x,)) + edge = to_edge_transform_and_lower( + ep, partitioner=[VulkanPartitioner()] + ) + graph = edge.exported_program().graph_module.graph + self.assertEqual(len(get_delegates(graph)), 1) + self.assertEqual(get_non_lowered_nodes(graph), []) if __name__ == "__main__": diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index b8c239483f2..d437a0a42e8 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -213,7 +213,7 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "ab15be30e7cfa2cb2f6fa7743d3b9f03535f5cc0b88d64d9b6bad8f777efda25", + "ef97dca2336315ee2c8b0f9e896c6aa082834ae94948bda3ccb42b1145f2bd27", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -952,9 +952,95 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "ce1777820bffe77e7cdda312f86a3fd41a090f8f282a6add66666446d06b1608", + "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", ) + def test_extrema_template_roundtrip_byte_identical(self) -> None: + extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" + template_path = extrema_dir / "extrema.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "amax": ( + "max", + "35fc059d7c72caa17f9cb1128823ecfd8f75be4ce24b6cd4f9629a97b52f64c0", + ), + "amin": ( + "min", + "8cb6035ae4d34eb2a6cc973d93d9847905722e967239c96033fccfe3a1943cb2", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (reduce_fn, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["REDUCE_FN"], reduce_fn) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = extrema_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (256, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/extrema/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "amax": "57f929b9f3087dc32403c3587884ce2ed4be2d03c4e80ff7035428b52e7e0e51", + "amin": "5dc947d4781a67df953b9c5970c4ab2119317b29657f983a9d308dfdf123dede", + } + for name, expected_hash in handler_hashes.items(): + handler = g.BACKEND_ROOT / f"runtime/ops/{name}/Reduce.cpp" + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", + ) + + def test_unary_template_roundtrip_byte_identical(self) -> None: + unary_dir = g.BACKEND_ROOT / "runtime/ops/unary" + template_path = unary_dir / "unary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "abs": "39d3c163fdf6a92286828f4b3217e00294e3ca5634a878ed5fd34e3b1cdf0a27", + "cos": "9df78873e5fae98d347c26db2a02b047ea3d5d2c93f0761cb9ac6995f9a71ab2", + "exp": "3171399bc36acf9c1cb2a03c2a31038318203c4c63ab03c4881df7a660346020", + "hardswish": "c874a15ef6cdaec71187296016cc2a1515f5e7c889b97dfa8fd4b278e6e2c3d5", + "neg": "8851b9f42d14153f6f04484fee2f8bf67bda26dea892ff48768e09e6ad49cee1", + "round": "8f3e0edbeb81aa50f35e691c78554e8057fa8d78fe8a86454f4f42e5e8871452", + "rsqrt": "108765d5a23b87473f34651875d08abf2a5fa8980bd92fc8cbe3617295097747", + "sin": "e5762804773659d348fddddcef4935807ae6fe7d92c92eb17a2f44aae8f2c5b9", + "sqrt": "008534ae365969f5c180b42e8d6d0b131df78f181e5435abbcafc3ffb8be8aac", + "tanh": "5bd7eb1c6411940d84a9b311884f35b39f15b82103b14bab02902290ed6b0339", + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + for name, expected_hash in expected.items(): + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **variants[name]}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + header = (unary_dir / f"{name}_wgsl.h").read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (256, 1, 1)) + self.assertEqual(entries[name].include, f"runtime/ops/unary/{name}_wgsl.h") + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + self.assertTrue({"clamp", "pow_scalar"}.isdisjoint(variants)) + def test_rms_norm_half_variant_is_type_correct(self) -> None: # A DTYPE=half expansion must emit compilable WGSL: `enable f16;`, an f32 # accumulator, loads widened to f32 for the reduction, and the store From 93658f4c60ba0ce2029a0c30196735726eb97d0f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:39 -0700 Subject: [PATCH 2/4] [ExecuTorch][WebGPU] Generate logical and arithmetic binary shader variants Pull Request resolved: https://github.com/pytorch/executorch/pull/21451 **Generate byte-identical logical and arithmetic binary variants from shared WGSL families** Logical AND/OR and four arithmetic binary kernels duplicated shader skeletons and broadcast logic. This consolidates logical AND/OR behind one packed-Boolean family and minimum/pow/floor_divide/mul into the existing binary family, with a permanent mixed-rank broadcast contract. Key changes: - Generate logical AND/OR from one operator-token manifest. - Generate minimum, pow, floor_divide, and mul beside the existing div/sub variants. - Lock same-shape and mixed-rank expressions, exact payloads/workgroups, PTE delegation, and broadcast boundary cases. No runtime C++ dispatch, bindings, pipeline construction, workgroups, or expanded shader payloads change. Four standalone WGSL inputs are removed, and future compatible variants require manifest entries instead of copied kernels. This follows the Vulkan binary-family pattern. Co-authored-with: Claude Code. ghstack-source-id: 411961479 @exported-using-ghexport Differential Revision: [D113979789](https://our.internmc.facebook.com/intern/diff/D113979789/) --- .../webgpu/runtime/WebGPUShaderRegistry.cpp | 12 +- .../binary_floor_divide_wgsl.h | 2 +- .../binary_minimum_wgsl.h | 2 +- .../ops/{mul => binary_op}/binary_mul_wgsl.h | 2 +- .../runtime/ops/binary_op/binary_op.wgsl | 25 +++- .../runtime/ops/binary_op/binary_op.yaml | 23 ++++ .../ops/{pow => binary_op}/binary_pow_wgsl.h | 2 +- .../runtime/ops/floor_divide/BinaryOp.cpp | 2 +- .../ops/floor_divide/binary_floor_divide.wgsl | 51 -------- .../runtime/ops/logical_and/LogicalAnd.cpp | 2 +- .../logical_and_wgsl.h | 2 +- .../logical_binary.wgsl} | 7 +- .../ops/logical_binary/logical_binary.yaml | 13 ++ .../logical_or_wgsl.h | 2 +- .../runtime/ops/logical_or/LogicalOr.cpp | 2 +- .../runtime/ops/logical_or/logical_or.wgsl | 25 ---- .../webgpu/runtime/ops/minimum/BinaryOp.cpp | 2 +- .../runtime/ops/minimum/binary_minimum.wgsl | 51 -------- backends/webgpu/runtime/ops/mul/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/mul/binary_mul.wgsl | 51 -------- backends/webgpu/runtime/ops/pow/BinaryOp.cpp | 2 +- .../webgpu/runtime/ops/pow/binary_pow.wgsl | 51 -------- backends/webgpu/test/op_tests/cases.py | 117 +++++++----------- .../webgpu/test/op_tests/test_generator.py | 115 ++++++++++++++++- backends/webgpu/test/ops/test_bitwise.py | 20 ++- backends/webgpu/test/ops/test_floor_divide.py | 4 +- backends/webgpu/test/ops/test_logical_and.py | 51 ++++---- backends/webgpu/test/ops/test_logical_or.py | 42 ++----- backends/webgpu/test/ops/test_minimum.py | 4 +- backends/webgpu/test/ops/test_pow.py | 4 +- backends/webgpu/test/test_wgsl_codegen.py | 116 +++++++++++++++-- 31 files changed, 402 insertions(+), 404 deletions(-) rename backends/webgpu/runtime/ops/{floor_divide => binary_op}/binary_floor_divide_wgsl.h (97%) rename backends/webgpu/runtime/ops/{minimum => binary_op}/binary_minimum_wgsl.h (97%) rename backends/webgpu/runtime/ops/{mul => binary_op}/binary_mul_wgsl.h (98%) rename backends/webgpu/runtime/ops/{pow => binary_op}/binary_pow_wgsl.h (98%) delete mode 100644 backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl rename backends/webgpu/runtime/ops/{logical_and => logical_binary}/logical_and_wgsl.h (96%) rename backends/webgpu/runtime/ops/{logical_and/logical_and.wgsl => logical_binary/logical_binary.wgsl} (73%) create mode 100644 backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml rename backends/webgpu/runtime/ops/{logical_or => logical_binary}/logical_or_wgsl.h (96%) delete mode 100644 backends/webgpu/runtime/ops/logical_or/logical_or.wgsl delete mode 100644 backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl delete mode 100644 backends/webgpu/runtime/ops/mul/binary_mul.wgsl delete mode 100644 backends/webgpu/runtime/ops/pow/binary_pow.wgsl diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index b475a6cd999..5374390722b 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -17,6 +17,10 @@ #include #include #include +#include +#include +#include +#include #include #include #include @@ -50,7 +54,6 @@ #include #include #include -#include #include #include #include @@ -65,20 +68,17 @@ #include #include #include -#include -#include +#include +#include #include -#include #include #include #include -#include #include #include #include #include #include -#include #include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h index fe2315df30f..dcd2b2e46d8 100644 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_floor_divide.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3 inline constexpr const char* kBinaryFloorDivideWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h index 88d9614ba59..c3f5f76c81f 100644 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_minimum.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: 929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972 inline constexpr const char* kBinaryMinimumWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h index c9f60dbd200..68784e82eb2 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_mul.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812 inline constexpr const char* kBinaryMulWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl index 98076ed98b5..4b42665013f 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl @@ -16,11 +16,14 @@ override wg_size: u32 = 64u; $if USE_ALPHA: override alpha: f32 = 1.0; -fn op(a: f32, b: f32) -> f32 { - return ${OP_EXPR}; -} +$if INLINE: + @compute @workgroup_size(wg_size, 1, 1) +$else: + fn op(a: f32, b: f32) -> f32 { + return ${OP_EXPR}; + } -@compute @workgroup_size(wg_size, 1, 1) + @compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { @@ -30,6 +33,8 @@ fn main( return; } + $if INLINE: + // Fast path: every input dim matches the output dim -> elementwise. 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] || @@ -38,10 +43,15 @@ fn main( } } if (same) { - output[idx] = op(input1[idx], input2[idx]); + $if INLINE: + output[idx] = ${SAME_EXPR}; + $else: + output[idx] = op(input1[idx], input2[idx]); return; } + $if INLINE: + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. var rem = idx; var l1: u32 = 0u; var l2: u32 = 0u; @@ -51,5 +61,8 @@ fn main( 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]); + $if INLINE: + output[idx] = ${BROADCAST_EXPR}; + $else: + output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml index 03040fcbcce..bd4cd452d35 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml @@ -2,6 +2,9 @@ binary_op: parameter_names_with_default_values: OP_EXPR: a + alpha * b USE_ALPHA: 1 + INLINE: 0 + SAME_EXPR: input1[idx] + input2[idx] + BROADCAST_EXPR: input1[l1] + input2[l2] shader_variants: - NAME: binary_div OP_EXPR: a / b @@ -9,3 +12,23 @@ binary_op: - NAME: binary_sub OP_EXPR: a - alpha * b USE_ALPHA: 1 + - NAME: binary_minimum + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: min(input1[idx], input2[idx]) + BROADCAST_EXPR: min(input1[l1], input2[l2]) + - NAME: binary_pow + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: pow(input1[idx], input2[idx]) + BROADCAST_EXPR: pow(input1[l1], input2[l2]) + - NAME: binary_floor_divide + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: floor(input1[idx] / input2[idx]) + BROADCAST_EXPR: floor(input1[l1] / input2[l2]) + - NAME: binary_mul + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: input1[idx] * input2[idx] + BROADCAST_EXPR: input1[l1] * input2[l2] diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h index 3532c091160..776d4b6693f 100644 --- a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_pow.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c inline constexpr const char* kBinaryPowWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp index a35a7587311..98a76228003 100644 --- a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl deleted file mode 100644 index 3b4edd788b0..00000000000 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@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; - -@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; - } - - // Fast path: every input dim matches the output dim -> elementwise. - 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] = floor(input1[idx] / input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - 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] = floor(input1[l1] / input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index b6a7e12830f..358cb29fff0 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h index 6a21a77a687..3c1f861f119 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_and.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20 inline constexpr const char* kLogicalAndWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl similarity index 73% rename from backends/webgpu/runtime/ops/logical_and/logical_and.wgsl rename to backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl index 9acb583f51c..fb5310cfde9 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl @@ -16,10 +16,13 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $if OP == "&": + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $else: + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. let w = gid.x + gid.y * (num_workgroups.x * wg_size); if (w >= params.num_words) { return; } - t_out[w] = t_a[w] & t_b[w]; + t_out[w] = t_a[w] ${OP} t_b[w]; } diff --git a/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml new file mode 100644 index 00000000000..65df4465b04 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml @@ -0,0 +1,13 @@ +# 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. + +logical_binary: + parameter_names_with_default_values: + OP: "&" + shader_variants: + - NAME: logical_and + - NAME: logical_or + OP: "|" diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h index d64898cb523..e61317d2ff1 100644 --- a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_or.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: 4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf inline constexpr const char* kLogicalOrWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp index d26f6b486d2..d750af11227 100644 --- a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl b/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl deleted file mode 100644 index d7e6176ba32..00000000000 --- a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl +++ /dev/null @@ -1,25 +0,0 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_a: array; -@group(0) @binding(2) var t_b: array; - -struct Params { - num_words: u32, - pad0: u32, - pad1: u32, - pad2: u32, -} -@group(0) @binding(3) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. - let w = gid.x + gid.y * (num_workgroups.x * wg_size); - if (w >= params.num_words) { - return; - } - t_out[w] = t_a[w] | t_b[w]; -} diff --git a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp index 3c150fc2580..f457bc9dc99 100644 --- a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl b/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl deleted file mode 100644 index e79cb2d2bcc..00000000000 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@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; - -@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; - } - - // Fast path: every input dim matches the output dim -> elementwise. - 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] = min(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - 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] = min(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index fdb7984b9e1..27bff334dd3 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -10,7 +10,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl deleted file mode 100644 index f82a16e4b21..00000000000 --- a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@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; - -@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; - } - - // Fast path: every input dim matches the output dim -> elementwise. - 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] = input1[idx] * input2[idx]; - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - 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] = input1[l1] * input2[l2]; -} diff --git a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp index d2f923880bf..1ac90015ffe 100644 --- a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl deleted file mode 100644 index 2114ef87fee..00000000000 --- a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@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; - -@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; - } - - // Fast path: every input dim matches the output dim -> elementwise. - 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] = pow(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - 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] = pow(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index aa8f000bceb..bbd6e13dc9d 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -38,10 +38,10 @@ ) from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_bitwise import ( + BITWISE_NOT_SHAPES, BitwiseAndModule, BitwiseNotModule, bw_gen_a, - bw_gen_b, ) from executorch.backends.webgpu.test.ops.test_cat import ( CatModule, @@ -71,14 +71,13 @@ make_qcs4w_linear_module, ) from executorch.backends.webgpu.test.ops.test_logical_and import ( - la_gen_a, - la_gen_b, + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, LogicalAndModule, ) from executorch.backends.webgpu.test.ops.test_logical_or import ( BitwiseOrModule, - lo_gen_a, - lo_gen_b, LogicalOrModule, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -253,7 +252,7 @@ def _rms_norm_suite() -> WebGPUTestSuite: @register_op_test("mul") def _mul_suite() -> WebGPUTestSuite: - # Full numeric coverage incl. broadcast (binary_mul.wgsl over a TensorMeta UBO); fp64 golden. + # Full binary_op-family numeric coverage, including broadcast. return WebGPUTestSuite( module_factory=lambda: MulModule(), cases=[ @@ -278,12 +277,19 @@ def _fn_config_suite(module_cls, configs) -> WebGPUTestSuite: @register_op_test("minimum") def _minimum_suite() -> WebGPUTestSuite: - # Same-shape numeric coverage (flat binary kernel; broadcast stays smoke). + # Same-shape and mixed-rank broadcast numeric coverage. return WebGPUTestSuite( module_factory=lambda: MinimumModule(), cases=[ Case(name="2d", inputs=((M1, M2), (M1, M2))), Case(name="3d", inputs=((S, S1, S2), (S, S1, S2))), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-3.0, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 4.0)), + ), + ), ], ) @@ -334,49 +340,33 @@ def _ge_suite() -> WebGPUTestSuite: return _compare_suite("ge") -@register_op_test("logical_and") -def _logical_and_suite() -> WebGPUTestSuite: - # out = (a>0) && (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), AND'd -> bool. Distinct a/b seeds so the masks - # differ (AND ~25% True, a real mix an OR mutant fails); all shapes numel % - # 4 == 0 (bool packs 4/word). float32 oracle (byte-exact bool golden). +def _logical_binary_suite(module_factory) -> WebGPUTestSuite: + # Every packed u32 receives all four Boolean pairs in byte order. def case(name, shape): return Case( name=name, construct={"shape": shape}, inputs=( - InputSpec(shape=shape, gen=la_gen_a), - InputSpec(shape=shape, gen=la_gen_b), + InputSpec(shape=shape, gen=logical_binary_gen_a), + InputSpec(shape=shape, gen=logical_binary_gen_b), ), ) return WebGPUTestSuite( - module_factory=lambda shape: LogicalAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + module_factory=module_factory, + cases=[case(name, shape) for name, shape in LOGICAL_BINARY_CASES], golden_dtype="float32", ) +@register_op_test("logical_and") +def _logical_and_suite() -> WebGPUTestSuite: + return _logical_binary_suite(lambda shape: LogicalAndModule(shape)) + + @register_op_test("bitwise_and") def _bitwise_and_suite() -> WebGPUTestSuite: - # bool bitwise AND == logical_and for canonical 0/1 (shares the handler). - # Two masks derived on-GPU from float inputs via gt.Tensor (baked zeros), - # distinct a/b seeds (AND ~25% True); all shapes numel % 4 == 0. - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=bw_gen_a), - InputSpec(shape=shape, gen=bw_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseAndModule(shape)) @register_op_test("bitwise_not") @@ -392,52 +382,22 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda shape: BitwiseNotModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case(name, shape) + for name, shape in zip(("2d", "3d", "sq"), BITWISE_NOT_SHAPES) + ], golden_dtype="float32", ) @register_op_test("logical_or") def _logical_or_suite() -> WebGPUTestSuite: - # out = (a>0) || (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), OR'd -> bool. Distinct a/b seeds (~50% each, - # independent -> OR ~75% True, a real mix an AND mutant fails); all shapes - # numel % 4 == 0. float32 oracle (byte-exact bool golden). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: LogicalOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: LogicalOrModule(shape)) @register_op_test("bitwise_or") def _bitwise_or_suite() -> WebGPUTestSuite: - # bool bitwise OR == logical_or for canonical 0/1 (shares the handler). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseOrModule(shape)) @register_op_test("pow") @@ -460,6 +420,13 @@ def _pow_suite() -> WebGPUTestSuite: InputSpec(shape=(S, S1, S2), gen=_unary_lin(-2.0, 3.0)), ), ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(0.1, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 3.0)), + ), + ), ], ) @@ -493,6 +460,14 @@ def _floor_divide_suite() -> WebGPUTestSuite: ), golden_fn=_floor_div_golden, ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-8.0, 8.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(0.5, 4.0)), + ), + golden_fn=_floor_div_golden, + ), ], ) diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index fabf79171c0..65f765812be 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -10,7 +10,15 @@ import torch from executorch.backends.webgpu.test.op_tests import generate_op_tests as g -from executorch.backends.webgpu.test.op_tests.test_suite import op_test_registry +from executorch.backends.webgpu.test.op_tests.test_suite import ( + InputSpec, + op_test_registry, +) +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) def _add_regular_case(): @@ -112,3 +120,108 @@ def test_manifest_schema_roundtrip(tmp_path): gd = e["golden"] assert {"path", "shape", "dtype", "output_index"} <= set(gd) assert gd["output_index"] == 0 + + +def test_logical_binary_case_contract(): + expected_cases = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), + ) + assert LOGICAL_BINARY_CASES == expected_cases + assert logical_binary_gen_a((8,)).tolist() == [ + -1.0, + -1.0, + 1.0, + 1.0, + -1.0, + -1.0, + 1.0, + 1.0, + ] + assert logical_binary_gen_b((8,)).tolist() == [ + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + ] + + for op in ("logical_and", "bitwise_and", "logical_or", "bitwise_or"): + suite = op_test_registry[op] + assert tuple((case.name, case.construct["shape"]) for case in suite.cases) == ( + expected_cases + ) + for case in suite.cases: + assert case.required is True + assert case.heavy is False + assert len(case.inputs) == 2 + assert case.inputs[0].gen is logical_binary_gen_a + assert case.inputs[1].gen is logical_binary_gen_b + + +def test_binary_shader_family_case_contract(): + expected = { + "minimum": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "pow": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "floor_divide": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "mul": ( + ("same", ((8, 32), (8, 32))), + ("bcast_lastdim", ((1, 1, 7, 896), (1, 1, 7, 1))), + ("bcast_firstdim", ((4, 4), (1, 4))), + ("bcast_4d_mixed", ((3, 5, 7, 11), (1, 5, 1, 11))), + ("mixedrank", ((4,), (3, 4))), + ), + } + + for op, cases in expected.items(): + suite = op_test_registry[op] + actual = tuple( + ( + case.name, + tuple( + spec.shape if isinstance(spec, InputSpec) else spec + for spec in case.inputs + ), + ) + for case in suite.cases + ) + assert actual == cases + + ranges = { + "minimum": ((-3.0, 3.0), (-2.0, 4.0)), + "pow": ((0.1, 3.0), (-2.0, 3.0)), + "floor_divide": ((-8.0, 8.0), (0.5, 4.0)), + } + for op, expected_ranges in ranges.items(): + case = next( + c for c in op_test_registry[op].cases if c.name == "broadcast_3d_2d" + ) + for spec, (start, end) in zip(case.inputs, expected_ranges): + assert isinstance(spec, InputSpec) and callable(spec.gen) + values = spec.gen(spec.shape).flatten() + assert torch.isclose(values[0], torch.tensor(start)) + assert torch.isclose(values[-1], torch.tensor(end)) + + assert all( + case.golden_fn is g.cases._floor_div_golden + for case in op_test_registry["floor_divide"].cases + ) diff --git a/backends/webgpu/test/ops/test_bitwise.py b/backends/webgpu/test/ops/test_bitwise.py index b61a37d9990..a8fee7a0095 100644 --- a/backends/webgpu/test/ops/test_bitwise.py +++ b/backends/webgpu/test/ops/test_bitwise.py @@ -19,6 +19,11 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -50,11 +55,10 @@ def g(shape): bw_gen_a = _bw_gen(0) -bw_gen_b = _bw_gen(1) # All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +BITWISE_NOT_SHAPES = ((4, 8), (2, 3, 8), (16, 16)) class BitwiseTest(unittest.TestCase): @@ -75,13 +79,17 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = bw_gen_a(shape) - b = bw_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(op="bitwise_and", case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( BitwiseAndModule(shape), (a, b), "bitwise_and", shape ) + + for shape in BITWISE_NOT_SHAPES: + with self.subTest(op="bitwise_not", shape=shape): + a = bw_gen_a(shape) self._assert_delegates( BitwiseNotModule(shape), (a,), "bitwise_not", shape ) diff --git a/backends/webgpu/test/ops/test_floor_divide.py b/backends/webgpu/test/ops/test_floor_divide.py index 6020480ff2c..141e8314cde 100644 --- a/backends/webgpu/test/ops/test_floor_divide.py +++ b/backends/webgpu/test/ops/test_floor_divide.py @@ -6,8 +6,8 @@ """`aten.div.Tensor_mode` module for the WebGPU op-test framework. -`FloorDivideModule` is imported by `cases.py`. Same-shape elementwise -`div(a, b, rounding_mode="floor")`. The kernel computes `floor(a/b)` mirroring +`FloorDivideModule` is imported by `cases.py`. Same-shape and broadcast +`div(a, b, rounding_mode="floor")` use `floor(a/b)`, mirroring the Vulkan `floor_divide` glsl (`floor(X/Y)`); this differs from torch's own fmod-corrected `div_floor` at rare fp boundaries, so the suite goldens against a `floor(a/b)` `golden_fn` (Vulkan-faithful), not this module's eager output. diff --git a/backends/webgpu/test/ops/test_logical_and.py b/backends/webgpu/test/ops/test_logical_and.py index 4c3939002d0..9cf474fd8de 100644 --- a/backends/webgpu/test/ops/test_logical_and.py +++ b/backends/webgpu/test/ops/test_logical_and.py @@ -4,17 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_and.default` module + configs for the WebGPU op-test framework. - -`LogicalAndModule` derives its two bool operands on-GPU from float inputs -(`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so -the only runtime inputs are the two float tensors (the op-test framework is -float-input-only). `a`/`b` use distinct seeds so the two bool masks differ (each -~50% True, independent -> AND ~25% True), a real mix that a wrong op (e.g. OR) -would fail. Output is bool (byte-exact golden). `LogicalAndTest` is the -export-delegation smoke test. -""" +"""Delegation coverage for logical AND with packed truth-table inputs.""" +import math import unittest import torch @@ -32,29 +24,40 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.logical_and(a > self.z, b > self.z) -def _la_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) +LOGICAL_BINARY_CASES = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), +) - return g +def _logical_binary_gen(pattern): + def generate(shape): + numel = math.prod(shape) + if numel == 0 or numel % len(pattern) != 0: + raise ValueError("logical-binary test shapes must have numel % 4 == 0") + return ( + torch.tensor(pattern, dtype=torch.float32) + .repeat(numel // len(pattern)) + .reshape(shape) + ) -la_gen_a = _la_gen(0) -la_gen_b = _la_gen(1) + return generate -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +logical_binary_gen_a = _logical_binary_gen((-1.0, -1.0, 1.0, 1.0)) +logical_binary_gen_b = _logical_binary_gen((-1.0, 1.0, -1.0, 1.0)) class LogicalAndTest(unittest.TestCase): def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = la_gen_a(shape) - b = la_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) ep = torch.export.export(LogicalAndModule(shape).eval(), (a, b)) edge = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] diff --git a/backends/webgpu/test/ops/test_logical_or.py b/backends/webgpu/test/ops/test_logical_or.py index d71c1493ea9..420fee26336 100644 --- a/backends/webgpu/test/ops/test_logical_or.py +++ b/backends/webgpu/test/ops/test_logical_or.py @@ -4,23 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_or.default` / `aten.bitwise_or.Tensor` (bool) modules + configs. - -Mirrors the logical_and/bitwise_and tests: the modules derive their two bool -operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` -against a baked zero buffer), so the only runtime inputs are the two float -tensors (the op-test framework is float-input-only). `a`/`b` use distinct seeds -so the two bool masks differ (each ~50% True, independent -> OR ~75% True), a -real mix that a wrong op (e.g. AND) would fail. `bitwise_or` on bool is identical -to `logical_or` (shares the handler). Output is bool (byte-exact golden). -`LogicalOrTest` is the export-delegation smoke test. -""" +"""Delegation coverage for logical and bitwise OR truth-table inputs.""" import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -42,23 +37,6 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.bitwise_or(a > self.z, b > self.z) -def _lo_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) - - return g - - -lo_gen_a = _lo_gen(0) -lo_gen_b = _lo_gen(1) - - -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] - - class LogicalOrTest(unittest.TestCase): def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ep = torch.export.export(mod.eval(), inputs) @@ -77,10 +55,10 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = lo_gen_a(shape) - b = lo_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( LogicalOrModule(shape), (a, b), "logical_or", shape ) diff --git a/backends/webgpu/test/ops/test_minimum.py b/backends/webgpu/test/ops/test_minimum.py index f8f8088fef5..ba466317d5b 100644 --- a/backends/webgpu/test/ops/test_minimum.py +++ b/backends/webgpu/test/ops/test_minimum.py @@ -6,8 +6,8 @@ """`aten.minimum.default` module for the WebGPU op-test framework. -`MinimumModule` is imported by `cases.py`. minimum is a same-shape elementwise -binary op mirroring the landed `add`/`mul` pattern (flat 2D-dispatch kernel). +`MinimumModule` is imported by `cases.py`. minimum is an elementwise binary op +with the same-shape fast path and broadcast indexing used by `mul`. """ import torch diff --git a/backends/webgpu/test/ops/test_pow.py b/backends/webgpu/test/ops/test_pow.py index 93596dc5d9f..8498bd4f7aa 100644 --- a/backends/webgpu/test/ops/test_pow.py +++ b/backends/webgpu/test/ops/test_pow.py @@ -6,8 +6,8 @@ """`aten.pow.Tensor_Tensor` module for the WebGPU op-test framework. -`PowModule` is imported by `cases.py`. Same-shape elementwise `pow(a, b)`; the -suite uses a POSITIVE base so `pow(neg, frac)` (NaN) is never exercised. +`PowModule` is imported by `cases.py`. The suite covers same-shape and broadcast +`pow(a, b)` with POSITIVE bases so `pow(neg, frac)` never produces NaN. """ import torch diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index d437a0a42e8..574c3869864 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -213,7 +213,11 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "ef97dca2336315ee2c8b0f9e896c6aa082834ae94948bda3ccb42b1145f2bd27", + "e502196846f0f8100f468e5d9f8f9c006b67e08df54e1e2e667daa2fc50d8844", + ) + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "492535b396833ad6ebfed29b093057e38f40b1182b5c7d6b3eb2f3577cab024e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -950,10 +954,6 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: entries["to_copy_int_to_float"].include, "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h", ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) def test_extrema_template_roundtrip_byte_identical(self) -> None: extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" @@ -1004,10 +1004,108 @@ def test_extrema_template_roundtrip_byte_identical(self) -> None: hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) + def test_logical_binary_template_roundtrip_byte_identical(self) -> None: + logical_dir = g.BACKEND_ROOT / "runtime/ops/logical_binary" + template_path = logical_dir / "logical_binary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "logical_and": ( + "&", + "cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20", + ), + "logical_or": ( + "|", + "4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (op, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["OP"], op) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = logical_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/logical_binary/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "logical_and": "eb85a8f97ee7640298a661da49feb08aa79b8c24d3d4458b71d24d3f01bc388d", + "logical_or": "bda18617f7077fee5a812c21cdc495c89542a1688f7e1ef6739ed01da343a66b", + } + for name, expected_hash in handler_hashes.items(): + handler = ( + g.BACKEND_ROOT / f"runtime/ops/{name}/Logical{name[8:].title()}.cpp" + ) + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + def test_binary_family_roundtrip_byte_identical(self) -> None: + binary_dir = g.BACKEND_ROOT / "runtime/ops/binary_op" + template_path = binary_dir / "binary_op.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "binary_div": ( + 0, + "e36b560fd623dd5337b9ae57acd8981c9c635b995d6021caf1331c182cd3f0cd", + ), + "binary_sub": ( + 0, + "63209ff70422a21fc340d9aadba0945bc259bba89bdf05db018a6507d01c7ae5", + ), + "binary_minimum": ( + 1, + "929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972", + ), + "binary_pow": ( + 1, + "a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c", + ), + "binary_floor_divide": ( + 1, + "baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3", + ), + "binary_mul": ( + 1, + "d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + + for name, (inline, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["INLINE"], inline) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header = (binary_dir / f"{name}_wgsl.h").read_text() + literal = header.split('R"(', 1)[1].split(')";', 1)[0] + self.assertEqual(literal, "\n" + expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(expanded), (64, 1, 1)) + self.assertIn(f"k{g.symbol_base(name)}WGSL", header) + self.assertEqual( + entries[name].include, + f"runtime/ops/binary_op/{name}_wgsl.h", + ) def test_unary_template_roundtrip_byte_identical(self) -> None: unary_dir = g.BACKEND_ROOT / "runtime/ops/unary" From 2f8ffe2d7cae886e0201b66c5b0a71fe4b80e074 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 09:35:41 -0700 Subject: [PATCH 3/4] [ExecuTorch][WebGPU] Preserve Linear and Q4 embedding params across dynamic resize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/21483 **Preserve immutable Linear and Q4 embedding shader controls across dynamic resize through typed parameter authorities.** **Problem** The Linear resize hook dropped has_bias, and the Q4 embedding resize path could drop is_linear_weight. Both defects produced correct build-time outputs but silently changed semantics after a live shape update. **Solution** - Before: build and resize populated control words independently, allowing immutable shader state to reset. - After: each build/resize pair calls one typed helper while recomputing only live counts and dispatch. **Implementation** - Linear.cpp — make_linear_params preserves the bias flag across vec4 and tiled routes. - EmbeddingQ4gsw.cpp — EmbeddingLayout and make_embedding_params preserve nibble layout across resize. - Mirrors Vulkan runtime/graph/ops/impl/Linear.cpp and EmbeddingQ4gsw.cpp, which carry immutable shader controls through dynamic dispatch. **Constraints** WGSL bytes, shader selection, bindings, uniform sizes, queue-write counts, dispatch formulas, and pipeline topology are unchanged. Production handler code is smaller because duplicated field population and scalar captures are removed. Co-authored-with: Claude Code. ghstack-source-id: 411961489 @exported-using-ghexport Differential Revision: [D113992326](https://our.internmc.facebook.com/intern/diff/D113992326/) --- .../ops/embedding_q4gsw/EmbeddingQ4gsw.cpp | 58 +++++--- backends/webgpu/runtime/ops/linear/Linear.cpp | 18 +-- .../webgpu/test/native/test_dynamic_shape.cpp | 40 ++++- .../test_dynamic_shape_export.py | 140 +++++++++++++----- 4 files changed, 183 insertions(+), 73 deletions(-) diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 64956121b17..b5ea0cdca4c 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -37,13 +37,37 @@ static_assert( sizeof(EmbeddingParams) == 32, "EmbeddingParams must be 32 bytes"); +struct EmbeddingLayout { + uint32_t embed_dim; + uint32_t blocks_per_row; + uint32_t group_size; + uint32_t groups_per_row; + uint32_t bytes_per_row; + bool is_linear_weight; +}; + +EmbeddingParams make_embedding_params( + const EmbeddingLayout& layout, + uint32_t num_indices, + uint32_t total_blocks) { + return { + layout.embed_dim, + layout.blocks_per_row, + num_indices, + layout.group_size, + layout.groups_per_row, + layout.bytes_per_row, + total_blocks, + layout.is_linear_weight ? 1u : 0u}; +} + // Resize hook body: recompute counts/dispatch; out = indices dims + // [embed_dim]. void resize_embedding_q4gsw( WebGPUGraph& g, int indices_id, int out_id, - EmbeddingParams params, + const EmbeddingLayout& layout, uint32_t wg_size, size_t dispatch_idx, WGPUBuffer params_buf) { @@ -52,17 +76,17 @@ void resize_embedding_q4gsw( if (ni == 0) { throw std::runtime_error("WebGPU embedding_q4gsw: zero indices"); } - const uint64_t total_blocks = ni * params.blocks_per_row; + const uint64_t total_blocks = ni * layout.blocks_per_row; if (total_blocks > UINT32_MAX) { throw std::runtime_error( "WebGPU embedding_q4gsw: total_blocks exceeds uint32"); } std::vector od = id; - od.push_back(static_cast(params.embed_dim)); + od.push_back(static_cast(layout.embed_dim)); g.set_cur_dims(out_id, od); - params.num_indices = static_cast(ni); - params.total_blocks = static_cast(total_blocks); - wgpuQueueWriteBuffer(g.queue(), params_buf, 0, ¶ms, sizeof(params)); + EmbeddingParams p = make_embedding_params( + layout, static_cast(ni), static_cast(total_blocks)); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); g.dispatch_at(dispatch_idx).workgroup_count_x = utils::compute_1d_workgroup_count( g.device(), @@ -166,15 +190,15 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t workgroup_count = utils::compute_1d_workgroup_count( device, static_cast(total_blocks), wg_size, "embedding_q4gsw"); - EmbeddingParams params = {}; - params.embed_dim = embed_dim; - params.blocks_per_row = blocks_per_row; - params.num_indices = num_indices; // std140 layout only; shader derives it - params.group_size = static_cast(group_size); - params.groups_per_row = groups_per_row; - params.bytes_per_row = bytes_per_row; - params.total_blocks = static_cast(total_blocks); - params.is_linear_weight = is_linear ? 1u : 0u; + const EmbeddingLayout layout = { + embed_dim, + blocks_per_row, + static_cast(group_size), + groups_per_row, + bytes_per_row, + is_linear}; + EmbeddingParams params = make_embedding_params( + layout, num_indices, static_cast(total_blocks)); WGPUBufferDescriptor uniform_desc = {}; uniform_desc.size = sizeof(EmbeddingParams); @@ -223,10 +247,10 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( indices_id, - [indices_id, out_id, params, wg_size, dispatch_idx, params_buf]( + [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf]( WebGPUGraph& g) { resize_embedding_q4gsw( - g, indices_id, out_id, params, wg_size, dispatch_idx, params_buf); + g, indices_id, out_id, layout, wg_size, dispatch_idx, params_buf); }); // Graph owns it so the resize hook can rewrite it; freed in the dtor. diff --git a/backends/webgpu/runtime/ops/linear/Linear.cpp b/backends/webgpu/runtime/ops/linear/Linear.cpp index 45cc27a515d..ba82fcc36a4 100644 --- a/backends/webgpu/runtime/ops/linear/Linear.cpp +++ b/backends/webgpu/runtime/ops/linear/Linear.cpp @@ -33,6 +33,11 @@ static_assert(sizeof(LinearParams) == 16, "LinearParams must be 16 bytes"); constexpr uint32_t kTile = 32u; +LinearParams +make_linear_params(uint32_t M, uint32_t N, uint32_t K, bool has_bias) { + return {M, N, K, has_bias ? 1u : 0u}; +} + // aten.linear (+ optional bias); shared-memory tiled GEMM. void linear_impl(WebGPUGraph& graph, const std::vector& args) { // args: [input, weight, bias?, out]; out is last. bias (arg 2) is a tensor @@ -82,11 +87,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { } } - LinearParams params = {}; - params.M = M; - params.N = N; - params.K = K; - params.has_bias = has_bias ? 1u : 0u; + LinearParams params = make_linear_params(M, N, K, has_bias); // Bias binding (binding 4); a 4-byte dummy satisfies it when None // (WGSL-gated). @@ -138,7 +139,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( in_id, - [in_id, out_id, M, N, K, dispatch_x, dispatch_idx, params_buf]( + [in_id, out_id, M, N, K, has_bias, dispatch_x, dispatch_idx, params_buf]( WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); const uint64_t numel = utils::numel_of(d); @@ -152,10 +153,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "WebGPU linear: live M is 0 or exceeds the build-time max"); } - LinearParams p = {}; - p.M = m; - p.N = N; - p.K = K; + LinearParams p = make_linear_params(m, N, K, has_bias); wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); g.dispatch_at(dispatch_idx).workgroup_count_x = dispatch_x; g.dispatch_at(dispatch_idx).workgroup_count_y = diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 0e72ce6234d..7b3c1d6b0a7 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -133,6 +133,7 @@ constexpr int kLinK = 64; constexpr int kLinAltK = 72; constexpr int kLinN = 128; constexpr int kLinNShmem = 2048; +constexpr int kFp32LinearN = 32; // Run at [m_rows, kLinK] on an already-loaded module (so it can be // reused across M without a fresh load), and compare to the golden. void run_linear( @@ -233,6 +234,14 @@ void check_linear_tiled(int m_rows) { run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); } +void check_fp32_linear_reused(const char* prefix, int k) { + Module module(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix << ".pte"; + for (int m_rows : {128, 32, 1, 128}) { + run_linear(module, m_rows, prefix, kFp32LinearN, k, 1e-3f); + } +} + constexpr int kQkvNq = 2048; constexpr int kQkvNk = 512; constexpr int kQkvNv = 512; @@ -941,6 +950,23 @@ TEST(DynamicShape, RmsMul) { } } +// I0: dynamic fp32 linear preserves bias across repeated resizes. +TEST(DynamicShape, Fp32LinearVec4BiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_vec4_bias", 64); +} + +TEST(DynamicShape, Fp32LinearVec4UnbiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_vec4_no_bias", 64); +} + +TEST(DynamicShape, Fp32LinearTiledBiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_tiled_bias", 63); +} + +TEST(DynamicShape, Fp32LinearTiledUnbiasedReusedGraph) { + check_fp32_linear_reused("dyn_linear_fp32_tiled_no_bias", 63); +} + // I: dynamic 4-bit quantized linear (prefill GEMM) at several M. TEST(DynamicShape, QuantizedLinear) { for (int m_rows : {128, 32, 1}) { @@ -1670,12 +1696,14 @@ TEST(DynamicShape, EmbeddingReusedGraph) { } } -// K3: linear-packed reuse must preserve nibble order across resizes. -TEST(DynamicShape, LinearPackedEmbeddingReusedGraph) { - Module m(g_dir + "/emb_dyn_linear.pte"); - ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn_linear.pte"; - for (int n : {16, 8, 1, 16}) { - run_embedding(m, n, "emb_dyn_linear"); +// K3: linear/nonlinear-packed reuse must preserve nibble order across resizes. +TEST(DynamicShape, EmbeddingLayoutsReusedGraph) { + for (const char* prefix : {"emb_dyn_linear", "emb_dyn_nonlinear"}) { + Module m(g_dir + "/" + prefix + ".pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load " << prefix << ".pte"; + for (int n : {16, 8, 1, 16}) { + run_embedding(m, n, prefix); + } } } diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index a99c9baf46e..b302b3c120d 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -297,7 +297,10 @@ def export_dynamic_shape_cases(out_dir: str) -> None: ) _write_goldens(rmsmul, "dyn_rmsmul", out_dir, [MAXS, 32, 1]) - # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM + # 2d) fp32 linear with a dynamic rows (M) dim. + export_dynamic_fp32_linear_cases(out_dir) + + # 2d.1) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM # (register-tiled N=128) + a shmem-GEMM-routed variant (N=2048). _export_dynamic_linear(out_dir) _export_dynamic_linear( @@ -327,7 +330,7 @@ def export_dynamic_shape_cases(out_dir: str) -> None: _export_static_sdpa(out_dir, 16, "static_sdpa_s16") # 2f) 4-bit embedding with a DYNAMIC token count (int64 indices). - _export_dynamic_embedding(out_dir) + export_dynamic_embedding_cases(out_dir) # 2g) Interleaved RoPE with a DYNAMIC seq-len S (two outputs xq/xk). _export_dynamic_rope(out_dir) @@ -370,6 +373,10 @@ def export_dynamic_shape_cases(out_dir: str) -> None: LIN_GROUP = 32 LIN_MAXM = 128 +FP32_LINEAR_N = 32 +FP32_LINEAR_MAXM = 128 +FP32_LINEAR_M_VALUES = (FP32_LINEAR_MAXM, 32, 1) + BK64_K = 2048 BK64_N = 2048 BK64_KV_N = 512 @@ -533,6 +540,46 @@ def _export_dynamic_linear( print(f" golden {prefix} M={m}") +def _export_dynamic_fp32_linear_case( + out_dir: str, + *, + k: int, + bias: bool, + prefix: str, +) -> None: + from executorch.backends.webgpu.test.ops.test_linear_fp32 import make_linear + + model = make_linear(k, FP32_LINEAR_N, bias=bias).eval() + x = _ramp((FP32_LINEAR_MAXM, k)) + m_dim = torch.export.Dim("m", min=1, max=FP32_LINEAR_MAXM) + ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + + weight = model.fc.weight.detach().double() + bias_value = model.fc.bias.detach().double() if model.fc.bias is not None else None + for m in FP32_LINEAR_M_VALUES: + xm = _ramp((m, k)) + golden = torch.nn.functional.linear(xm.double(), weight, bias_value) + base = os.path.join(out_dir, f"{prefix}.S{m}") + xm.detach().numpy().astype(" None: + os.makedirs(out_dir, exist_ok=True) + for route, k in (("vec4", 64), ("tiled", 63)): + for bias in (True, False): + suffix = "bias" if bias else "no_bias" + _export_dynamic_fp32_linear_case( + out_dir, + k=k, + bias=bias, + prefix=f"dyn_linear_fp32_{route}_{suffix}", + ) + + def _make_bk64_model( *, k: int = BK64_K, @@ -1306,52 +1353,64 @@ def _export_static_sdpa(out_dir: str, s: int, prefix: str) -> None: EMB_MAXN = 16 -class _LinearPackedEmbedding(torch.nn.Module): - def __init__(self) -> None: +def _write_embedding_goldens( + out_dir: str, + prefix: str, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + is_linear: bool, +) -> None: + for n in [EMB_MAXN, 8, 1]: + idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB + golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, is_linear + ) + idx.detach().numpy().astype(" None: super().__init__() packed = torch.arange(EMB_VOCAB * (EMB_DIM // 2), dtype=torch.int64).reshape( EMB_VOCAB, EMB_DIM // 2 ) self.register_buffer("weight", (packed % 256).to(torch.uint8)) self.register_buffer("scales", torch.ones(EMB_VOCAB, EMB_DIM // EMB_GROUP)) + self.is_linear_weight = is_linear_weight def forward(self, indices: torch.Tensor) -> torch.Tensor: return torch.ops.et_vk.embedding_q4gsw.default( - self.weight, self.scales, EMB_GROUP, indices, True + self.weight, + self.scales, + EMB_GROUP, + indices, + self.is_linear_weight, ) -def _write_embedding_goldens( - out_dir: str, - prefix: str, - weight: torch.Tensor, - scales: torch.Tensor, - group_size: int, - is_linear: bool, +def _write_packed_embedding_goldens( + out_dir: str, prefix: str, model: _PackedEmbedding ) -> None: for n in [EMB_MAXN, 8, 1]: idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB - golden = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, is_linear - ) - if is_linear: - nonlinear_golden = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, False - ) - if torch.equal(golden, nonlinear_golden): - raise RuntimeError( - "emb_dyn_linear fixture does not distinguish nibble packing" - ) + golden = model(idx) idx.detach().numpy().astype(" None: +def export_dynamic_embedding_cases(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, @@ -1370,21 +1429,22 @@ def _export_dynamic_embedding(out_dir: str) -> None: weight, scales, group_size = _quant_params(qm) _write_embedding_goldens(out_dir, "emb_dyn", weight, scales, group_size, False) - linear_model = _LinearPackedEmbedding().eval() - _export( - linear_model, - (idx_max,), - ({0: n_dim},), - os.path.join(out_dir, "emb_dyn_linear.pte"), - ) - _write_embedding_goldens( - out_dir, - "emb_dyn_linear", - linear_model.weight, - linear_model.scales, - EMB_GROUP, - True, - ) + packed_models = { + "emb_dyn_linear": _PackedEmbedding(True).eval(), + "emb_dyn_nonlinear": _PackedEmbedding(False).eval(), + } + linear_golden = packed_models["emb_dyn_linear"](idx_max) + nonlinear_golden = packed_models["emb_dyn_nonlinear"](idx_max) + if torch.equal(linear_golden, nonlinear_golden): + raise RuntimeError("embedding layout fixtures must distinguish nibble order") + for prefix, model in packed_models.items(): + _export( + model, + (idx_max,), + ({0: n_dim},), + os.path.join(out_dir, f"{prefix}.pte"), + ) + _write_packed_embedding_goldens(out_dir, prefix, model) # Dynamic RoPE: xq/xk + freqs all share a dynamic seq-len S. From 037005ead25bbc2c4d11ed7b362ff95bcf31755f Mon Sep 17 00:00:00 2001 From: pytorchbot Date: Fri, 7 Aug 2026 10:47:11 -0700 Subject: [PATCH 4/4] [ExecuTorch][Vulkan] Preserve persistent buffer mutations (#21662) This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: https://github.com/pytorch/executorch/pull/21597 by @JCNTH ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/204/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/204/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/203/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/204/orig Differential Revision: [D114936148](https://our.internmc.facebook.com/intern/diff/D114936148/) @diff-train-skip-merge cc @SS-JIA @manuelcandales @digantdesai @cbilgin --------- Co-authored-by: Julian Ng-Thow-Hing Co-authored-by: Julian Ng-Thow-Hing <107437036+JCNTH@users.noreply.github.com> --- .../serialization/vulkan_graph_builder.py | 48 +++- backends/vulkan/test/test_serialization.py | 138 +++++++++- .../test/test_vulkan_compile_options.py | 25 +- backends/vulkan/vulkan_preprocess.py | 5 + backends/webgpu/runtime/WebGPUDispatchMath.h | 2 +- backends/webgpu/runtime/WebGPUGraph.cpp | 98 +++++--- backends/webgpu/runtime/WebGPUGraph.h | 2 + .../webgpu/runtime/WebGPUShaderRegistry.cpp | 18 +- .../webgpu/runtime/ops/compare/Compare.cpp | 25 +- .../webgpu/runtime/ops/compare/compare.wgsl | 30 +-- .../webgpu/runtime/ops/compare/compare_wgsl.h | 32 +-- .../webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp | 237 ++++++++++++++++-- .../webgpu/runtime/ops/conv1d_dw/conv1d.wgsl | 53 ++++ .../runtime/ops/conv1d_dw/conv1d_wgsl.h | 77 ++++++ .../runtime/ops/expand_copy/ExpandCopy.cpp | 38 ++- .../runtime/ops/expand_copy/expand_copy.wgsl | 6 +- .../ops/expand_copy/expand_copy_wgsl.h | 8 +- backends/webgpu/runtime/ops/gelu/Gelu.cpp | 52 +++- backends/webgpu/runtime/ops/gelu/gelu.wgsl | 14 +- backends/webgpu/runtime/ops/gelu/gelu_wgsl.h | 16 +- .../webgpu/runtime/ops/to_copy/ToCopy.cpp | 196 ++++++++++++++- .../ops/to_copy/to_copy_bool_to_float.wgsl | 24 ++ .../ops/to_copy/to_copy_bool_to_float_wgsl.h | 48 ++++ .../test/native/test_compute_dispatch.cpp | 170 ++++++++++--- .../webgpu/test/native/test_dynamic_shape.cpp | 108 ++++++++ .../webgpu/test/native/test_webgpu_utils.cpp | 8 + backends/webgpu/test/op_tests/cases.py | 140 ++++++++++- .../webgpu/test/op_tests/generate_op_tests.py | 15 +- .../webgpu/test/op_tests/op_test_driver.cpp | 16 +- .../webgpu/test/op_tests/test_generator.py | 21 ++ backends/webgpu/test/op_tests/test_suite.py | 3 + .../test_dynamic_shape_export.py | 89 +++++++ backends/webgpu/test/ops/test_conv1d_pw.py | 82 +++++- backends/webgpu/test/ops/test_to_copy.py | 60 ++++- backends/webgpu/test/test_wgsl_codegen.py | 16 +- 35 files changed, 1720 insertions(+), 200 deletions(-) create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl create mode 100644 backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl create mode 100644 backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index ca5ab196dd2..46e01e701b1 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -9,7 +9,7 @@ import logging import operator from types import NoneType -from typing import cast, List, Optional, Union +from typing import cast, Dict, List, Optional, Union import executorch.backends.vulkan.serialization.vulkan_graph_schema as vk_graph_schema import torch @@ -28,6 +28,7 @@ ) from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.tensor import TensorSpec from torch._export.utils import get_buffer, get_param, is_buffer, is_param from torch.export import ExportedProgram @@ -49,11 +50,41 @@ def __init__( delegate_mapping_builder: DelegateMappingBuilder, downcast_64_bit: bool = True, force_fp16: bool = False, + alias_buffer_mutations: bool = False, ) -> None: self.program = program self.delegate_mapping_builder = delegate_mapping_builder self.downcast_64_bit = downcast_64_bit self.force_fp16 = force_fp16 + self.buffer_mutation_inputs: Dict[str, Node] = {} + self.buffer_mutation_user_outputs: set[str] = set() + if alias_buffer_mutations: + nodes_by_name = { + node.name: node for node in program.graph_module.graph.nodes + } + buffer_inputs_by_target: Dict[str, Node] = {} + for name, target in program.graph_signature.inputs_to_buffers.items(): + if name not in nodes_by_name: + continue + buffer_input = nodes_by_name[name] + prepack = next( + ( + user + for user in buffer_input.users + if user.op == "call_function" + and user.target == exir_ops.edge.et_vk.prepack.default + ), + None, + ) + buffer_inputs_by_target[target] = prepack or buffer_input + self.buffer_mutation_inputs = { + output_name: buffer_inputs_by_target[target] + for output_name, target in program.graph_signature.buffers_to_mutate.items() + if target in buffer_inputs_by_target + } + self.buffer_mutation_user_outputs = set( + program.graph_signature.user_outputs + ) self.chain = [] self.values = [] self.input_ids = [] @@ -160,6 +191,16 @@ def maybe_add_constant_tensor(self, node: Node) -> int: return constant_id def create_node_value(self, node: Node) -> int: + if node.name in self.buffer_mutation_inputs: + input_node = self.buffer_mutation_inputs[node.name] + if input_node not in self.node_to_value_ids: + raise AssertionError( + "Cannot alias a buffer mutation before its input is serialized" + ) + value_id = self.node_to_value_ids[input_node] + self.node_to_value_ids[node] = value_id + return value_id + # If the node has been marked as a scalar tensor, create a SymInt instead of a tensor if is_symint_node(node) or node.meta.get("etvk_is_scalar_tensor", False): new_id = self.create_symint_value() @@ -448,7 +489,10 @@ def process_output_node(self, node: Node) -> None: ) # Mutable buffers outputs are not included as an output to the # delegate call. Skip marking them as an output. - if is_mutable_buffer_node(out_node, self.program): + if out_node.name in self.buffer_mutation_inputs: + if out_node.name not in self.buffer_mutation_user_outputs: + continue + elif is_mutable_buffer_node(out_node, self.program): continue self.output_ids.append(self.node_to_value_ids[out_node]) diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index c373f5216d2..71a6980635a 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -9,9 +9,14 @@ import ctypes import random import unittest -from typing import List +from types import SimpleNamespace +from typing import List, Tuple +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 import torch +from executorch.backends.vulkan.serialization import ( + vulkan_graph_builder as graph_builder_module, +) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( IntList, @@ -30,6 +35,137 @@ class TestSerialization(unittest.TestCase): + def _build_mutation_program( + self, prepack: bool, shared_user_output: bool = False + ) -> Tuple[SimpleNamespace, torch.fx.Node, torch.fx.Node, torch.fx.Node]: + graph = torch.fx.Graph() + state = graph.placeholder("state") + user_input = graph.placeholder("user_input") + state.meta["spec"] = graph_builder_module.TensorSpec.from_tensor(torch.zeros(4)) + user_input.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + + state_value = state + if prepack: + state_value = graph.call_function( + graph_builder_module.exir_ops.edge.et_vk.prepack.default, + (state,), + ) + state_value.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.zeros(4) + ) + + mutation = graph.call_function( + torch.ops.aten.add.Tensor, (state_value, user_input) + ) + mutation.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + user_output = mutation + if not shared_user_output: + user_output = graph.call_function( + torch.ops.aten.mul.Tensor, (user_input, 2.0) + ) + user_output.meta["spec"] = graph_builder_module.TensorSpec.from_tensor( + torch.ones(4) + ) + graph.output((mutation, user_output)) + + graph_module = torch.fx.GraphModule({}, graph) + signature = SimpleNamespace( + buffers_to_mutate={mutation.name: "state"}, + inputs_to_buffers={state.name: "state"}, + inputs_to_lifted_tensor_constants={}, + inputs_to_parameters={}, + non_persistent_buffers=set(), + user_outputs=(user_output.name,), + ) + program = SimpleNamespace( + constants={}, + graph_module=graph_module, + graph_signature=signature, + state_dict={"state": torch.zeros(4)}, + ) + return program, state_value, mutation, user_output + + def test_alias_buffer_mutations_is_opt_in(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, user_output = ( + self._build_mutation_program(prepack) + ) + + default_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + ) + default_graph = default_builder.build_graph() + self.assertNotEqual( + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + default_graph.output_ids, + [ + default_builder.node_to_value_ids[mutation], + default_builder.node_to_value_ids[user_output], + ], + ) + + explicit_false_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=False, + ) + self.assertEqual(default_graph, explicit_false_builder.build_graph()) + + aliasing_builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + aliasing_graph = aliasing_builder.build_graph() + self.assertEqual( + aliasing_builder.node_to_value_ids[mutation], + aliasing_builder.node_to_value_ids[state_value], + ) + self.assertEqual( + aliasing_graph.output_ids, + [aliasing_builder.node_to_value_ids[user_output]], + ) + + def test_alias_buffer_mutations_preserves_shared_user_output(self) -> None: + for prepack in (False, True): + with self.subTest(prepack=prepack): + program, state_value, mutation, _ = self._build_mutation_program( + prepack, shared_user_output=True + ) + builder = graph_builder_module.VkGraphBuilder( + program, + graph_builder_module.DelegateMappingBuilder( + generated_identifiers=True + ), + alias_buffer_mutations=True, + ) + + graph = builder.build_graph() + + self.assertEqual( + builder.node_to_value_ids[mutation], + builder.node_to_value_ids[state_value], + ) + self.assertEqual( + graph.output_ids, + [builder.node_to_value_ids[mutation]], + ) + def _generate_random_const_tensors(self, num_tensors: int) -> List[torch.Tensor]: """ Helper function to generate `num_tensor` buffers of random sizes and random contents, diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index f45cfdf12d9..ff507bb3197 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/backends/vulkan/test/test_vulkan_compile_options.py @@ -41,6 +41,10 @@ def test_skip_memory_planning_round_trips(self) -> None: round_tripped = self._round_trip({"skip_memory_planning": True}) self.assertTrue(round_tripped.get("skip_memory_planning")) + def test_alias_buffer_mutations_round_trips(self) -> None: + round_tripped = self._round_trip({"alias_buffer_mutations": True}) + self.assertTrue(round_tripped.get("alias_buffer_mutations")) + def test_force_fp16_round_trips(self) -> None: round_tripped = self._round_trip({"force_fp16": True}) self.assertTrue(round_tripped.get("force_fp16")) @@ -105,15 +109,15 @@ def build_graph(): ), patch( "executorch.backends.vulkan.vulkan_preprocess.VkGraphBuilder", return_value=graph_builder, - ), patch( + ) as graph_builder_factory, patch( "executorch.backends.vulkan.vulkan_preprocess.serialize_vulkan_graph", return_value=b"vk_graph", ): result = VulkanBackend.preprocess(program, parse_compile_options(options)) - return result.data_store_output, externalize_pte_data + return result.data_store_output, externalize_pte_data, graph_builder_factory def test_external_constants_default_keeps_constants_inline(self) -> None: - output, externalize_pte_data = self._preprocess_named_data({}) + output, externalize_pte_data, _ = self._preprocess_named_data({}) self.assertEqual(output.buffers, [b"constant"]) self.assertEqual(output.pte_data, {"constant": DataEntry(0, 16, None)}) @@ -121,7 +125,7 @@ def test_external_constants_default_keeps_constants_inline(self) -> None: externalize_pte_data.assert_not_called() def test_external_constants_option_externalizes_constants(self) -> None: - output, externalize_pte_data = self._preprocess_named_data( + output, externalize_pte_data, _ = self._preprocess_named_data( {"external_constants_max_data_bytes": 16} ) @@ -131,8 +135,21 @@ def test_external_constants_option_externalizes_constants(self) -> None: self.assertEqual(list(next(iter(output.external_data.values()))), ["constant"]) externalize_pte_data.assert_called_once_with(16, "vulkan_constants") + def test_alias_buffer_mutations_reaches_graph_builder(self) -> None: + for options, expected in ( + ({}, False), + ({"alias_buffer_mutations": True}, True), + ): + with self.subTest(options=options): + _, _, graph_builder_factory = self._preprocess_named_data(options) + self.assertIs( + graph_builder_factory.call_args.kwargs["alias_buffer_mutations"], + expected, + ) + def test_unset_options_are_absent(self) -> None: round_tripped = self._round_trip({}) + self.assertNotIn("alias_buffer_mutations", round_tripped) self.assertNotIn("small_texture_limits", round_tripped) self.assertNotIn("skip_memory_planning", round_tripped) self.assertNotIn("external_constants_max_data_bytes", round_tripped) diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index d3954240880..f7d6955ce26 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -130,6 +130,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: if spec.key == "skip_memory_planning": options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "alias_buffer_mutations": + options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "external_constants_max_data_bytes": options[spec.key] = _parse_external_constants_max_data_bytes(spec.value) @@ -172,6 +175,7 @@ def preprocess( # noqa: C901 ) downcast_64_bit = compile_options.get("downcast_64_bit", True) force_fp16 = compile_options.get("force_fp16", False) + alias_buffer_mutations = compile_options.get("alias_buffer_mutations", False) program = unsafe_remove_auto_functionalized_pass(program) @@ -258,6 +262,7 @@ def preprocess( # noqa: C901 DelegateMappingBuilder(generated_identifiers=True), downcast_64_bit=downcast_64_bit, force_fp16=force_fp16, + alias_buffer_mutations=alias_buffer_mutations, ) vk_graph = graph_builder.build_graph() external_constants_max_data_bytes = compile_options.get( diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 561bebc6b87..60638b499bb 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -24,7 +24,7 @@ namespace executorch::backends::webgpu::utils { // Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). template inline T div_up(T a, T b) { - return (a + b - 1) / b; + return a / b + (a % b != 0); } // Product of a tensor's dims; the same accumulation was duplicated per-op. diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 58d7178ddbf..d1e1d625ad7 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -212,6 +212,31 @@ bool vk_datatype_is_int(vkgraph::VkDataType dtype) { } } +size_t storage_buffer_size(size_t nbytes) { + const size_t at_least_four = std::max(nbytes, size_t(4)); + if (at_least_four > std::numeric_limits::max() - 3u) { + throw std::runtime_error("WebGPU: storage buffer size overflows alignment"); + } + return (at_least_four + 3u) & ~size_t(3); +} + +void write_storage_buffer( + WGPUQueue queue, + WGPUBuffer buffer, + const void* data, + size_t nbytes) { + if (nbytes == 0u) { + return; + } + if (nbytes % 4u == 0u) { + wgpuQueueWriteBuffer(queue, buffer, 0, data, nbytes); + return; + } + std::vector padded(storage_buffer_size(nbytes), 0u); + std::memcpy(padded.data(), data, nbytes); + wgpuQueueWriteBuffer(queue, buffer, 0, padded.data(), padded.size()); +} + // Normalize a possibly-negative dim against rank; throws (fail-loud) if OOR. int normalize_dim(int dim, int rank, const char* op) { if (dim < 0) { @@ -276,7 +301,7 @@ WebGPUGraph::WebGPUGraph() = default; WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = nbytes > 0 ? nbytes : 4; + buf_desc.size = storage_buffer_size(nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -286,7 +311,7 @@ WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { } WGPUBuffer WebGPUGraph::acquire_scratch(size_t nbytes) { - nbytes = nbytes > 0 ? nbytes : 4; + nbytes = storage_buffer_size(nbytes); // Best-fit reuse: smallest free slot with size in [nbytes, 2*nbytes] -- the // 2x cap stops a large Cmax-sized buffer from backing a tiny request. Never // reuse an in_use slot (co-live safety). @@ -890,6 +915,7 @@ void WebGPUGraph::build( throw std::runtime_error("WebGPU: tensor byte size overflows"); } tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); + tensor.is_bool = vk_tensor->datatype() == vkgraph::VkDataType::BOOL; tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; // Live dims start == max (serialized upper bound); resize_input shrinks @@ -911,7 +937,7 @@ void WebGPUGraph::build( tensor.cur_nbytes = tensor.nbytes; tensor_mem_obj_ids_[i] = -1; WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -936,10 +962,9 @@ void WebGPUGraph::build( std::memcpy(&value, src + e * sizeof(float), sizeof(float)); converted[e] = executorch::runtime::etensor::Half(value); } - wgpuQueueWriteBuffer( + write_storage_buffer( queue_, tensor.buffer, - 0, converted.data(), converted.size() * sizeof(converted[0])); }; @@ -1013,7 +1038,7 @@ void WebGPUGraph::build( prepack_src_ids.count(i) != 0 && direct_use_ids.count(i) == 0; if (!defer) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -1110,7 +1135,7 @@ void WebGPUGraph::build( shared_buffers_.resize(shared_buffer_sizes_.size(), nullptr); for (size_t id = 0; id < shared_buffer_sizes_.size(); id++) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(shared_buffer_sizes_[id], size_t(4)); + buf_desc.size = storage_buffer_size(shared_buffer_sizes_[id]); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -1138,7 +1163,7 @@ void WebGPUGraph::build( // Create staging buffer for output readback WGPUBufferDescriptor staging_desc = {}; - staging_desc.size = std::max(tensors_[oid].nbytes, size_t(4)); + staging_desc.size = storage_buffer_size(tensors_[oid].nbytes); staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; staging_desc.mappedAtCreation = false; output_staging_buffers_.push_back( @@ -1313,7 +1338,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { cs.nbytes, "WebGPU: inline constant exceeds constant data"); if (cs.nbytes != 0) { - wgpuQueueWriteBuffer(queue_, dst, 0, data, cs.nbytes); + write_storage_buffer(queue_, dst, data, cs.nbytes); } } else if (cs.nbytes == 0) { return; @@ -1327,7 +1352,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { throw std::runtime_error( "WebGPU: named constant '" + cs.named_key + "' undersized"); } - wgpuQueueWriteBuffer(queue_, dst, 0, buf->data(), cs.nbytes); + write_storage_buffer(queue_, dst, buf->data(), cs.nbytes); buf->Free(); } else { throw std::runtime_error("WebGPU: constant has no source"); @@ -1419,7 +1444,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { // Fast path: host and GPU element types match byte-for-byte. if (in.nbytes == live_nbytes) { - wgpuQueueWriteBuffer(queue_, tensor.buffer, 0, in.data, live_nbytes); + write_storage_buffer(queue_, tensor.buffer, in.data, live_nbytes); continue; } @@ -1439,8 +1464,21 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { #endif narrowed[e] = static_cast(src[e]); } - wgpuQueueWriteBuffer( - queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); + continue; + } + + // Require an explicit fp32 host dtype, not merely "not int64": inferring + // the narrow from the 2:1 byte ratio alone would silently reinterpret a + // same-sized non-fp32 host buffer (e.g. a stale int32) as fp32. + if (in.host_is_fp32 && buffer_is_fp16 && in.nbytes == live_nbytes * 2) { + const size_t numel = live_nbytes / sizeof(uint16_t); + const float* src = static_cast(in.data); + std::vector narrowed(numel); + for (size_t e = 0; e < numel; e++) { + narrowed[e] = executorch::runtime::etensor::Half(src[e]); + } + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); continue; } @@ -1684,10 +1722,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { } for (size_t i = 0; i < output_copies_.size(); i++) { - const size_t copy_nbytes = tensors_[output_ids_[i]].cur_nbytes; - if (!plan.copy_outputs[i] || copy_nbytes == 0) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { continue; } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); @@ -1715,12 +1754,13 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { return 1; } - // GPU timestamp queries assume one submit; chunked execute is multi-submit. +#ifdef WGPU_BACKEND_ENABLE_PROFILING if (should_timestamp_query()) { throw std::runtime_error( "WebGPU: WEBGPU_TIMESTAMP_QUERY is incompatible with chunked execute " "(multi-submit); disable chunking to use GPU timestamp queries"); } +#endif // WGPU_BACKEND_ENABLE_PROFILING for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size(); chunk_index++) { @@ -1759,10 +1799,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { if (chunk_index + 1 == plan.dispatch_chunks.size()) { for (size_t i = 0; i < output_copies_.size(); i++) { - const size_t copy_nbytes = tensors_[output_ids_[i]].cur_nbytes; - if (!plan.copy_outputs[i] || copy_nbytes == 0) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { continue; } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); @@ -1813,13 +1854,13 @@ void WebGPUGraph::copy_outputs( continue; } const auto& tensor = tensors_[output_ids_[i]]; - const size_t map_nbytes = tensor.cur_nbytes; - if (map_nbytes == 0) { + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { continue; } const size_t dst_nbytes = outputs[i].nbytes; const bool is_double_width = - dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; const bool widen_fp16 = is_double_width && !tensor.is_int && tensor.elem_size == 2; const bool widen_int32 = @@ -1832,7 +1873,7 @@ void WebGPUGraph::copy_outputs( if (outputs[i].host_is_fp32 && buffer_is_fp16 && !widen_fp16) { throw std::runtime_error("WebGPU: fp16 output buffer size mismatch"); } - if (dst_nbytes != map_nbytes && !widen_fp16 && !widen_int32) { + if (dst_nbytes != logical_nbytes && !widen_fp16 && !widen_int32) { throw std::runtime_error("WebGPU: output buffer size mismatch"); } } @@ -1842,13 +1883,14 @@ void WebGPUGraph::copy_outputs( continue; } const auto& tensor = tensors_[output_ids_[i]]; - const size_t map_nbytes = tensor.cur_nbytes; - if (map_nbytes == 0) { + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { continue; } + const size_t map_nbytes = storage_buffer_size(logical_nbytes); const size_t dst_nbytes = outputs[i].nbytes; const bool is_double_width = - dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; const bool widen_fp16 = is_double_width && !tensor.is_int && tensor.elem_size == 2; const bool widen_int32 = @@ -1887,7 +1929,7 @@ void WebGPUGraph::copy_outputs( const auto* src = static_cast(mapped); auto* dst = static_cast(outputs[i].data); - const size_t n = map_nbytes / sizeof(*src); + const size_t n = logical_nbytes / sizeof(*src); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } @@ -1895,12 +1937,12 @@ void WebGPUGraph::copy_outputs( // int64 host output backed by an int32 GPU buffer: widen (sign-extend). const int32_t* src = static_cast(mapped); int64_t* dst = static_cast(outputs[i].data); - const size_t n = map_nbytes / sizeof(int32_t); + const size_t n = logical_nbytes / sizeof(int32_t); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } } else { - std::memcpy(outputs[i].data, mapped, map_nbytes); + std::memcpy(outputs[i].data, mapped, logical_nbytes); } wgpuBufferUnmap(output_staging_buffers_[i]); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index a634ebb4172..23ce9df03ed 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -39,6 +39,8 @@ struct WebGPUTensor { // Serialized (GPU-side) element type, used to narrow wider host inputs. size_t elem_size = 0; bool is_int = false; + // Exact BOOL tag for byte-packed WGSL storage. + bool is_bool = false; // Exactly int8 (not uint8/bool), so int8-only ops can guard their dtype. bool is_int8 = false; }; diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 5374390722b..480944ea93d 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -124,6 +125,7 @@ #include #include #include +#include #include #include #include @@ -151,7 +153,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -362,6 +364,13 @@ constexpr std::array kShaderRegistry = {{ kConstantPadNdWorkgroupSizeY, kConstantPadNdWorkgroupSizeZ, }, + { + "conv1d", + kConv1dWGSL, + kConv1dWorkgroupSizeX, + kConv1dWorkgroupSizeY, + kConv1dWorkgroupSizeZ, + }, { "conv1d_dw", kConv1dDwWGSL, @@ -1034,6 +1043,13 @@ constexpr std::array kShaderRegistry = {{ kTanhWorkgroupSizeY, kTanhWorkgroupSizeZ, }, + { + "to_copy_bool_to_float", + kToCopyBoolToFloatWGSL, + kToCopyBoolToFloatWorkgroupSizeX, + kToCopyBoolToFloatWorkgroupSizeY, + kToCopyBoolToFloatWorkgroupSizeZ, + }, { "to_copy_float_to_int", kToCopyFloatToIntWGSL, diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index 99200cb2425..1870ee85213 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -58,13 +58,13 @@ void compare_impl( in2_tensor.elem_size != 4) { throw std::runtime_error("compare: fp32 inputs only"); } - if (!out_tensor.is_int || out_tensor.elem_size != 1) { + if (!out_tensor.is_bool || out_tensor.elem_size != 1) { throw std::runtime_error("compare: output must be a 1-byte bool tensor"); } const uint64_t numel = out_tensor.nbytes; - // out bool packed 4/word (array); numel%4==0 gates the readback map. - if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { - throw std::runtime_error("compare: numel must be a nonzero mult of 4"); + // out bool is byte-packed into ceil(numel / 4) u32 storage words. + if (numel == 0u || numel > UINT32_MAX) { + throw std::runtime_error("compare: numel must be nonzero and fit u32"); } const uint64_t in_numel = in1_tensor.nbytes / sizeof(float); if (in1_tensor.nbytes != in2_tensor.nbytes || in_numel != numel) { @@ -75,7 +75,7 @@ void compare_impl( params.num_elements = static_cast(numel); params.op = op; - const uint32_t words = static_cast(numel / 4u); + const uint32_t words = static_cast((numel + 3u) / 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX); utils::WgCount workgroup_count = @@ -85,9 +85,7 @@ void compare_impl( wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(CompareParams)); - graph.add_uniform_buffer_bytes(sizeof(CompareParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // out (rw storage) + in1/in2 (ro storage) + params (uniform). utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( @@ -97,7 +95,7 @@ void compare_impl( {0, WGPUBufferBindingType_Storage, out_tensor.buffer, - out_tensor.nbytes}, + static_cast(words) * sizeof(uint32_t)}, {1, WGPUBufferBindingType_ReadOnlyStorage, in1_tensor.buffer, @@ -127,9 +125,8 @@ void compare_impl( WebGPUGraph& g) { const auto& d = g.cur_dims(in1_id); const uint64_t n = utils::numel_of(d); - if (n == 0u || n % 4u != 0u || n > UINT32_MAX || - utils::numel_of(g.cur_dims(in2_id)) != n) { - throw std::runtime_error("compare(resize): numel must be a mult of 4"); + if (n == 0u || n > UINT32_MAX || utils::numel_of(g.cur_dims(in2_id)) != n) { + throw std::runtime_error("compare(resize): invalid numel"); } g.set_cur_dims(out_id, d); CompareParams p = {}; @@ -137,14 +134,12 @@ void compare_impl( p.op = op; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), static_cast(n / 4u), wg_size, "compare"); + g.device(), static_cast((n + 3u) / 4u), wg_size, "compare"); 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, resize); graph.add_tensor_resize_hook(in2_id, resize); - - graph.own_uniform_buffer(uniform_buffer); } void eq_op(WebGPUGraph& graph, const std::vector& args) { diff --git a/backends/webgpu/runtime/ops/compare/compare.wgsl b/backends/webgpu/runtime/ops/compare/compare.wgsl index a16568e5533..a650a6298b5 100644 --- a/backends/webgpu/runtime/ops/compare/compare.wgsl +++ b/backends/webgpu/runtime/ops/compare/compare.wgsl @@ -16,27 +16,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/compare/compare_wgsl.h b/backends/webgpu/runtime/ops/compare/compare_wgsl.h index 672c99b62d8..c1c4ac23e4f 100644 --- a/backends/webgpu/runtime/ops/compare/compare_wgsl.h +++ b/backends/webgpu/runtime/ops/compare/compare_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from compare.wgsl - DO NOT EDIT. -// wgsl-sha256: 241e7e6762b1eded07d28a3767936c970f509f6591e7cbf599d0b1eb61efb181 +// wgsl-sha256: 8f330b5a1e29a1fb8135e64600dadb1dc64c98f8f5371357f8de73a1808b76d6 inline constexpr const char* kCompareWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var input1: array; @@ -33,27 +33,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index ae86ccffba8..64d8c2d2380 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include @@ -49,8 +50,24 @@ uint32_t conv1d_out_len( int64_t stride, int64_t padding, int64_t dilation) { - return static_cast( - (in_len + 2 * padding - dilation * (k - 1) - 1) / stride + 1); + if (in_len <= 0 || k <= 0 || stride <= 0 || padding < 0 || dilation <= 0) { + throw std::runtime_error("conv1d: invalid geometry parameter"); + } + constexpr int64_t kMaxShaderIndex = std::numeric_limits::max(); + if (in_len > kMaxShaderIndex || k > kMaxShaderIndex || + stride > kMaxShaderIndex || padding > kMaxShaderIndex || + dilation > kMaxShaderIndex) { + throw std::runtime_error("conv1d: geometry parameter exceeds i32"); + } + const int64_t numerator = in_len + 2 * padding - dilation * (k - 1) - 1; + if (numerator < 0) { + throw std::runtime_error("conv1d: kernel exceeds padded input"); + } + const int64_t out_len = numerator / stride + 1; + if (static_cast(out_len) > UINT32_MAX) { + throw std::runtime_error("conv1d: output length exceeds u32"); + } + return static_cast(out_len); } int64_t first_int(const std::vector& v) { @@ -71,6 +88,24 @@ static_assert( sizeof(Conv1dPwParams) == 32, "Conv1dPwParams must match the WGSL Params struct (32 bytes)"); +struct Conv1dParams { + uint32_t in_channels; + uint32_t out_channels; + uint32_t in_len; + uint32_t out_len; + uint32_t kernel_size; + uint32_t stride; + uint32_t padding; + uint32_t dilation; + uint32_t numel; + uint32_t has_bias; +}; +static_assert( + sizeof(Conv1dParams) == 40, + "Conv1dParams must match the WGSL Params struct (40 bytes)"); +constexpr uint64_t kMaxConv1dDispatchElements = + static_cast(std::numeric_limits::max()); + // Pointwise conv1d (K=1, groups=1): a per-position matmul over channels. void add_conv1d_pw_node( WebGPUGraph& graph, @@ -203,6 +238,154 @@ void add_conv1d_pw_node( graph.own_uniform_buffer(params_buf); } +// General groups=1 conv1d. Voxtral uses K=3 with stride 1 then 2. +void add_conv1d_node( + WebGPUGraph& graph, + int in_id, + int weight_id, + int bias_id, + int out_id, + uint32_t stride, + uint32_t padding, + uint32_t dilation) { + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + if (!utils::is_fp32_tensor(in) || !utils::is_fp32_tensor(weight) || + !utils::is_fp32_tensor(out)) { + throw std::runtime_error("conv1d: input, weight, and output must be fp32"); + } + + const uint32_t expected_out_len = conv1d_out_len( + in.dims.at(2), weight.dims.at(2), stride, padding, dilation); + const uint32_t batch = static_cast(in.dims.at(0)); + const uint32_t in_channels = static_cast(in.dims.at(1)); + const uint32_t in_len = static_cast(in.dims.at(2)); + const uint32_t out_channels = static_cast(out.dims.at(1)); + const uint32_t out_len = static_cast(out.dims.at(2)); + const uint32_t kernel_size = static_cast(weight.dims.at(2)); + if (out.dims.at(0) != in.dims.at(0) || out_len != expected_out_len || + weight.dims.at(0) != out.dims.at(1) || + weight.dims.at(1) != in.dims.at(1)) { + throw std::runtime_error("conv1d: shape mismatch"); + } + + const uint64_t in_numel = utils::check_fp32(in, "conv1d", "input"); + const uint64_t out_numel = utils::check_fp32(out, "conv1d", "output"); + const uint64_t weight_numel = utils::check_fp32(weight, "conv1d", "weight"); + if (in_numel != static_cast(batch) * in_channels * in_len || + out_numel != static_cast(batch) * out_channels * out_len || + weight_numel != + static_cast(out_channels) * in_channels * kernel_size || + in_numel > UINT32_MAX || weight_numel > UINT32_MAX || + out_numel > kMaxConv1dDispatchElements) { + throw std::runtime_error("conv1d: fp32 byte-size or u32 mismatch"); + } + if (has_bias) { + const auto& bias = graph.get_tensor(bias_id); + if (!utils::is_fp32_tensor(bias) || bias.dims.size() != 1 || + bias.dims.at(0) != out.dims.at(1) || + utils::check_fp32(bias, "conv1d", "bias") != out_channels) { + throw std::runtime_error("conv1d: bias shape mismatch"); + } + } + + Conv1dParams params = {}; + params.in_channels = in_channels; + params.out_channels = out_channels; + params.in_len = in_len; + params.out_len = out_len; + params.kernel_size = kernel_size; + params.stride = stride; + params.padding = padding; + params.dilation = dilation; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dWorkgroupSizeX); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.numel, wg_size, "conv1d"); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); + WGPUBuffer params_buf = graph.create_params_buffer(params); + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : weight.nbytes; + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv1dWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(Conv1dParams)}, + }, + &wg_size_constant, + 1); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "conv1d", + workgroup_count.y}); + + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + has_bias, + wg_size, + dispatch_idx, + params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + if (dims.size() != 3 || dims[0] <= 0 || dims[1] <= 0 || dims[2] <= 0 || + dims[1] != static_cast(in_channels)) { + throw std::runtime_error("conv1d(resize): input shape changed"); + } + Conv1dParams p = {}; + p.in_channels = in_channels; + p.out_channels = out_channels; + p.in_len = static_cast(dims[2]); + p.out_len = + conv1d_out_len(dims[2], kernel_size, stride, padding, dilation); + p.kernel_size = kernel_size; + p.stride = stride; + p.padding = padding; + p.dilation = dilation; + const uint64_t input_numel = utils::numel(dims); + const uint64_t numel = utils::numel( + {dims[0], out_channels, static_cast(p.out_len)}); + if (input_numel > UINT32_MAX || numel > kMaxConv1dDispatchElements) { + throw std::runtime_error( + "conv1d(resize): tensor numel exceeds shader index range"); + } + p.numel = static_cast(numel); + p.has_bias = has_bias ? 1u : 0u; + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d(resize)"); + g.set_cur_dims( + out_id, {dims[0], out_channels, static_cast(p.out_len)}); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); +} + // depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). void convolution_impl(WebGPUGraph& graph, const std::vector& args) { // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). @@ -242,38 +425,44 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { const bool transposed = graph.get_bool(transposed_id); const int64_t groups = graph.get_int(groups_id); - // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. - if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && - first_int(graph.get_int_list(stride_id)) == 1 && - first_int(graph.get_int_list(padding_id)) == 0) { - add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); - return; - } - - // Otherwise only the depthwise config (groups==C, weight [C,1,K]). - if (transposed || groups != static_cast(channels) || - weight_tensor.dims.at(0) != static_cast(channels) || - weight_tensor.dims.at(1) != 1) { - throw std::runtime_error( - "convolution: only depthwise or pointwise conv1d supported"); - } - const int64_t stride_i = first_int(graph.get_int_list(stride_id)); const int64_t padding_i = first_int(graph.get_int_list(padding_id)); const int64_t dilation_i = first_int(graph.get_int_list(dilation_id)); - if (stride_i < 1) { - throw std::runtime_error("convolution: stride must be >= 1"); + if (stride_i < 1 || stride_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: stride must fit positive i32"); } - if (padding_i < 0) { - throw std::runtime_error("convolution: padding must be >= 0"); + if (padding_i < 0 || padding_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: padding must fit nonnegative i32"); } - if (dilation_i < 1) { - throw std::runtime_error("convolution: dilation must be >= 1"); + if (dilation_i < 1 || dilation_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: dilation must fit positive i32"); } const uint32_t stride = static_cast(stride_i); const uint32_t padding = static_cast(padding_i); const uint32_t dilation = static_cast(dilation_i); + // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. + if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && + stride_i == 1 && padding_i == 0) { + add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); + return; + } + + const bool is_depthwise = !transposed && + groups == static_cast(channels) && + weight_tensor.dims.at(0) == static_cast(channels) && + weight_tensor.dims.at(1) == 1; + if (!is_depthwise && !transposed && groups == 1) { + add_conv1d_node( + graph, in_id, weight_id, bias_id, out_id, stride, padding, dilation); + return; + } + + if (!is_depthwise) { + throw std::runtime_error( + "convolution: only depthwise, pointwise, or groups=1 conv1d supported"); + } + uint64_t out_numel = 1; for (int64_t d : out_tensor.dims) { out_numel *= static_cast(d); diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl new file mode 100644 index 00000000000..c51594d2005 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl @@ -0,0 +1,53 @@ +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +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.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h new file mode 100644 index 00000000000..e29af80459b --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h @@ -0,0 +1,77 @@ +/* + * 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 conv1d.wgsl - DO NOT EDIT. +// wgsl-sha256: 7bc955b7f43473aab96222e7a1228973b65e38e2e6a19e2b9793e0ed7f3768d9 +inline constexpr const char* kConv1dWGSL = R"( +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +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.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} +)"; + +inline constexpr uint32_t kConv1dWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp index 5bfc0fa3bf1..9f6a4e68ba6 100644 --- a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -14,6 +14,7 @@ #include +#include #include namespace executorch::backends::webgpu { @@ -34,6 +35,22 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (graph.get_value_type(args.at(1)) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic target sizes are unsupported"); + } + for (int64_t target_size : graph.get_int_list(args.at(1))) { + if (target_size == -1) { + throw std::runtime_error( + "WebGPU expand_copy: inferred target sizes are unsupported"); + } + } + if (graph.tensor_has_dynamic_dims(in_id) || + graph.tensor_has_dynamic_dims(out_id)) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic shapes are unsupported"); + } + TensorMeta out_meta; TensorMeta in_meta; fill_tensor_meta(out_tensor, &out_meta); @@ -44,21 +61,24 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "expand_copy: non-fp32 operand (nbytes != numel*4)"); } + if (out_meta.numel > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU expand_copy: element count exceeds the flattened 2D dispatch " + "limit"); + } uint32_t wg_size = utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "expand_copy"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer out_meta_buf = - utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); - WGPUBuffer in_meta_buf = - utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); - graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); + WGPUBuffer out_meta_buf = graph.create_params_buffer(out_meta); + WGPUBuffer in_meta_buf = graph.create_params_buffer(in_meta); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -78,10 +98,8 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - wgpuBufferRelease(out_meta_buf); - wgpuBufferRelease(in_meta_buf); + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); } } // namespace diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl index 053311a69f4..fab4df15a90 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl @@ -13,8 +13,10 @@ struct TensorMeta { 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 >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h index 83c5881604f..f1449f61793 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from expand_copy.wgsl - DO NOT EDIT. -// wgsl-sha256: 99953670bea89e42bc9c689ab80addfd9a331442c8c8f1a5b0c39dbe11c19370 +// wgsl-sha256: b3c032ab961ffde245fc44289b67df3b5e4ca93eedb9ada2f20a3eaa6f10e9c6 inline constexpr const char* kExpandCopyWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -30,8 +30,10 @@ struct TensorMeta { 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 >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/Gelu.cpp b/backends/webgpu/runtime/ops/gelu/Gelu.cpp index 943012515ff..023ed7b5d88 100644 --- a/backends/webgpu/runtime/ops/gelu/Gelu.cpp +++ b/backends/webgpu/runtime/ops/gelu/Gelu.cpp @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -42,13 +43,18 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { const auto& out_tensor = graph.get_tensor(out_id); utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu"); - uint32_t num_elements = - static_cast(out_tensor.nbytes / sizeof(float)); + const uint64_t num_elements64 = out_tensor.nbytes / sizeof(float); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu: element count exceeds the flattened 2D dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom). uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, num_vec4_threads, wg_size, "gelu"); WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size); @@ -56,9 +62,7 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { GeluParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(GeluParams)); - graph.add_uniform_buffer_bytes(sizeof(GeluParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // input (read storage) + output (storage) + params. The exact/approximate // choice is baked into the compiled pipeline via the entry point (mirrors @@ -85,10 +89,38 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { 1, exact ? "main_erf" : "main_tanh"); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t num_elements64 = utils::numel_of(dims); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu(resize): element count exceeds the flattened 2D " + "dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); + g.set_cur_dims(out_id, dims); + + GeluParams params = {}; + params.num_elements = num_elements; + wgpuQueueWriteBuffer( + g.queue(), params_buf, 0, ¶ms, sizeof(GeluParams)); + + const uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); + const utils::WgCount resized_workgroup_count = + utils::compute_2d_workgroup_count( + g.device(), num_vec4_threads, wg_size, "gelu(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = + resized_workgroup_count.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = + resized_workgroup_count.y; + }); } } // namespace diff --git a/backends/webgpu/runtime/ops/gelu/gelu.wgsl b/backends/webgpu/runtime/ops/gelu/gelu.wgsl index 4f7eb68bc96..9583ef81551 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu.wgsl +++ b/backends/webgpu/runtime/ops/gelu/gelu.wgsl @@ -33,8 +33,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -49,8 +52,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h index 6da12e229af..f8af0f8d2c3 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h +++ b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from gelu.wgsl - DO NOT EDIT. -// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f +// wgsl-sha256: 96570753688590fa009ee5503f754cf3eb572dcb3dcae6818220fe06fe3139ee inline constexpr const char* kGeluWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -50,8 +50,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -66,8 +69,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp index 05f7ff6fe24..54d80016296 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -10,12 +10,14 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -68,9 +70,7 @@ void add_convert_op( ConvertParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); - graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -113,8 +113,185 @@ void add_convert_op( wg_size, "to_copy(resize)"); }); +} + +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + 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)"); + }); +} + +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + 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)"); + }); - // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } @@ -129,6 +306,12 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.is_bool != out_tensor.is_bool && in_tensor.is_int && + out_tensor.is_int) { + throw std::runtime_error( + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain. if (in_tensor.is_int == out_tensor.is_int && in_tensor.elem_size == out_tensor.elem_size) { @@ -137,7 +320,10 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { } // int<->float = numeric convert (mirrors Vulkan add_view_copy_convert_node). - if (in_tensor.is_int && !out_tensor.is_int) { + if (in_tensor.is_bool && !out_tensor.is_int && + out_tensor.elem_size == sizeof(float)) { + add_bool_to_float_op(graph, in_id, out_id); + } else if (in_tensor.is_int && !out_tensor.is_int) { add_convert_op( graph, in_id, 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 new file mode 100644 index 00000000000..239730de65d --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl @@ -0,0 +1,24 @@ +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} 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 new file mode 100644 index 00000000000..ef7e40976f3 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h @@ -0,0 +1,48 @@ +/* + * 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 to_copy_bool_to_float.wgsl - DO NOT EDIT. +// wgsl-sha256: 29fd43b2f638489e9b8d72b2cc9140d07174c750cbdc291e63793319a2fa5961 +inline constexpr const char* kToCopyBoolToFloatWGSL = R"( +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} +)"; + +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeX = 256; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 6ed7229604b..949ad3c0de6 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -153,42 +153,90 @@ void expect_dual_q4_topology(WebGPUGraph& graph) { 1); } -TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { - const WebGPUShaderInfo& sigmoid = get_webgpu_shader_info("sigmoid"); - EXPECT_EQ(sigmoid.name, "sigmoid"); - EXPECT_NE(sigmoid.source, nullptr); - EXPECT_GT(sigmoid.workgroup_size_x, 0u); - EXPECT_THROW( - get_webgpu_shader_info("not_a_registered_shader"), std::runtime_error); -} +struct Conv1dRouteCase { + const char* name; + std::vector input_dims; + std::vector weight_dims; + std::vector output_dims; + int64_t stride; + int64_t padding; + int64_t dilation; + int64_t groups; + const char* expected_kernel; +}; + +void build_conv1d_route_graph( + WebGPUGraph& graph, + const Conv1dRouteCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](const std::vector& dims, int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + auto add_int_list = [&](int64_t value) { + const int id = static_cast(values.size()); + const std::vector items = {value}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::IntList, + vk::CreateIntListDirect(fbb, &items).Union())); + return id; + }; + + const int input = add_tensor(test_case.input_dims, 0); + const int weight = add_tensor(test_case.weight_dims, 1); + const int bias = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + const int stride = add_int_list(test_case.stride); + const int padding = add_int_list(test_case.padding); + const int dilation = add_int_list(test_case.dilation); + const int transposed = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Bool, vk::CreateBool(fbb, false).Union())); + const int output_padding = add_int_list(0); + const int groups = add_int(test_case.groups); + const int output = add_tensor(test_case.output_dims, 2); + const std::vector args = { + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.convolution.default", &args)); + const std::vector input_ids = { + static_cast(input), static_cast(weight)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); -TEST(WebGPUQ4RouteSignal, PreservesStaticAndRecordsBothDynamicSignals) { - WebGPUGraph static_graph; - static_graph.set_device(g_device); - build_q4_route_graph(static_graph, Q4RouteSignal::Static); - const auto static_dispatches = q4_dispatches(static_graph); - ASSERT_EQ(static_dispatches.size(), 1); - EXPECT_EQ(static_graph.num_dispatches(), 1); - EXPECT_NE(static_dispatches[0]->pipeline, nullptr); - EXPECT_NE(static_dispatches[0]->bind_group, nullptr); - EXPECT_FALSE(static_graph.has_dynamic_shapes()); - EXPECT_FALSE(static_graph.config().record_q4gsw_decode_route); - - WebGPUGraph legacy_graph; - legacy_graph.set_device(g_device); - build_q4_route_graph(legacy_graph, Q4RouteSignal::LegacyGraphMarker); - EXPECT_TRUE(legacy_graph.has_dynamic_shapes()); - EXPECT_FALSE(legacy_graph.config().record_q4gsw_decode_route); - EXPECT_EQ(legacy_graph.num_dispatches(), 3); - expect_dual_q4_topology(legacy_graph); - - WebGPUGraph explicit_graph; - explicit_graph.set_device(g_device); - build_q4_route_graph(explicit_graph, Q4RouteSignal::ExplicitOption); - EXPECT_FALSE(explicit_graph.has_dynamic_shapes()); - EXPECT_TRUE(explicit_graph.config().record_q4gsw_decode_route); - EXPECT_EQ(explicit_graph.num_dispatches(), 2); - expect_dual_q4_topology(explicit_graph); + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { @@ -899,6 +947,58 @@ TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { } } +TEST(WebGPUToCopyValidation, RejectsBoolAndByteIntegerConversions) { + ASSERT_TRUE(webgpu_operator_registry().has_op("aten._to_copy.default")); + namespace vk = vkgraph; + struct TestCase { + const char* name; + vk::VkDataType input_dtype; + vk::VkDataType output_dtype; + }; + const TestCase cases[] = { + {"bool_to_int8", vk::VkDataType::BOOL, vk::VkDataType::INT8}, + {"bool_to_uint8", vk::VkDataType::BOOL, vk::VkDataType::UINT8}, + {"int8_to_bool", vk::VkDataType::INT8, vk::VkDataType::BOOL}, + {"uint8_to_bool", vk::VkDataType::UINT8, vk::VkDataType::BOOL}, + }; + for (const TestCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + ::flatbuffers::FlatBufferBuilder fbb; + const std::vector dims = {4}; + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.input_dtype, &dims, -1, 0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.output_dtype, &dims, -1, 1) + .Union())); + const std::vector args = {0, 1}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten._to_copy.default", &args)); + const std::vector input_ids = {0}; + const std::vector output_ids = {1}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + FAIL() << test_case.name << " unexpectedly built"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + EXPECT_EQ(graph.memory_stats().num_dispatches, 0); + } +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 7b3c1d6b0a7..cf16cd7e87c 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -18,6 +18,8 @@ // F dyn_rms_chain (rms(rms(x))) at 3 S -> golden (resize CASCADE, DD-4) // G rms+residual H rms*x I dyn_linear J sdpa_dyn K emb_dyn L rope_dyn // M dyn_sigmoid N dyn_select (select_copy(0,-1), dynamic S) +// O ONE dyn_conv1d graph reused across live input lengths +// P ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. // // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else @@ -33,6 +35,7 @@ #include #include +#include #include #include #include @@ -127,6 +130,74 @@ void check_s(Module& m, const std::string& prefix, int s) { << " golden.size=" << golden.size() << ")"; } +constexpr int kConv1dInChannels = 3; +constexpr int kConv1dOutChannels = 4; +constexpr int kConv1dKernel = 3; +constexpr int kConv1dStride = 2; +constexpr int kConv1dPadding = 1; +constexpr int kConv1dDilation = 2; + +void check_conv1d(Module& module, int length) { + const std::string prefix = g_dir + "/dyn_conv1d.S" + std::to_string(length); + auto input = read_bin(prefix + ".input.bin"); + auto golden = read_bin(prefix + ".golden.bin"); + ASSERT_EQ(input.size(), static_cast(kConv1dInChannels * length)); + ASSERT_FALSE(golden.empty()); + auto tensor = + make_tensor_ptr({1, kConv1dInChannels, length}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "conv1d length=" << length << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const int output_length = (length + 2 * kConv1dPadding - + kConv1dDilation * (kConv1dKernel - 1) - 1) / + kConv1dStride + + 1; + ASSERT_EQ(output.dim(), 3); + ASSERT_EQ(output.size(0), 1); + ASSERT_EQ(output.size(1), kConv1dOutChannels); + ASSERT_EQ(output.size(2), output_length); + const size_t numel = static_cast(kConv1dOutChannels * output_length); + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + const float error = max_err(got, golden); + EXPECT_LT(error, 1e-3f) << "conv1d length=" << length << " max_err=" << error; +} + +constexpr int kGeluOld1dDispatchCap = 4 * 64 * 65535; +constexpr int kGelu2dDispatchBoundary = kGeluOld1dDispatchCap + 1; +constexpr int kGeluPatternSize = 257; + +void check_gelu_2d(Module& module, int elements) { + std::array golden = {}; + std::vector input(static_cast(elements)); + for (int i = 0; i < kGeluPatternSize; i++) { + const float value = -4.0f + 8.0f * i / (kGeluPatternSize - 1); + golden[i] = 0.5f * value * (1.0f + std::erf(value * 0.7071067811865476f)); + } + for (int i = 0; i < elements; i++) { + input[i] = -4.0f + 8.0f * (i % kGeluPatternSize) / (kGeluPatternSize - 1); + } + auto tensor = make_tensor_ptr({elements}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "gelu elements=" << elements << " forward failed"; + const auto& output = result.get()[0].toTensor(); + ASSERT_EQ(output.dim(), 1); + ASSERT_EQ(output.size(0), elements); + ASSERT_EQ(output.numel(), elements); + const float* data = output.const_data_ptr(); + float error = 0.0f; + for (int i = 0; i < elements; i++) { + error = std::fmax(error, std::fabs(data[i] - golden[i % kGeluPatternSize])); + } + EXPECT_LT(error, 1e-4f) << "gelu elements=" << elements + << " max_err=" << error; +} + // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; @@ -893,6 +964,43 @@ TEST(DynamicShape, RmsNormReusedGraph) { } } +TEST(DynamicShape, Conv1dReusedGraph) { + Module module(g_dir + "/dyn_conv1d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_conv1d.pte"; + for (int length : {16, 9, 5, 16}) { + check_conv1d(module, length); + } +} + +TEST(DynamicShape, GeluCrosses2dDispatchBoundary) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + Module module(g_dir + "/dyn_gelu_2d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_gelu_2d.pte"; + for (int elements : + {kGelu2dDispatchBoundary, + kGeluOld1dDispatchCap, + kGelu2dDispatchBoundary}) { + check_gelu_2d(module, elements); + } +} + +TEST(DynamicShape, ExpandCopyRejectsDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy.pte"; + ASSERT_TRUE(std::ifstream(path).good()) << "missing dyn_expand_copy.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + +TEST(DynamicShape, ExpandCopyRejectsInferredDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy_inferred.pte"; + ASSERT_TRUE(std::ifstream(path).good()) + << "missing dyn_expand_copy_inferred.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + // C2: grow-only reuse — one loaded rms graph run smallest -> largest, so the // FIRST resize grows the dispatch (every other reuse test starts at MAXS and // only shrinks; this catches a hook with a shrink-only short-circuit). diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp index edc0f315294..a839d224b16 100644 --- a/backends/webgpu/test/native/test_webgpu_utils.cpp +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -14,8 +14,16 @@ #include +#include + using namespace executorch::backends::webgpu; +TEST(WebGPUUtils, DivUpDoesNotOverflowAtUint32Max) { + constexpr uint32_t kMax = std::numeric_limits::max(); + EXPECT_EQ(utils::div_up(kMax, 4u), 1073741824u); + EXPECT_EQ(utils::div_up(kMax, kMax), 1u); +} + TEST(WebGPUUtils, DispatchGridStaysOneDimUnderCeiling) { utils::DispatchGrid g = utils::compute_dispatch_grid_from_limits(1000u, 256u, 65535u, "test"); diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index bbd6e13dc9d..6a019ab8e9f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -53,8 +53,16 @@ CompareModule, ) from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule -from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule +from executorch.backends.webgpu.test.ops.test_conv1d_pw import ( + Conv1dModule, + Conv1dPwModule, + GENERAL_CONFIGS as _CONV1D_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ConvWithClampModule +from executorch.backends.webgpu.test.ops.test_expand_copy import ( + CONFIGS as _EXPAND_COPY_CONFIGS, + ExpandCopyModule, +) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule @@ -153,8 +161,13 @@ ) from executorch.backends.webgpu.test.ops.test_to_copy import ( + bool_tail_input, + compare_to_copy_input_a, + compare_to_copy_input_b, + CompareToCopyBoolToFloatModule, to_copy_float_input, to_copy_int_input, + ToCopyBoolToFloatModule, ToCopyFloatToIntToFloatModule, ToCopyIntToFloatModule, ) @@ -191,6 +204,39 @@ def _add_factory(variant: str = "regular") -> torch.nn.Module: }[variant]() +@register_op_test("to_copy_bool_to_float") +def _to_copy_bool_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=CompareToCopyBoolToFloatModule, + cases=[ + Case( + inputs=( + InputSpec((n,), gen=compare_to_copy_input_a), + InputSpec((n,), gen=compare_to_copy_input_b), + ), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + +@register_op_test("to_copy_bool_input_to_float") +def _to_copy_bool_input_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=ToCopyBoolToFloatModule, + cases=[ + Case( + inputs=(InputSpec((n,), gen=bool_tail_input),), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + @register_op_test("add") def _add_suite() -> WebGPUTestSuite: # Same-shape numeric coverage only: broadcast adds stay export-smoke in @@ -295,10 +341,7 @@ def _minimum_suite() -> WebGPUTestSuite: def _compare_suite(op: str) -> WebGPUTestSuite: - # Elementwise fp32 comparison -> bool (byte-exact golden). The two inputs use - # DIFFERENT discrete-range seeds so a!=b (real lt/gt mix) while colliding - # often (eq/le/ge ties); all shapes have numel % 4 == 0 (bool output packs 4 - # bytes/word). Same-shape only (flat kernel; broadcast=smoke). + # Distinct inputs and tail shapes cover byte-exact packed BOOL output. def case(name, shape): return Case( name=name, @@ -310,7 +353,14 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda: CompareModule(op), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case("tail_1", (1,)), + case("tail_5", (5,)), + case("tail_67", (67,)), + case("2d", (4, 8)), + case("3d", (2, 3, 8)), + case("sq", (16, 16)), + ], golden_dtype="bool", ) @@ -612,6 +662,54 @@ def case(name, C, L, kernel, stride, padding, dilation, bias): case("k3s2p1", 4, 8, 3, 2, 1, 1, True), case("dil2", 3, 10, 3, 1, 2, 2, True), case("k5_nobias", 5, 7, 5, 1, 0, 1, False), + case("single_channel_route", 1, 8, 3, 1, 1, 1, True), + ], + atol=1e-3, + rtol=1e-3, + ) + + +@register_op_test("conv1d") +def _conv1d_suite() -> WebGPUTestSuite: + # General NCL conv1d; neighboring suites cover the retained fast paths. + def case(name, cfg): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + return Case( + name=name, + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + inputs=((n, ic, length),), + ) + + dynamic_cfg = _CONV1D_CONFIGS["voxtral_stride1"] + n, ic, oc, length, kernel, stride, padding, dilation, bias = dynamic_cfg + dynamic_length = torch.export.Dim("manifest_conv1d_length", min=7, max=length) + return WebGPUTestSuite( + module_factory=Conv1dModule, + cases=[ + *[case(name, cfg) for name, cfg in _CONV1D_CONFIGS.items()], + Case( + name="dynamic_length_10_to_7", + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + export_inputs=((n, ic, length),), + inputs=((n, ic, 7),), + dynamic_shapes=({2: dynamic_length},), + ), ], atol=1e-3, rtol=1e-3, @@ -988,12 +1086,36 @@ def _cat_suite() -> WebGPUTestSuite: N as _GELU_N, ) +_GELU_2D_DISPATCH_BOUNDARY = 4 * 64 * 65535 + 1 +_EXPAND_COPY_2D_DISPATCH_BOUNDARY = 64 * 65535 + 1 + def _gelu_full_range(_shape) -> torch.Tensor: # Reuse the deterministic linspace(-6, 6) spanning negatives/zero/positives. return _gelu_det_input() +@register_op_test("expand_copy") +def _expand_copy_suite() -> WebGPUTestSuite: + cases = [ + Case(name=name, construct={"shape": out_shape}, inputs=(in_shape,)) + for name, (in_shape, out_shape) in _EXPAND_COPY_CONFIGS.items() + ] + cases.append( + Case( + name="dispatch_2d_boundary", + construct={"shape": (_EXPAND_COPY_2D_DISPATCH_BOUNDARY,)}, + inputs=((1,),), + heavy=True, + ) + ) + return WebGPUTestSuite( + module_factory=ExpandCopyModule, + cases=cases, + golden_dtype="float32", + ) + + @register_op_test("gelu") def _gelu_suite() -> WebGPUTestSuite: # erf ("none") is the Florence-2/BART + PyTorch default; tanh is the approx. @@ -1015,6 +1137,12 @@ def _gelu_suite() -> WebGPUTestSuite: construct={"approximate": "none"}, inputs=(InputSpec(shape=(_GELU_N,), gen=_gelu_full_range),), ), + Case( + name="erf_dispatch_2d_boundary", + construct={"approximate": "none"}, + inputs=(InputSpec(shape=(_GELU_2D_DISPATCH_BOUNDARY,), gen="ramp"),), + heavy=True, + ), ], atol=1e-4, rtol=1e-3, diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 72f819f94ce..33b7940a9be 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -39,6 +39,8 @@ def _materialize(spec) -> torch.Tensor: shape, gen = spec, "randn" if callable(gen): _t = gen(shape) + if _t.dtype == torch.bool: + return _t return ( _t.to(torch.int32) if not _t.is_floating_point() else _t.to(torch.float32) ) @@ -53,13 +55,17 @@ def _materialize(spec) -> torch.Tensor: def export_case(suite: WebGPUTestSuite, case) -> tuple[torch.nn.Module, tuple, object]: - """Build the module + forward inputs and export to an ExecuTorch program.""" + """Build the module and export it, returning the live runtime inputs.""" module = suite.module_factory(**case.construct) # Seed so an unseeded-randn input is reproducible across generations (the golden uses # the SAME tensor, so this only affects which bytes a case sees, never pass/fail). torch.manual_seed(0) inputs = tuple(_materialize(s) for s in case.inputs) - ep = torch.export.export(module, inputs) + export_inputs = inputs + if case.export_inputs is not None: + torch.manual_seed(0) + export_inputs = tuple(_materialize(s) for s in case.export_inputs) + ep = torch.export.export(module, export_inputs, dynamic_shapes=case.dynamic_shapes) prog = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] ).to_executorch() @@ -169,7 +175,10 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d input_entries: list[dict] = [] for i, t in enumerate(inputs): rel = f"{case_id}.in{i}.bin" - if t.dtype == torch.int32: + if t.dtype == torch.bool: + _write_int8(t.to(torch.int8), os.path.join(out_dir, rel)) + in_dtype = "bool" + elif t.dtype == torch.int32: t.detach().cpu().numpy().astype(" +#include #include #include #include @@ -61,7 +62,15 @@ class OpCase : public ::testing::Test { const size_t n = numel(in.shape); std::vector sizes( in.shape.begin(), in.shape.end()); - if (in.dtype == "int32") { + if (in.dtype == "bool") { + auto data = load_int8_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + std::vector raw(data.begin(), data.end()); + tensors.push_back(make_tensor_ptr( + std::move(sizes), + std::move(raw), + executorch::aten::ScalarType::Bool)); + } else if (in.dtype == "int32") { auto data = load_int32_bin(in.path, n); ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); @@ -96,10 +105,11 @@ class OpCase : public ::testing::Test { auto golden = load_int8_bin(e_.golden.path, gn); ASSERT_FALSE(golden.empty()) << "missing/short golden: " << e_.golden.path; - const bool* out_p = out_tensor.const_data_ptr(); + ASSERT_EQ(out_tensor.scalar_type(), executorch::aten::ScalarType::Bool); + const uint8_t* out_p = out_tensor.const_data_ptr(); int mism = -1; for (size_t i = 0; i < gn; i++) { - if (static_cast(out_p[i]) != golden[i]) { + if (out_p[i] != static_cast(golden[i])) { mism = static_cast(i); break; } diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index 65f765812be..ec4125a2818 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -61,6 +61,27 @@ def test_generate_case_writes_artifacts(tmp_path): assert entry["golden"]["output_index"] == 0 +def test_export_case_separates_upper_bound_from_runtime_inputs(monkeypatch): + suite = op_test_registry["conv1d"] + case = next(c for c in suite.cases if c.name == "dynamic_length_10_to_7") + export_shapes = [] + exported_dynamic_shapes = [] + real_export = torch.export.export + + def capture_export(module, inputs, **kwargs): + export_shapes.append(tuple(inputs[0].shape)) + exported_dynamic_shapes.append(kwargs.get("dynamic_shapes")) + return real_export(module, inputs, **kwargs) + + monkeypatch.setattr(torch.export, "export", capture_export) + _module, runtime_inputs, prog = g.export_case(suite, case) + + assert export_shapes == [(1, 4, 10)] + assert exported_dynamic_shapes == [case.dynamic_shapes] + assert tuple(runtime_inputs[0].shape) == (1, 4, 7) + assert g._has_vulkan_delegate(prog) + + def test_generate_manifest(tmp_path): g.generate(str(tmp_path), ops=["add"]) manifest = tmp_path / "manifest.json" diff --git a/backends/webgpu/test/op_tests/test_suite.py b/backends/webgpu/test/op_tests/test_suite.py index f2714125c84..17542cd2e55 100644 --- a/backends/webgpu/test/op_tests/test_suite.py +++ b/backends/webgpu/test/op_tests/test_suite.py @@ -58,6 +58,9 @@ class Case: required: bool = True heavy: bool = False golden_fn: Callable | None = None + # Optional upper-bound export inputs; `inputs` stay live manifest tensors. + export_inputs: tuple[Input, ...] | None = None + dynamic_shapes: object | None = None def __post_init__(self) -> None: # Mirror kQ4gswConfigs: every heavy config is required=False (export-gated, never FAILs on absence). diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index b302b3c120d..d9f666622ee 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -19,6 +19,8 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule +from executorch.backends.webgpu.test.ops.test_gelu import GeluModule from executorch.exir import to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes @@ -181,6 +183,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.select(0, -1) +class DynamicExpandCopyModule(torch.nn.Module): + """Dynamic expand_copy is rejected until its TensorMeta can be resized.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, x.shape[1])).clone() + + +class DynamicExpandCopyInferredModule(torch.nn.Module): + """Dynamic expand_copy whose -1 target hides symbolic provenance.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, -1)).clone() + + def _ramp(shape) -> torch.Tensor: n = 1 for d in shape: @@ -252,9 +268,75 @@ def _write_goldens(model, prefix: str, out_dir: str, s_values) -> None: print(f" golden {prefix} S={s}") +def export_dynamic_conv1d_cases(out_dir: str) -> None: + """Write one dynamic Conv1d program and live-length runtime fixtures.""" + os.makedirs(out_dir, exist_ok=True) + max_length = 16 + lengths = (max_length, 9, 5) + model = Conv1dModule( + in_channels=3, + out_channels=4, + kernel_size=3, + stride=2, + padding=1, + dilation=2, + bias=True, + ).eval() + length_dim = torch.export.Dim("conv1d_length", min=5, max=max_length) + _export( + model, + (_ramp((1, 3, max_length)),), + {"x": {2: length_dim}}, + os.path.join(out_dir, "dyn_conv1d.pte"), + ) + for length in lengths: + x = _ramp((1, 3, length)) + with torch.no_grad(): + golden = model(x) + prefix = os.path.join(out_dir, f"dyn_conv1d.S{length}") + x.detach().numpy().astype(" None: + """Write a dynamic GELU fixture crossing the old 1D dispatch cap.""" + os.makedirs(out_dir, exist_ok=True) + max_elements = 4 * 64 * 65535 + 1 + model = GeluModule("none").eval() + elements_dim = torch.export.Dim("gelu_elements", min=1024, max=max_elements) + _export( + model, + (torch.empty((max_elements,), dtype=torch.float32),), + {"x": {0: elements_dim}}, + os.path.join(out_dir, "dyn_gelu_2d.pte"), + ) + + +def export_dynamic_expand_copy_rejection_case(out_dir: str) -> None: + """Write a dynamic expand_copy graph that the runtime must reject at load.""" + model = DynamicExpandCopyModule().eval() + elements_dim = torch.export.Dim("expand_elements", min=1, max=8) + _export( + model, + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy.pte"), + ) + _export( + DynamicExpandCopyInferredModule().eval(), + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy_inferred.pte"), + ) + + def export_dynamic_shape_cases(out_dir: str) -> None: """Write the dynamic + static .pte's and per-S goldens for the native test.""" os.makedirs(out_dir, exist_ok=True) + export_dynamic_conv1d_cases(out_dir) + export_dynamic_expand_copy_rejection_case(out_dir) + if os.environ.get("WEBGPU_TEST_HEAVY"): + export_dynamic_gelu_boundary_cases(out_dir) s_dim = torch.export.Dim("s", min=1, max=MAXS) # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). @@ -1536,6 +1618,13 @@ def test_export_dynamic_rms(self) -> None: self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.pte"))) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.S1.golden.bin"))) expected = [ + "dyn_conv1d.pte", + "dyn_conv1d.S16.input.bin", + "dyn_conv1d.S16.golden.bin", + "dyn_conv1d.S9.input.bin", + "dyn_conv1d.S9.golden.bin", + "dyn_conv1d.S5.input.bin", + "dyn_conv1d.S5.golden.bin", "dyn_linear_bk64.pte", "dyn_linear_bk64.S512.input.bin", "dyn_linear_bk64.S512.golden.bin", diff --git a/backends/webgpu/test/ops/test_conv1d_pw.py b/backends/webgpu/test/ops/test_conv1d_pw.py index a0de8988f81..07036fe4a42 100644 --- a/backends/webgpu/test/ops/test_conv1d_pw.py +++ b/backends/webgpu/test/ops/test_conv1d_pw.py @@ -27,6 +27,16 @@ "batch2": (2, 3, 4, 5, True), } +# name -> N, C_in, C_out, L, K, stride, padding, dilation, bias +GENERAL_CONFIGS = { + "voxtral_stride1": (1, 4, 6, 10, 3, 1, 0, 1, True), + "voxtral_stride2": (1, 6, 5, 10, 3, 2, 0, 1, True), + "no_bias": (1, 3, 2, 9, 3, 1, 0, 1, False), + "padded": (1, 3, 4, 9, 3, 1, 1, 1, True), + "dilated": (1, 2, 3, 11, 3, 1, 2, 2, True), + "batch2": (2, 3, 4, 8, 3, 2, 1, 1, True), +} + class Conv1dPwModule(torch.nn.Module): def __init__(self, in_channels, out_channels, bias) -> None: @@ -42,6 +52,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dModule(torch.nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias, + ) -> None: + super().__init__() + g = torch.Generator().manual_seed(0) + self.conv = torch.nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + with torch.no_grad(): + self.conv.weight.normal_(generator=g) + if bias: + self.conv.bias.normal_(generator=g) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + def _det_input(shape): g = torch.Generator().manual_seed(1) return torch.randn(*shape, generator=g, dtype=torch.float32) @@ -54,6 +95,14 @@ def _lower(cfg): return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) +def _lower_general(cfg, dynamic_shapes=None): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + module = Conv1dModule(ic, oc, kernel, stride, padding, dilation, bias).eval() + inputs = (_det_input((n, ic, length)),) + ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + def _delegated(et) -> bool: return any( d.id == "VulkanBackend" @@ -63,9 +112,17 @@ def _delegated(et) -> bool: def _op_delegated(edge, op_substr: str) -> bool: - # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + # Require the op in a delegate, not merely absent from the host graph. + from executorch.exir.lowered_backend_module import get_lowered_submodules + gm = edge.exported_program().graph_module - return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + if any(op_substr in str(getattr(n, "target", "")) for n in gm.graph.nodes): + return False + return any( + op_substr in str(getattr(dn, "target", "")) + for _, lowered, _ in get_lowered_submodules(gm) + for dn in lowered.original_module.graph_module.graph.nodes + ) class Conv1dPwTest(unittest.TestCase): @@ -82,3 +139,24 @@ def test_export_delegates(self) -> None: _op_delegated(edge, "convolution"), f"conv1d not delegated (fell back to CPU) for {name}", ) + + +class Conv1dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in GENERAL_CONFIGS.items(): + with self.subTest(name=name): + edge = _lower_general(cfg) + self.assertTrue( + _delegated(edge.to_executorch()), + f"Expected a VulkanBackend delegate (conv1d {name})", + ) + self.assertTrue( + _op_delegated(edge, "convolution"), + f"conv1d not delegated (fell back to CPU) for {name}", + ) + + def test_dynamic_length_export_delegates(self) -> None: + length = torch.export.Dim("conv1d_length", min=5, max=16) + edge = _lower_general(GENERAL_CONFIGS["padded"], dynamic_shapes=({2: length},)) + self.assertTrue(_delegated(edge.to_executorch())) + self.assertTrue(_op_delegated(edge, "convolution")) diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py index 1fa2375f248..54b400ea9ef 100644 --- a/backends/webgpu/test/ops/test_to_copy.py +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -46,6 +46,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.float32, copy=True) +class ToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class ToCopyInt8ToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class CompareToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a > b).to(torch.float32) + + def to_copy_int_input(shape: tuple[int, ...]) -> torch.Tensor: n = math.prod(shape) return (torch.arange(n, dtype=torch.int32) - n // 2).reshape(shape) @@ -61,14 +76,32 @@ def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor: return pattern.repeat(repeats)[:n].reshape(shape) -def _lower(model: torch.nn.Module, x: torch.Tensor): - ep = torch.export.export(model.eval(), (x,)) +def bool_tail_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([True, False, True, True, False, False, True]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_a(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_b(shape: tuple[int, ...]) -> torch.Tensor: + return torch.zeros(shape, dtype=torch.float32) + + +def _lower(model: torch.nn.Module, *inputs: torch.Tensor): + ep = torch.export.export(model.eval(), inputs) edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) return ep, edge -def _export(model: torch.nn.Module, x: torch.Tensor): - _, edge = _lower(model, x) +def _export(model: torch.nn.Module, *inputs: torch.Tensor): + _, edge = _lower(model, *inputs) return edge.to_executorch() @@ -142,3 +175,22 @@ def test_float_passthrough_delegates(self) -> None: self.assertTrue( _delegated(et), "Expected a VulkanBackend delegate (to_copy float->float)" ) + + def test_bool_to_float_delegates(self) -> None: + x = bool_tail_input((5,)) + ep, edge = _lower(ToCopyBoolToFloatModule(), x) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_compare_bool_to_float_delegates(self) -> None: + a = compare_to_copy_input_a((5,)) + b = compare_to_copy_input_b((5,)) + ep, edge = _lower(CompareToCopyBoolToFloatModule(), a, b) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_int8_to_float_does_not_delegate(self) -> None: + x = torch.tensor([-2, 0, 3], dtype=torch.int8) + self.assertFalse(_delegated(_export(ToCopyInt8ToFloatModule(), x))) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 574c3869864..9990297b9f0 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -77,6 +77,16 @@ def _function_source(text: str, name: str) -> str: class WgslCodegenTest(unittest.TestCase): + def test_compare_word_count_does_not_overflow_u32(self) -> None: + source = (g.BACKEND_ROOT / "runtime/ops/compare/compare.wgsl").read_text() + expression = "(params.num_elements - 1u) / 4u + 1u" + self.assertIn(expression, source) + for num_elements in (1, 4, 5, (1 << 32) - 3, (1 << 32) - 2, (1 << 32) - 1): + self.assertEqual( + (num_elements - 1) // 4 + 1, + (num_elements + 3) // 4, + ) + def test_registry_entries_match_concrete_headers(self) -> None: entries = g.registry_entries() names = [entry.name for entry in entries] @@ -210,14 +220,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), 134) + self.assertEqual(len(outputs), 136) self.assertEqual( digest.hexdigest(), - "e502196846f0f8100f468e5d9f8f9c006b67e08df54e1e2e667daa2fc50d8844", + "0512f8d258952e446ffaedcb653b6a3a720eccf8a6b5327d95fd454a912214a3", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "492535b396833ad6ebfed29b093057e38f40b1182b5c7d6b3eb2f3577cab024e", + "28aaa7a8d3e916df43e407120e91d487d0d51cbc5ca93c56bd822d25d109890e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: