[Fix][WebGPU] Validate and bound symbolic stack allocations

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";

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;

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)

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 main(A: T.Buffer((8,), "float32"), B: T.Buffer((8,), "float32")):
     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)
+
+
+@pytest.mark.parametrize("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()