diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 7c844eb0c0ba..4165ffbddee8 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -66,6 +66,7 @@ Var GetSimdgroupBufferVar(const Expr& data) { void CodeGenMetal::InitFuncState(const PrimFunc& f) { CodeGenC::InitFuncState(f); + analyzer_ = arith::Analyzer(); // analyze the data; for (Var arg : f->params) { if (arg->ty.as()) { @@ -330,6 +331,9 @@ void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os) } void CodeGenMetal::VisitStmt_(const BindNode* op) { + if (auto prim_value = op->value.as()) { + analyzer_->Bind(op->var, prim_value.value()); + } const auto* pointer_type = op->var->ty.as(); if (pointer_type == nullptr || pointer_type->storage_scope.empty()) { return CodeGenC::VisitStmt_(op); @@ -361,10 +365,9 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) { this->PrintIndent(); // 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 auto* dim_imm = dim.as(); - int64_t dim_size = dim_imm ? dim_imm->value : analyzer->const_int_bound(dim)->max_value; + 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) diff --git a/src/backend/metal/codegen/codegen_metal.h b/src/backend/metal/codegen/codegen_metal.h index 1f7097bf514a..662ed903e090 100644 --- a/src/backend/metal/codegen/codegen_metal.h +++ b/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 #include #include @@ -63,6 +64,7 @@ class CodeGenMetal final : public CodeGenC { using CodeGenC::PrintType; private: + arith::Analyzer analyzer_; std::unordered_map simdgroup_dtype_; int thread_index_bits_{32}; int thread_work_dim_{0}; diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc b/src/backend/webgpu/codegen/codegen_webgpu.cc index e0262876688a..4c78ed109e28 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -168,6 +168,7 @@ std::string CodeGenWebGPU::Finish() { void CodeGenWebGPU::InitFuncState(const PrimFunc& f) { CodeGenC::InitFuncState(f); + analyzer_ = arith::Analyzer(); workgroup_memory_bytes_ = 0; // analyze the data; for (Var arg : f->params) { @@ -643,6 +644,9 @@ void CodeGenWebGPU::VisitExpr_(const TensorLoadNode* op, std::ostream& os) { // } void CodeGenWebGPU::VisitStmt_(const BindNode* op) { + if (auto prim_value = op->value.as()) { + analyzer_->Bind(op->var, prim_value.value()); + } // use ssa form. if (print_ssa_form_) { std::string value = PrintExpr(op->value); @@ -721,10 +725,9 @@ void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); std::string vid = AllocVarID(op->buffer.get()); size_t constant_size = 1; - arith::Analyzer analyzer; for (const auto& dim : op->buffer->shape) { const auto* dim_imm = dim.as(); - int64_t dim_size = dim_imm ? dim_imm->value : analyzer->const_int_bound(dim)->max_value; + 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. diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h b/src/backend/webgpu/codegen/codegen_webgpu.h index 078fbc6f4a29..9b1cc5ffbd00 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.h +++ b/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 #include #include @@ -86,6 +87,8 @@ class CodeGenWebGPU final : public CodeGenC { void VisitStmt_(const ContinueNode* op) final; private: + arith::Analyzer analyzer_; + /*! * \brief Enforce value to be U32. */ diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py index 150d1c40cec3..4e54558c6d1b 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/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,40 @@ 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) + + def test_bounded_uint64_symbolic_stack_allocation(): @I.ir_module class Module: diff --git a/tests/python/codegen/test_target_codegen_webgpu.py b/tests/python/codegen/test_target_codegen_webgpu.py index e6dcf1c2d691..7873ec51d120 100644 --- a/tests/python/codegen/test_target_codegen_webgpu.py +++ b/tests/python/codegen/test_target_codegen_webgpu.py @@ -67,6 +67,74 @@ def main(n: T.int32): assert re.search(r"\bvar\s+\w+\s*:\s*array;", 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" if scope == "shared" else r"\bvar" + assert len(re.findall(declaration + r"\s+\w+\s*:\s*array;", 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("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\s+\w+\s*:\s*array;", 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: