diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index fc22e24896ab..5325cefe6d78 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include @@ -67,6 +67,7 @@ Var GetSimdgroupBufferVar(const Expr& data) { void CodeGenMetal::InitFuncState(const PrimFunc& f) { CodeGenC::InitFuncState(f); + analyzer_ = sym::Analyzer(); // analyze the data; for (Var arg : f->params) { if (arg->ty.as()) { @@ -331,6 +332,11 @@ void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os) } void CodeGenMetal::Dispatch_(const BindNode* op) { + // Stateful reads cannot be substituted after the underlying state changes. + if (auto prim_value = op->value.as(); + prim_value && SideEffect(prim_value.value()) <= CallEffectKind::kPure) { + 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::Dispatch_(op); @@ -362,11 +368,10 @@ void CodeGenMetal::Dispatch_(const AllocBufferNode* op) { this->PrintIndent(); // Compute a compile-time upper bound on the number of buffer elements. size_t constant_size = 1; - sym::Analyzer analyzer; for (const auto& dim : op->buffer->shape) { const auto* dim_imm = dim.as(); int64_t dim_size = - dim_imm ? static_cast(dim_imm->value) : analyzer->const_int_bound(dim)->max_value; + dim_imm ? static_cast(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 != sym::ConstIntBound::kPosInf) diff --git a/src/backend/metal/codegen/codegen_metal.h b/src/backend/metal/codegen/codegen_metal.h index 7fad4eb398b4..96b89bb15142 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: + sym::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 ee9de61624ec..07ad151d32a9 100644 --- a/src/backend/webgpu/codegen/codegen_webgpu.cc +++ b/src/backend/webgpu/codegen/codegen_webgpu.cc @@ -27,7 +27,7 @@ #include #include #include -#include +#include #include #include @@ -171,6 +171,7 @@ std::string CodeGenWebGPU::Finish() { void CodeGenWebGPU::InitFuncState(const PrimFunc& f) { CodeGenC::InitFuncState(f); + analyzer_ = sym::Analyzer(); workgroup_memory_bytes_ = 0; // analyze the data; for (Var arg : f->params) { @@ -648,6 +649,11 @@ void CodeGenWebGPU::Dispatch_(const TensorLoadNode* op, std::ostream& os) { // } void CodeGenWebGPU::Dispatch_(const BindNode* op) { + // Stateful reads cannot be substituted after the underlying state changes. + if (auto prim_value = op->value.as(); + prim_value && SideEffect(prim_value.value()) <= CallEffectKind::kPure) { + analyzer_->Bind(op->var, prim_value.value()); + } // use ssa form. if (print_ssa_form_) { std::string value = PrintExpr(op->value); @@ -726,11 +732,10 @@ void CodeGenWebGPU::Dispatch_(const AllocBufferNode* op) { TVM_FFI_ICHECK(op->buffer.defined()); std::string vid = AllocVarID(op->buffer.get()); size_t constant_size = 1; - sym::Analyzer analyzer; for (const auto& dim : op->buffer->shape) { const auto* dim_imm = dim.as(); int64_t dim_size = - dim_imm ? static_cast(dim_imm->value) : analyzer->const_int_bound(dim)->max_value; + dim_imm ? static_cast(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 e7e4bc1d06d0..d9f10bc16a55 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 Dispatch_(const ContinueNode* op) final; private: + sym::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..8f6c1a00c387 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,83 @@ 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) + + +@pytest.mark.parametrize("scope", ["local", "shared"]) +@pytest.mark.parametrize("bounded", [True, False]) +def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded): + @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, + } + ) + state = T.alloc_buffer((1,), "int32", scope="local") + state[0] = 0 + snapshot: T.let[T.int32] = state[0] + state[0] = 32 + difference: T.let[T.int32] = state[0] - snapshot + # The snapshot is immutable, but the buffer it read has changed. + # Substituting the load would incorrectly reduce this extent to 1. + scratch = T.alloc_buffer( + (T.min(T.max(difference, 1), 32 if bounded else 2147483647),), + "float32", + scope=scope, + ) + scratch[31] = 1.0 + + if bounded: + source = _build_metal(Module).inspect_source() + storage = "threadgroup" if scope == "shared" else "thread" + assert f"{storage} float scratch[32]" in source + assert "scratch[31] =" 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..42c52144eda2 100644 --- a/tests/python/codegen/test_target_codegen_webgpu.py +++ b/tests/python/codegen/test_target_codegen_webgpu.py @@ -67,6 +67,116 @@ 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("scope", ["local", "shared"]) +@pytest.mark.parametrize("bounded", [True, False]) +def test_allocation_bound_does_not_substitute_buffer_load(scope, bounded): + @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("webgpu"), + "tirx.is_global_func": True, + } + ) + state = T.alloc_buffer((1,), "int32", scope="local") + state[0] = 0 + snapshot: T.let[T.int32] = state[0] + state[0] = 32 + difference: T.let[T.int32] = state[0] - snapshot + # The snapshot is immutable, but the buffer it read has changed. + # Substituting the load would incorrectly reduce this extent to 1. + scratch = T.alloc_buffer( + (T.min(T.max(difference, 1), 32 if bounded else 2147483647),), + "float32", + scope=scope, + ) + scratch[31] = 1.0 + + if bounded: + source = _build_webgpu(Module).inspect_source() + declaration = r"var" if scope == "shared" else r"\bvar" + assert re.search(declaration + r"\s+scratch\s*:\s*array;", source) + assert "scratch[31" in source + 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: