From 730fcb889bed760ac4bcae779e89805154b5d658 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 13 Sep 2026 15:26:33 -0400 Subject: [PATCH 1/2] [FIX][Metal] Retain bound expressions for symbolic allocations Track primitive Bind expressions in a per-function analyzer so allocation bounds remain provable after common subexpression elimination. Continue rejecting extents whose only upper bound is the integer type limit. Test chained bounded and unbounded bindings. Update the vector-ramp test to use an explicit T.let declaration required by the current TIRx parser. --- src/backend/metal/codegen/codegen_metal.cc | 7 ++-- src/backend/metal/codegen/codegen_metal.h | 2 ++ .../codegen/test_target_codegen_metal.py | 36 ++++++++++++++++++- 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 7c844eb0c0ba..4165ffbddee8 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -66,6 +66,7 @@ Var GetSimdgroupBufferVar(const Expr& data) { void CodeGenMetal::InitFuncState(const PrimFunc& f) { CodeGenC::InitFuncState(f); + analyzer_ = arith::Analyzer(); // analyze the data; for (Var arg : f->params) { if (arg->ty.as()) { @@ -330,6 +331,9 @@ void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os) } void CodeGenMetal::VisitStmt_(const BindNode* op) { + if (auto prim_value = op->value.as()) { + analyzer_->Bind(op->var, prim_value.value()); + } const auto* pointer_type = op->var->ty.as(); if (pointer_type == nullptr || pointer_type->storage_scope.empty()) { return CodeGenC::VisitStmt_(op); @@ -361,10 +365,9 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { this->PrintIndent(); // Compute a compile-time upper bound on the number of buffer elements. size_t constant_size = 1; - arith::Analyzer analyzer; for (const auto& dim : op->buffer->shape) { const auto* dim_imm = dim.as(); - int64_t dim_size = dim_imm ? dim_imm->value : analyzer->const_int_bound(dim)->max_value; + int64_t dim_size = dim_imm ? dim_imm->value : analyzer_->const_int_bound(dim)->max_value; if (dim_imm == nullptr) { // An integer dtype's intrinsic maximum is not a program-derived allocation bound. TVM_FFI_ICHECK(dim_size != arith::ConstIntBound::kPosInf) diff --git a/src/backend/metal/codegen/codegen_metal.h b/src/backend/metal/codegen/codegen_metal.h index 1f7097bf514a..662ed903e090 100644 --- a/src/backend/metal/codegen/codegen_metal.h +++ b/src/backend/metal/codegen/codegen_metal.h @@ -24,6 +24,7 @@ #ifndef TVM_TARGET_METAL_CODEGEN_METAL_H_ #define TVM_TARGET_METAL_CODEGEN_METAL_H_ +#include #include #include @@ -63,6 +64,7 @@ class CodeGenMetal final : public CodeGenC { using CodeGenC::PrintType; private: + arith::Analyzer analyzer_; std::unordered_map simdgroup_dtype_; int thread_index_bits_{32}; int thread_work_dim_{0}; diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py index 150d1c40cec3..4e54558c6d1b 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/tests/python/codegen/test_target_codegen_metal.py @@ -140,7 +140,7 @@ def main(A: T.Buffer((1, 2), "int32")): for i in T.thread_binding(1, thread="threadIdx.x"): with T.sblock("block"): tx = T.axis.spatial(1, i) - r = T.ramp(tx, 3, 2) + r: T.let = T.ramp(tx, 3, 2) A[0, T.ramp(0, 1, 2)] = r f = tvm.compile(IRModule, target=target) @@ -410,6 +410,40 @@ def main(n: T.int32): assert "thread float scratch[128]" in source +@pytest.mark.parametrize("bounded", [True, False]) +def test_bound_symbolic_stack_allocation(bounded): + limit = 64 if bounded else 2147483647 + + @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("metal"), + "tirx.kernel_launch_params": [], + "tirx.is_global_func": True, + } + ) + # Common subexpression elimination can hoist the bounded extent. + extent: T.let[T.int32] = T.min(n, limit) + elements: T.let[T.int32] = extent * 2 + scratch = T.alloc_buffer((elements,), "float32", scope="local") + T.evaluate(scratch.data) + + if bounded: + source = _build_metal(Module).inspect_source() + assert "thread float scratch[128]" in source + else: + with pytest.raises( + tvm.error.InternalError, + match="Metal allocation extent requires a finite compile-time upper bound", + ): + _build_metal(Module) + + def test_bounded_uint64_symbolic_stack_allocation(): @I.ir_module class Module: From 0d7f7fc5bbe098036d42512f9e95a4e2efc78bfd Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 13 Sep 2026 17:48:45 -0400 Subject: [PATCH 2/2] [FIX][WebGPU] Retain bound expressions for symbolic allocations Keep a per-function arithmetic analyzer and record primitive variable bindings before resolving allocation extents. This preserves finite bounds hoisted into temporary variables, including chains of bindings, for local and workgroup allocations. Cover bounded and unbounded aliases in both storage scopes and verify that resolved workgroup allocation sizes still obey the target memory limit. --- src/backend/webgpu/codegen/codegen_webgpu.cc | 7 +- src/backend/webgpu/codegen/codegen_webgpu.h | 3 + .../codegen/test_target_codegen_webgpu.py | 68 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc index e0262876688a..4c78ed109e28 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -168,6 +168,7 @@ std::string CodeGenWebGPU::Finish() { void CodeGenWebGPU::InitFuncState(const PrimFunc& f) { CodeGenC::InitFuncState(f); + analyzer_ = arith::Analyzer(); workgroup_memory_bytes_ = 0; // analyze the data; for (Var arg : f->params) { @@ -643,6 +644,9 @@ void CodeGenWebGPU::VisitExpr_(const TensorLoadNode* op, std::ostream& os) { // } void CodeGenWebGPU::VisitStmt_(const BindNode* op) { + if (auto prim_value = op->value.as()) { + analyzer_->Bind(op->var, prim_value.value()); + } // use ssa form. if (print_ssa_form_) { std::string value = PrintExpr(op->value); @@ -721,10 +725,9 @@ 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 auto* dim_imm = dim.as(); - int64_t dim_size = dim_imm ? dim_imm->value : analyzer->const_int_bound(dim)->max_value; + 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. diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h b/src/backend/webgpu/codegen/codegen_webgpu.h index 078fbc6f4a29..9b1cc5ffbd00 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.h +++ b/src/backend/webgpu/codegen/codegen_webgpu.h @@ -27,6 +27,7 @@ #ifndef TVM_TARGET_WEBGPU_CODEGEN_WEBGPU_H_ #define TVM_TARGET_WEBGPU_CODEGEN_WEBGPU_H_ +#include #include #include @@ -86,6 +87,8 @@ class CodeGenWebGPU final : public CodeGenC { void VisitStmt_(const ContinueNode* op) final; private: + arith::Analyzer analyzer_; + /*! * \brief Enforce value to be U32. */ diff --git a/tests/python/codegen/test_target_codegen_webgpu.py b/tests/python/codegen/test_target_codegen_webgpu.py index e6dcf1c2d691..7873ec51d120 100644 --- a/tests/python/codegen/test_target_codegen_webgpu.py +++ b/tests/python/codegen/test_target_codegen_webgpu.py @@ -67,6 +67,74 @@ def main(n: T.int32): assert re.search(r"\bvar\s+\w+\s*:\s*array;", source) +@pytest.mark.parametrize("scope", ["local", "shared"]) +@pytest.mark.parametrize("bounded", [True, False]) +def test_bound_symbolic_allocation(scope, bounded): + limit = 64 if bounded else 2147483647 + + @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, + } + ) + # Common subexpression elimination can hoist the bounded extent. + extent: T.let[T.int32] = T.min(n, limit) + first = T.alloc_buffer((extent * 2,), "float32", scope=scope) + elements: T.let[T.int32] = extent * 2 + second = T.alloc_buffer((elements,), "float32", scope=scope) + first[0] = 1.0 + second[0] = first[0] + + if bounded: + source = _build_webgpu(Module).inspect_source() + declaration = r"var" if scope == "shared" else r"\bvar" + assert len(re.findall(declaration + r"\s+\w+\s*:\s*array;", source)) == 2 + else: + with pytest.raises( + tvm.error.InternalError, + match="WebGPU allocation extent requires a finite compile-time upper bound", + ): + _build_webgpu(Module) + + +@pytest.mark.parametrize("target_limit", [512, 496]) +def test_bound_symbolic_workgroup_allocation_respects_target_limit(target_limit): + @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, + } + ) + extent: T.let[T.int32] = T.min(n, 64) + elements: T.let[T.int32] = extent * 2 + scratch = T.alloc_buffer((elements,), "float32", scope="shared") + scratch[0] = 1.0 + + target = {"kind": "webgpu", "max_shared_memory_per_block": target_limit} + if target_limit == 512: + source = _build_webgpu(Module, target).inspect_source() + assert re.search(r"var\s+\w+\s*:\s*array;", source) + else: + with pytest.raises( + tvm.error.InternalError, + match=r"WebGPU workgroup allocations use 512 bytes, .* supports only 496 bytes", + ): + _build_webgpu(Module, target) + + def test_unbounded_symbolic_stack_allocation_rejected(): @I.ir_module class Module: