This is an automated email from the ASF dual-hosted git repository.
akaashrp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 043b57fbe9 [Backend][WebGPU] Fix off-by-one in the grid-packing guard
(#20424)
043b57fbe9 is described below
commit 043b57fbe9b84aba6c0525fc77b624ef8b7bbe2f
Author: Julian Ng-Thow-Hing <[email protected]>
AuthorDate: Sat Oct 3 01:33:34 2026 -0400
[Backend][WebGPU] Fix off-by-one in the grid-packing guard (#20424)
`packGridDimX` is the number of workgroups requested along x, so valid
packed ids are `[0, packGridDimX)`. The emitted guard only returned for
`id > packGridDimX`, letting `id == packGridDimX` run.
The web runtime folds `blockIdx.x` into `blockIdx.z` to stay under the
65535 workgroups-per-dimension limit and pads x up when it is odd, so
that id really is launched: for N = 100277 it dispatches (50139, 1, 2) =
100278 workgroups, and the extra one wrote `B[N]` past the end of an
N-row buffer.
Fix: `>` → `>=`. The codegen test asserts the emitted guard; it fails on
the current codegen and passes with this change.
`tests/python/codegen/test_target_codegen_webgpu.py`: 22 passed.
pre-commit clean.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Fable 5.1 <[email protected]>
---
src/backend/webgpu/codegen/codegen_webgpu.cc | 4 ++--
tests/python/codegen/test_target_codegen_webgpu.py | 16 ++++++++++++++++
2 files changed, 18 insertions(+), 2 deletions(-)
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 5cce127f7b..5c9a2fb399 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -335,8 +335,8 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const
PrimFunc& f, bool skip_re
<< " @builtin(num_workgroups) gridDim : vec3<u32>,\n"
<< " @builtin(local_invocation_id) threadIdx : vec3<u32>\n"
<< ") {\n";
- // skip out of bound grids
- this->stream << " if (blockIdx.z * gridDim.x + blockIdx.x > " // NOLINT(*)
+ // skip out of bound grids; valid packed ids are [0, packGridDimX)
+ this->stream << " if (blockIdx.z * gridDim.x + blockIdx.x >= " // NOLINT(*)
<< val_pod_args << "." << packGridDimX << ") { return; }\n";
// the function scope.
int func_scope = this->BeginScope();
diff --git a/tests/python/codegen/test_target_codegen_webgpu.py
b/tests/python/codegen/test_target_codegen_webgpu.py
index e10af436b0..b3791795b1 100644
--- a/tests/python/codegen/test_target_codegen_webgpu.py
+++ b/tests/python/codegen/test_target_codegen_webgpu.py
@@ -364,5 +364,21 @@ def test_workgroup_allocation_uses_target_limit():
_build_webgpu(Module, {"kind": "webgpu", "max_shared_memory_per_block":
65536})
+def test_grid_pack_guard_rejects_id_equal_to_workgroup_count():
+ """The runtime pads the launch when it folds x into z, so id ==
packGridDimX must return."""
+
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(B: T.Buffer((8,), "int32")):
+ for i in T.thread_binding(8, thread="blockIdx.x"):
+ for j in T.thread_binding(1, thread="threadIdx.x"):
+ B[i] = i
+
+ executable = tvm.compile(Module, target="webgpu")
+ source = executable.mod.imports[0].inspect_source("wgsl")
+ assert re.search(r"blockIdx\.x >= \w+\.packGridDimX\) \{ return; \}",
source)
+
+
if __name__ == "__main__":
tvm.testing.main()
