[Fix][Metal][WebGPU] Preserve variable bindings when resolving allocation bounds

src/backend/metal/codegen/codegen_metal.cc

@@ -27,7 +27,7 @@
 #include <tvm/ffi/container/map.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/runtime/logging.h>
-#include <tvm/sym/analyzer.h>
+#include <tvm/tirx/analysis.h>
 #include <tvm/tirx/transform.h>
 
 #include <algorithm>
@@ -67,6 +67,7 @@ Var GetSimdgroupBufferVar(const Expr& data) {
 
 void CodeGenMetal::InitFuncState(const PrimFunc& f) {
   CodeGenC::InitFuncState(f);
+  analyzer_ = sym::Analyzer();
   // analyze the data;
   for (Var arg : f->params) {
     if (arg->ty.as<PointerTypeNode>()) {
@@ -331,6 +332,11 @@ void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os)
 }
 
 void CodeGenMetal::Dispatch_(const BindNode* op) {
+  // Stateful reads cannot be substituted after the underlying state changes.
+  if (auto prim_value = op->value.as<PrimExpr>();
+      prim_value && SideEffect(prim_value.value()) <= CallEffectKind::kPure) {
+    analyzer_->Bind(op->var, prim_value.value());
+  }
   const auto* pointer_type = op->var->ty.as<PointerTypeNode>();
   if (pointer_type == nullptr || pointer_type->storage_scope.empty()) {
     return CodeGenC::Dispatch_(op);
@@ -362,11 +368,10 @@ void CodeGenMetal::Dispatch_(const AllocBufferNode* op) {
   this->PrintIndent();
   // Compute a compile-time upper bound on the number of buffer elements.
   size_t constant_size = 1;
-  sym::Analyzer analyzer;
   for (const auto& dim : op->buffer->shape) {
     const auto* dim_imm = dim.as<IntImmNode>();
     int64_t dim_size =
-        dim_imm ? static_cast<int64_t>(dim_imm->value) : analyzer->const_int_bound(dim)->max_value;
+        dim_imm ? static_cast<int64_t>(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 != sym::ConstIntBound::kPosInf)

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 <tvm/sym/analyzer.h>
 #include <tvm/target/codegen.h>
 
 #include <string>
@@ -63,6 +64,7 @@ class CodeGenMetal final : public CodeGenC {
   using CodeGenC::PrintType;
 
  private:
+  sym::Analyzer analyzer_;
   std::unordered_map<const VarNode*, std::string> simdgroup_dtype_;
   int thread_index_bits_{32};
   int thread_work_dim_{0};

src/backend/webgpu/codegen/codegen_webgpu.cc

@@ -27,7 +27,7 @@
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/prim/builtin.h>
 #include <tvm/support/io.h>
-#include <tvm/sym/analyzer.h>
+#include <tvm/tirx/analysis.h>
 #include <tvm/tirx/builtin.h>
 #include <tvm/tirx/transform.h>
 
@@ -171,6 +171,7 @@ std::string CodeGenWebGPU::Finish() {
 
 void CodeGenWebGPU::InitFuncState(const PrimFunc& f) {
   CodeGenC::InitFuncState(f);
+  analyzer_ = sym::Analyzer();
   workgroup_memory_bytes_ = 0;
   // analyze the data;
   for (Var arg : f->params) {
@@ -648,6 +649,11 @@ void CodeGenWebGPU::Dispatch_(const TensorLoadNode* op, std::ostream& os) {  //
 }
 
 void CodeGenWebGPU::Dispatch_(const BindNode* op) {
+  // Stateful reads cannot be substituted after the underlying state changes.
+  if (auto prim_value = op->value.as<PrimExpr>();
+      prim_value && SideEffect(prim_value.value()) <= CallEffectKind::kPure) {
+    analyzer_->Bind(op->var, prim_value.value());
+  }
   // use ssa form.
   if (print_ssa_form_) {
     std::string value = PrintExpr(op->value);
@@ -726,11 +732,10 @@ void CodeGenWebGPU::Dispatch_(const AllocBufferNode* op) {
   TVM_FFI_ICHECK(op->buffer.defined());
   std::string vid = AllocVarID(op->buffer.get());
   size_t constant_size = 1;
-  sym::Analyzer analyzer;
   for (const auto& dim : op->buffer->shape) {
     const auto* dim_imm = dim.as<IntImmNode>();
     int64_t dim_size =
-        dim_imm ? static_cast<int64_t>(dim_imm->value) : analyzer->const_int_bound(dim)->max_value;
+        dim_imm ? static_cast<int64_t>(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.

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 <tvm/sym/analyzer.h>
 #include <tvm/target/codegen.h>
 
 #include <cstddef>
@@ -86,6 +87,8 @@ class CodeGenWebGPU final : public CodeGenC {
   void Dispatch_(const ContinueNode* op) final;
 
  private:
+  sym::Analyzer analyzer_;
+
   /*!
    * \brief Enforce value to be U32.
    */

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,83 @@ 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)
+
+
+@pytest.mark.parametrize("scope", ["local", "shared"])
+@pytest.mark.parametrize("bounded", [True, False])
+def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded):
+    @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("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "tirx.is_global_func": True,
+                }
+            )
+            state = T.alloc_buffer((1,), "int32", scope="local")
+            state[0] = 0
+            snapshot: T.let[T.int32] = state[0]
+            state[0] = 32
+            difference: T.let[T.int32] = state[0] - snapshot
+            # The snapshot is immutable, but the buffer it read has changed.
+            # Substituting the load would incorrectly reduce this extent to 1.
+            scratch = T.alloc_buffer(
+                (T.min(T.max(difference, 1), 32 if bounded else 2147483647),),
+                "float32",
+                scope=scope,
+            )
+            scratch[31] = 1.0
+
+    if bounded:
+        source = _build_metal(Module).inspect_source()
+        storage = "threadgroup" if scope == "shared" else "thread"
+        assert f"{storage} float scratch[32]" in source
+        assert "scratch[31] =" 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:

tests/python/codegen/test_target_codegen_webgpu.py

@@ -67,6 +67,116 @@ def main(n: T.int32):
     assert re.search(r"\bvar\s+\w+\s*:\s*array<f32,\s*128>;", 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<workgroup>" if scope == "shared" else r"\bvar"
+        assert len(re.findall(declaration + r"\s+\w+\s*:\s*array<f32,\s*128>;", 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("scope", ["local", "shared"])
+@pytest.mark.parametrize("bounded", [True, False])
+def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded):
+    @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,
+                }
+            )
+            state = T.alloc_buffer((1,), "int32", scope="local")
+            state[0] = 0
+            snapshot: T.let[T.int32] = state[0]
+            state[0] = 32
+            difference: T.let[T.int32] = state[0] - snapshot
+            # The snapshot is immutable, but the buffer it read has changed.
+            # Substituting the load would incorrectly reduce this extent to 1.
+            scratch = T.alloc_buffer(
+                (T.min(T.max(difference, 1), 32 if bounded else 2147483647),),
+                "float32",
+                scope=scope,
+            )
+            scratch[31] = 1.0
+
+    if bounded:
+        source = _build_webgpu(Module).inspect_source()
+        declaration = r"var<workgroup>" if scope == "shared" else r"\bvar"
+        assert re.search(declaration + r"\s+scratch\s*:\s*array<f32,\s*32>;", source)
+        assert "scratch[31" in source
+    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<workgroup>\s+\w+\s*:\s*array<f32,\s*128>;", 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: