This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c3acd8c110 [Fix][Metal] Bound symbolic stack allocations (#20236)
c3acd8c110 is described below
commit c3acd8c110bdfdb8ecccf9e8c1cf0fe16e018159
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Thu Sep 3 12:06:17 2026 -0700
[Fix][Metal] Bound symbolic stack allocations (#20236)
Use analyzer-proven finite upper bounds when emitting fixed-size Metal
arrays. Reject unbounded, nonpositive, and overflowing extents while
allowing shapes such as `min(dynamic_extent, constant_limit)`.
---
src/backend/metal/codegen/codegen_metal.cc | 27 +++-
tests/python/codegen/test_target_codegen_metal.py | 149 ++++++++++++++++++++++
2 files changed, 171 insertions(+), 5 deletions(-)
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 6b07baae96..7c844eb0c0 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/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;
diff --git a/tests/python/codegen/test_target_codegen_metal.py
b/tests/python/codegen/test_target_codegen_metal.py
index f0d9998b4c..150d1c40ce 100644
--- a/tests/python/codegen/test_target_codegen_metal.py
+++ b/tests/python/codegen/test_target_codegen_metal.py
@@ -384,6 +384,155 @@ def test_codegen_simdgroup_buffer_data():
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)
+
+
[email protected]("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."""
