diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 6b07baae963b..7c844eb0c0ba 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 @@ -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(); - 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) { + // 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()) { + 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(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..150d1c40cec3 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/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."""