From bf2c0c8de4299844f505a85c22d8c67d598c9e06 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Thu, 13 Aug 2026 00:18:01 -0400 Subject: [PATCH 1/3] [FIX][METAL] Bound symbolic stack allocations 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 | 20 ++++++-- .../codegen/test_target_codegen_metal.py | 51 +++++++++++++++++++ 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 6b07baae963b..a9569851d8b2 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 #include #include #include @@ -31,6 +32,7 @@ #include #include +#include #include #include #include @@ -359,12 +361,22 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { this->PrintIndent(); // Compute constant_size from buffer shape 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; + 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(); + TVM_FFI_ICHECK(dtype_max && 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(dim_size), + std::numeric_limits::max() / constant_size) + << "Metal allocation element count is too large to represent"; + constant_size *= static_cast(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 f0d9998b4ce2..2d69b1a68f7f 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/tests/python/codegen/test_target_codegen_metal.py @@ -384,6 +384,57 @@ 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_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_codegen_pointer_byte_offsets_preserve_storage_scope(): """Pointer byte offsets should preserve the source Metal address space.""" From 981096bfa18020b7819b39b8c85fe6d0bab4d39a Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 30 Aug 2026 12:34:38 -0400 Subject: [PATCH 2/3] [TEST][METAL] Cover symbolic allocation bound failures Exercise rejection of nonpositive allocation upper bounds and size_t element-count overflow in Metal code generation. Clarify that symbolic allocation sizing uses program-derived compile-time upper bounds rather than integer dtype limits. --- src/backend/metal/codegen/codegen_metal.cc | 3 +- .../codegen/test_target_codegen_metal.py | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index a9569851d8b2..d85d68ab1d38 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -359,7 +359,7 @@ 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) { @@ -367,6 +367,7 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { 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) << "Metal allocation extent requires a finite compile-time upper bound, but got " << dim; } diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py index 2d69b1a68f7f..47ae18cff3dd 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/tests/python/codegen/test_target_codegen_metal.py @@ -435,6 +435,58 @@ def main(n: T.int32): _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.""" From c13fab94fec57043f23c091810e29627083d30f0 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Tue, 1 Sep 2026 19:27:07 -0400 Subject: [PATCH 3/3] [FIX][METAL] Handle bounded uint64 stack allocations Treat the analyzer positive-infinity sentinel independently from an integer dtype maximum. This allows finite uint64 extents whose dtype maximum is represented as LargeUIntImm while continuing to reject genuinely unbounded extents. Add bounded and unbounded uint64 codegen regressions alongside the existing narrower-integer coverage. --- src/backend/metal/codegen/codegen_metal.cc | 8 +++- .../codegen/test_target_codegen_metal.py | 46 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index d85d68ab1d38..7c844eb0c0ba 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -366,10 +366,14 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { 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) + 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()) { + 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; diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py index 47ae18cff3dd..150d1c40cec3 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/tests/python/codegen/test_target_codegen_metal.py @@ -410,6 +410,27 @@ def main(n: T.int32): 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: @@ -435,6 +456,31 @@ def main(n: T.int32): _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