From 023bd4eaab1c8ff31249cc46d4a0b4a9f498ad15 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Tue, 11 Aug 2026 07:25:10 -0400 Subject: [PATCH] [FIX][WEBGPU] Validate and bound symbolic stack allocations --- src/backend/webgpu/codegen/codegen_webgpu.cc | 67 +++++- src/backend/webgpu/codegen/codegen_webgpu.h | 4 + src/backend/webgpu/codegen/target_kind.cc | 1 + .../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 fe67febc4577..a82f9f96f140 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -31,6 +31,7 @@ #include #include +#include #include #include #include @@ -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()) { @@ -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(); - 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(); + 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(); + // 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(dim_size), + std::numeric_limits::max() / constant_size) + << "WebGPU allocation element count is too large to represent"; + constant_size *= static_cast(dim_size); + } + + size_t element_stride = GetWgslArrayElementStride(op->buffer->dtype); + TVM_FFI_ICHECK_LE(constant_size, std::numeric_limits::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::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::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("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(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 " << 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 a79144260f16..079df4505a00 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.h +++ b/src/backend/webgpu/codegen/codegen_webgpu.h @@ -29,6 +29,7 @@ #include +#include #include #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 ced5d90fe4fe..abb40fdc29f7 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("max_num_threads", refl::DefaultValue(256)) + .add_attr_option("max_shared_memory_per_block", refl::DefaultValue(32768)) .add_attr_option("supports_subgroups", refl::DefaultValue(false)) .add_attr_option("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 96bfc6b90d85..e6dcf1c2d691 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 main(A: T.Buffer((8,), "float32"), B: T.Buffer((8,), "float32")): assert "var 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;", 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\s+\w+\s*:\s*array;", 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()