[Fix][Metal] Bound symbolic stack allocations

src/backend/metal/codegen/codegen_metal.cc

@@ -22,6 +22,7 @@
  */
 #include "codegen_metal.h"
 
+#include <tvm/arith/analyzer.h>
 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/container/array.h>
 #include <tvm/ffi/container/map.h>
@@ -31,6 +32,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <limits>
 #include <sstream>
 #include <string>
 #include <unordered_map>
@@ -357,14 +359,29 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
   std::string vid = AllocVarID(op->buffer.get());
 
   this->PrintIndent();
-  // Compute constant_size from buffer shape
+  // 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 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;
+    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) {
+      // An integer dtype's intrinsic maximum is not a program-derived allocation bound.
+      TVM_FFI_ICHECK(dim_size != arith::ConstIntBound::kPosInf)
+          << "Metal allocation extent requires a finite compile-time upper bound, but got " << dim;
+      if (const auto* dtype_max = max_value(dim.ty()).as<IntImmNode>()) {
+        TVM_FFI_ICHECK_LT(dim_size, dtype_max->value)
+            << "Metal allocation extent requires a finite compile-time upper bound, but got "
+            << dim;
+      }
+    }
+    TVM_FFI_ICHECK_GT(dim_size, 0)
+        << "Metal 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)
+        << "Metal allocation element count is too large to represent";
+    constant_size *= static_cast<size_t>(dim_size);
   }
-  TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now";
 
   auto scope = op->buffer.scope();
   alloc_storage_scope_[op->buffer.get()] = scope;

tests/python/codegen/test_target_codegen_metal.py

@@ -384,6 +384,155 @@ def kernel():
     assert "simdgroup_multiply_accumulate(" in source
 
 
+def _build_metal(mod):
+    build = tvm.get_global_func("target.build.metal")
+    return build(mod, tvm.target.Target("metal"))
+
+
+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("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "tirx.is_global_func": True,
+                }
+            )
+            scratch = T.alloc_buffer((T.min(n, 64), 2), "float32", scope="local")
+            T.evaluate(scratch.data)
+
+    source = _build_metal(Module).inspect_source()
+    assert "thread float scratch[128]" in source
+
+
+def test_bounded_uint64_symbolic_stack_allocation():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(n: T.uint64):
+            T.func_attr(
+                {
+                    "calling_conv": 2,
+                    "global_symbol": "main",
+                    "target": T.target("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "tirx.is_global_func": True,
+                }
+            )
+            scratch = T.alloc_buffer((T.min(n, T.uint64(64)),), "float32", scope="local")
+            T.evaluate(scratch.data)
+
+    source = _build_metal(Module).inspect_source()
+    assert "thread float scratch[64]" in 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("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "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="Metal allocation extent requires a finite compile-time upper bound",
+    ):
+        _build_metal(Module)
+
+
+def test_unbounded_uint64_symbolic_stack_allocation_rejected():
+    @I.ir_module
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(n: T.uint64):
+            T.func_attr(
+                {
+                    "calling_conv": 2,
+                    "global_symbol": "main",
+                    "target": T.target("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "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="Metal allocation extent requires a finite compile-time upper bound",
+    ):
+        _build_metal(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("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "tirx.is_global_func": True,
+                }
+            )
+            scratch = T.alloc_buffer((extent,), "float32", scope="local")
+            T.evaluate(scratch.data)
+
+    with pytest.raises(
+        tvm.error.InternalError,
+        match="Metal allocation extent requires a positive compile-time upper bound",
+    ):
+        _build_metal(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("metal"),
+                    "tirx.kernel_launch_params": [],
+                    "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="Metal allocation element count is too large to represent"
+    ):
+        _build_metal(Module)
+
+
 def test_codegen_pointer_byte_offsets_preserve_storage_scope():
     """Pointer byte offsets should preserve the source Metal address space."""