This is an automated email from the ASF dual-hosted git repository.
tlopex 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 cdef3efe41 [Fix][WebGPU] Validate and bound symbolic stack allocations
(#20132)
cdef3efe41 is described below
commit cdef3efe414fa4fe2aeaac52e8a8a3ac7097b900
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Aug 18 22:52:41 2026 -0700
[Fix][WebGPU] Validate and bound symbolic stack allocations (#20132)
Allow WebGPU code generation for local and workgroup allocations whose
symbolic extents have finite compile-time upper bounds. WGSL requires
array lengths to be compile-time constants, so codegen uses arithmetic
analysis to derive and emit a static upper bound for expressions such as
`min(n, 64)` (unbounded or nonpositive extents are rejected).
This PR:
- Detects element-count, byte-size, and aggregate workgroup-memory
overflow
- Computes allocation sizes using WGSL array element stride
- Applies WebGPU's per-workgroup-variable 16-byte size rounding
- Tracks total workgroup storage used by each generated entry point
- Enforces `max_shared_memory_per_block`, corresponding to WebGPU's
`maxComputeWorkgroupStorageSize`
- Sets the default WebGPU limit to 32 KiB to match TVM's JavaScript
runtime contract while permitting target-specific overrides
---
src/backend/webgpu/codegen/codegen_webgpu.cc | 67 ++++++-
src/backend/webgpu/codegen/codegen_webgpu.h | 4 +
src/backend/webgpu/codegen/target_kind.cc | 1 +
tests/python/codegen/test_target_codegen_webgpu.py | 216 +++++++++++++++++++++
4 files changed, 283 insertions(+), 5 deletions(-)
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index fe67febc45..a82f9f96f1 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -31,6 +31,7 @@
#include <tvm/tirx/transform.h>
#include <algorithm>
+#include <limits>
#include <optional>
#include <string>
#include <unordered_map>
@@ -49,6 +50,26 @@
namespace tvm {
namespace codegen {
+namespace {
+
+size_t GetWgslArrayElementStride(const PrimType& dtype) {
+ if (dtype == PrimType::Bool()) {
+ return 4;
+ }
+
+ int lanes = dtype.lanes();
+ if (dtype.MatchesCode(DLDataTypeCode::kDLInt) && dtype.bits() == 8 && lanes
== 4) {
+ return 4;
+ }
+
+ size_t scalar_bytes = (dtype.bits() + 7) / 8;
+ // WGSL arrays use the alignment-rounded size as their element stride. In
+ // particular, a three-lane vector has the same stride as a four-lane vector.
+ return scalar_bytes * (lanes == 3 ? 4 : lanes);
+}
+
+} // namespace
+
// WebGPU Info
struct WebGPUWorkGroupInfo {
int workgroup_size[3] = {1, 1, 1};
@@ -146,6 +167,7 @@ std::string CodeGenWebGPU::Finish() {
void CodeGenWebGPU::InitFuncState(const PrimFunc& f) {
CodeGenC::InitFuncState(f);
+ workgroup_memory_bytes_ = 0;
// analyze the data;
for (Var arg : f->params) {
if (arg->ty.as<PointerTypeNode>()) {
@@ -696,15 +718,50 @@ void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op)
{
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get());
size_t constant_size = 1;
+ arith::Analyzer analyzer;
for (const auto& dim : op->buffer->shape) {
- const IntImmNode* dim_imm = dim.as<IntImmNode>();
- TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack allocation
for now";
- constant_size *= dim_imm->value;
- }
- TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack
allocation for now";
+ const auto* dim_imm = dim.as<IntImmNode>();
+ int64_t dim_size = dim_imm ? dim_imm->value :
analyzer->const_int_bound(dim)->max_value;
+ if (dim_imm == nullptr) {
+ const auto* dtype_max = max_value(dim.ty()).as<IntImmNode>();
+ // An integer dtype's intrinsic maximum is not a program-derived
allocation bound.
+ TVM_FFI_ICHECK(dtype_max && dim_size < dtype_max->value)
+ << "WebGPU allocation extent requires a finite compile-time upper
bound, but got " << dim;
+ }
+ TVM_FFI_ICHECK_GT(dim_size, 0)
+ << "WebGPU allocation extent requires a positive compile-time upper
bound, but got " << dim;
+ TVM_FFI_ICHECK_LE(static_cast<uint64_t>(dim_size),
+ std::numeric_limits<size_t>::max() / constant_size)
+ << "WebGPU allocation element count is too large to represent";
+ constant_size *= static_cast<size_t>(dim_size);
+ }
+
+ size_t element_stride = GetWgslArrayElementStride(op->buffer->dtype);
+ TVM_FFI_ICHECK_LE(constant_size, std::numeric_limits<size_t>::max() /
element_stride)
+ << "WebGPU allocation byte size is too large to represent";
+ size_t allocation_bytes = constant_size * element_stride;
auto storage_scope = runtime::StorageScope::Create(op->buffer.scope());
if (storage_scope.rank == runtime::StorageRank::kShared) {
+ // WebGPU rounds the size of each workgroup variable up to 16 bytes before
+ // summing the storage used by an entry point.
+ constexpr size_t kWorkgroupVariableAlignment = 16;
+ TVM_FFI_ICHECK_LE(allocation_bytes,
+ std::numeric_limits<size_t>::max() -
(kWorkgroupVariableAlignment - 1))
+ << "WebGPU workgroup allocation size is too large to represent";
+ size_t workgroup_variable_bytes =
+ (allocation_bytes + kWorkgroupVariableAlignment - 1) &
~(kWorkgroupVariableAlignment - 1);
+ TVM_FFI_ICHECK_LE(workgroup_variable_bytes,
+ std::numeric_limits<size_t>::max() -
workgroup_memory_bytes_)
+ << "Total WebGPU workgroup allocation size is too large to represent";
+ workgroup_memory_bytes_ += workgroup_variable_bytes;
+ int64_t limit =
target_->GetAttr<int64_t>("max_shared_memory_per_block").value();
+ TVM_FFI_ICHECK_GT(limit, 0) << "WebGPU max_shared_memory_per_block must be
positive";
+ TVM_FFI_ICHECK_LE(workgroup_memory_bytes_, static_cast<uint64_t>(limit))
+ << "WebGPU workgroup allocations use " << workgroup_memory_bytes_
+ << " bytes, but the target supports only " << limit
+ << " bytes. If the adapter supports this allocation, set "
+ "max_shared_memory_per_block in the WebGPU target configuration.";
this->decl_stream << "var<workgroup> " << vid << " : array<";
PrintType(op->buffer->dtype, this->decl_stream);
this->decl_stream << ", " << constant_size << ">;\n";
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h
b/src/backend/webgpu/codegen/codegen_webgpu.h
index a79144260f..079df4505a 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.h
+++ b/src/backend/webgpu/codegen/codegen_webgpu.h
@@ -29,6 +29,7 @@
#include <tvm/target/codegen.h>
+#include <cstddef>
#include <string>
#include "../../../target/source/codegen_c.h"
@@ -99,6 +100,9 @@ class CodeGenWebGPU final : public CodeGenC {
// whether enable subgroups
bool enable_subgroups_{false};
+ /*! \brief Total bytes declared in the WGSL workgroup address space. */
+ size_t workgroup_memory_bytes_{0};
+
/*! \brief the header stream for function label and enable directive if any,
goes before any other
* declaration */
std::ostringstream header_stream;
diff --git a/src/backend/webgpu/codegen/target_kind.cc
b/src/backend/webgpu/codegen/target_kind.cc
index ced5d90fe4..abb40fdc29 100644
--- a/src/backend/webgpu/codegen/target_kind.cc
+++ b/src/backend/webgpu/codegen/target_kind.cc
@@ -56,6 +56,7 @@ void RegisterTargetKind() {
TVM_REGISTER_TARGET_KIND("webgpu", kDLWebGPU)
.add_attr_option<int64_t>("max_num_threads", refl::DefaultValue(256))
+ .add_attr_option<int64_t>("max_shared_memory_per_block",
refl::DefaultValue(32768))
.add_attr_option<bool>("supports_subgroups", refl::DefaultValue(false))
.add_attr_option<int64_t>("thread_warp_size", refl::DefaultValue(1))
.set_target_canonicalizer(UpdateWebGPUAttrs)
diff --git a/tests/python/codegen/test_target_codegen_webgpu.py
b/tests/python/codegen/test_target_codegen_webgpu.py
index 96bfc6b90d..e6dcf1c2d6 100644
--- a/tests/python/codegen/test_target_codegen_webgpu.py
+++ b/tests/python/codegen/test_target_codegen_webgpu.py
@@ -15,6 +15,10 @@
# specific language governing permissions and limitations
# under the License.
+import re
+
+import pytest
+
import tvm
import tvm.testing
from tvm.script import ir as I
@@ -38,5 +42,217 @@ def test_codegen_buffer_access_modes():
assert "var<storage, read_write> B_ptr" in source
+def _build_webgpu(mod, target="webgpu"):
+ build = tvm.get_global_func("target.build.webgpu")
+ return build(mod, tvm.target.Target(target))
+
+
+def test_bounded_symbolic_stack_allocation():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(n: T.int32):
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer((T.min(n, 64), 2), "float32",
scope="local")
+ T.evaluate(scratch.data)
+
+ source = _build_webgpu(Module).inspect_source()
+ assert re.search(r"\bvar\s+\w+\s*:\s*array<f32,\s*128>;", source)
+
+
+def test_unbounded_symbolic_stack_allocation_rejected():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(n: T.int32):
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer((n,), "float32", scope="local")
+ scratch[0] = 1.0
+ T.evaluate(scratch[0])
+
+ with pytest.raises(
+ tvm.error.InternalError,
+ match="WebGPU allocation extent requires a finite compile-time upper
bound",
+ ):
+ _build_webgpu(Module)
+
+
[email protected]("extent", [0, -1])
+def test_nonpositive_stack_allocation_rejected(extent):
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main():
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer((extent,), "float32", scope="local")
+ T.evaluate(scratch.data)
+
+ with pytest.raises(
+ tvm.error.InternalError,
+ match="WebGPU allocation extent requires a positive compile-time upper
bound",
+ ):
+ _build_webgpu(Module)
+
+
+def test_stack_allocation_element_count_overflow_rejected():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(n: T.int32, m: T.int32, k: T.int32):
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer(
+ (T.min(n, 1 << 30), T.min(m, 1 << 30), T.min(k, 1 << 30)),
+ "uint8",
+ scope="local",
+ )
+ T.evaluate(scratch.data)
+
+ with pytest.raises(
+ tvm.error.InternalError, match="WebGPU allocation element count is too
large to represent"
+ ):
+ _build_webgpu(Module)
+
+
+def test_stack_allocation_byte_size_overflow_rejected():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(n: T.int32, m: T.int32):
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer(
+ (T.min(n, 1 << 30), T.min(m, 1 << 30), 4), "float32",
scope="local"
+ )
+ T.evaluate(scratch.data)
+
+ with pytest.raises(
+ tvm.error.InternalError, match="WebGPU allocation byte size is too
large to represent"
+ ):
+ _build_webgpu(Module)
+
+
+def test_workgroup_allocation_at_target_limit():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main():
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer((8192,), "float32", scope="shared")
+ scratch[0] = 1.0
+
+ source = _build_webgpu(Module).inspect_source()
+ assert re.search(r"var<workgroup>\s+\w+\s*:\s*array<f32,\s*8192>;", source)
+
+
+def test_total_workgroup_allocation_above_target_limit_rejected():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main():
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ first = T.alloc_buffer((4096,), "float32", scope="shared")
+ second = T.alloc_buffer((4097,), "float32", scope="shared")
+ first[0] = 1.0
+ second[0] = 2.0
+
+ with pytest.raises(
+ tvm.error.InternalError,
+ match=r"WebGPU workgroup allocations use 32784 bytes, .* supports only
32768 bytes",
+ ):
+ _build_webgpu(Module)
+
+
+def test_workgroup_allocation_accounts_for_declaration_alignment():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main():
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ first = T.alloc_buffer((1,), "float32", scope="shared")
+ second = T.alloc_buffer((1,), "float32", scope="shared")
+ first[0] = 1.0
+ second[0] = 2.0
+
+ with pytest.raises(
+ tvm.error.InternalError,
+ match=r"WebGPU workgroup allocations use 32 bytes, .* supports only 16
bytes",
+ ):
+ _build_webgpu(Module, {"kind": "webgpu",
"max_shared_memory_per_block": 16})
+
+
+def test_workgroup_allocation_uses_target_limit():
+ @I.ir_module
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main():
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "main",
+ "target": T.target("webgpu"),
+ "tirx.is_global_func": True,
+ }
+ )
+ scratch = T.alloc_buffer((16384,), "float32", scope="shared")
+ scratch[0] = 1.0
+
+ _build_webgpu(Module, {"kind": "webgpu", "max_shared_memory_per_block":
65536})
+
+
if __name__ == "__main__":
tvm.testing.main()
