From 50d25ad0a3b4bca564faa630dc579aa994415139 Mon Sep 17 00:00:00 2001 From: yhna940 <100389977+yhna940@users.noreply.github.com> Date: Mon, 26 Oct 2026 07:39:57 +0000 Subject: [PATCH] [Fix][S-TIR] Preserve block predicates across compute_at compute_at rebuilt a moved block from tensor regions and dropped its realize predicate. After rfactor of a non-divisible split, that predicate is the tail guard, so the write-back read past the input. Carry the predicate onto the regenerated realize. reverse_compute_at uses the same path. --- src/s_tir/schedule/primitive/compute_at.cc | 132 +++++++++++++++++- .../schedule/test_tir_schedule_compute_at.py | 112 +++++++++++++++ 2 files changed, 242 insertions(+), 2 deletions(-) diff --git a/src/s_tir/schedule/primitive/compute_at.cc b/src/s_tir/schedule/primitive/compute_at.cc index 0f801bee21c2..33c6a61b4451 100644 --- a/src/s_tir/schedule/primitive/compute_at.cc +++ b/src/s_tir/schedule/primitive/compute_at.cc @@ -18,6 +18,7 @@ */ #include #include +#include #include #include @@ -239,6 +240,114 @@ struct BlockVarDomainInfo { } }; +/*! + * \brief Rewrite a block predicate from the loop nest being removed into the regenerated nest. + * + * `split` of a non-divisible loop stores the tail guard on the block realize (`k0 * factor + k1 < + * extent`). `rfactor` keeps that guard, but the guard is not part of the tensor region + * `compute_at` uses to rebuild loops, so the rebuilt realize would otherwise read out of bounds. + * The result is the original predicate with removed loop vars replaced by the new iter values, or + * nullopt when a removed loop var cannot be eliminated. + */ +ffi::Optional CarryBlockPredicate(const SBlockRealize& old_realize, + const ffi::Array& new_iter_values, + const ffi::Map& old_loop_domains, + const ffi::Map& kept_loop_domains, + sym::AnalyzerObj* analyzer) { + PrimExpr predicate = old_realize->predicate; + if (IsOne(predicate)) { + return IntImm::Bool(true); + } + const SBlock& block = old_realize->block; + int n = static_cast(block->iter_vars.size()); + TVM_FFI_ICHECK_EQ(n, static_cast(old_realize->iter_values.size())); + TVM_FFI_ICHECK_EQ(n, static_cast(new_iter_values.size())); + + auto dropped_loop = [&](const Var& var) -> bool { + return old_loop_domains.find(var) != old_loop_domains.end() && + kept_loop_domains.find(var) == kept_loop_domains.end(); + }; + auto apply_subst = [&](PrimExpr expr, const ffi::Map& subst) -> PrimExpr { + auto f_substitute = [&subst](const Var& var) -> ffi::Expected> { + if (auto repl = subst.Get(var)) return ffi::Any(*std::move(repl)); + return ffi::Unchanged(); + }; + PrimExpr rewritten = + ffi::StructuralMap(std::move(expr), f_substitute) + .as_or_throw(); + return analyzer->Simplify(std::move(rewritten)); + }; + auto free_of_dropped_loops = [&](const PrimExpr& expr) -> bool { + for (const Var& var : UndefinedVars(expr)) { + if (dropped_loop(var)) { + return false; + } + } + return true; + }; + + // Trivial bindings are loop vars. This is the form `rfactor` gives the factored block. + { + ffi::Map trivial; + bool trivial_ok = true; + for (int i = 0; i < n; ++i) { + const auto* loop_var = old_realize->iter_values[i].as(); + if (loop_var == nullptr) { + trivial_ok = false; + break; + } + Var key = ffi::GetRef(loop_var); + const PrimExpr& replacement = new_iter_values[i]; + if (ffi::Optional previous = trivial.Get(key)) { + if (!analyzer->CanProveEqual(previous.value(), replacement)) { + trivial_ok = false; + break; + } + } else { + trivial.Set(key, replacement); + } + } + if (trivial_ok) { + PrimExpr rewritten = apply_subst(predicate, trivial); + if (free_of_dropped_loops(rewritten)) { + return rewritten; + } + } + } + + // Non-trivial affine bindings: express the removed loops in the block iters, then in the new + // iter values. Ignore the predicate here so a tail guard does not make the binding look + // non-bijective. + ffi::Map input_iters; + for (const auto& kv : old_loop_domains) { + input_iters.Set(kv.first.as_or_throw(), kv.second); + } + sym::Analyzer analyzer_ref = ffi::GetRef(analyzer); + sym::IterMapResult detected = sym::DetectIterMap( + old_realize->iter_values, input_iters, IntImm::Bool(true), sym::IterMapLevel::Bijective, + analyzer_ref, /*simplify_trivial_iterators=*/false); + if (detected->indices.size() != static_cast(n)) { + return std::nullopt; + } + ffi::Array block_vars; + block_vars.reserve(n); + for (const IterVar& iter_var : block->iter_vars) { + block_vars.push_back(iter_var->var); + } + ffi::Map loop_to_block_var = + sym::InverseAffineIterMap(detected->indices, block_vars); + PrimExpr in_block_vars = apply_subst(predicate, loop_to_block_var); + ffi::Map block_var_to_new; + for (int i = 0; i < n; ++i) { + block_var_to_new.Set(block->iter_vars[i]->var, new_iter_values[i]); + } + PrimExpr rewritten = apply_subst(std::move(in_block_vars), block_var_to_new); + if (!free_of_dropped_loops(rewritten)) { + return std::nullopt; + } + return rewritten; +} + /*! * \brief A helper to reconstruct the block scope where the given block is moved under the given * loop, and the given block's induced loop nest is regenerated to satisfy the required region. @@ -262,9 +371,14 @@ class ScopeReconstructor : public StmtExprMutator { * \param iter_doms The domain of each block var * \param analyzer The arithmetic analyzer * \param preserve_unit_loops Whether to generate unit loops where the loop extent is 1 + * \param old_realize The realize being moved, whose predicate is carried onto the new nest + * \param old_loop_domains Loop domains surrounding `old_realize` before the move + * \param kept_loop_domains Loop domains that remain outside the regenerated nest */ void MakeNewLoop(int insert_position, std::vector iter_doms, - sym::AnalyzerObj* analyzer, bool preserve_unit_loops) { + sym::AnalyzerObj* analyzer, bool preserve_unit_loops, + const SBlockRealize& old_realize, const ffi::Map& old_loop_domains, + const ffi::Map& kept_loop_domains) { int n_iters = iter_doms.size(); ffi::Array loop_vars; ffi::Array loop_extents; @@ -303,6 +417,14 @@ class ScopeReconstructor : public StmtExprMutator { } } } + // Keep a tail guard (or any other realize predicate) that the region-derived loops do not + // imply. `reverse_compute_at` shares this path. + if (ffi::Optional carried = CarryBlockPredicate( + old_realize, iter_values, old_loop_domains, kept_loop_domains, analyzer)) { + if (!analyzer->CanProve(carried.value())) { + predicate = predicate && carried.value(); + } + } this->new_block_realize_ = SBlockRealize(std::move(iter_values), analyzer->Simplify(predicate), std::move(block_)); Stmt new_subtree = this->new_block_realize_.value(); @@ -781,9 +903,15 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s /*required_regions=*/std::move(required_regions), /*analyzer=*/analyzer); // Step 6. Create the new scope according to the iteration domain + SBlockRealize old_realize = GetSBlockRealize(self, block_sref); + ffi::Map old_loop_domains = + LoopDomainOfSRefTreePath(ffi::GetRef(block_sref->parent)); + ffi::Map kept_loop_domains = LoopDomainOfSRefTreePath(loop_sref); reconstructor->MakeNewLoop(/*insert_position=*/insert_position, /*iter_doms=*/std::move(iter_doms), - /*analyzer=*/analyzer, /*preserve_unit_loops=*/preserve_unit_loops); + /*analyzer=*/analyzer, /*preserve_unit_loops=*/preserve_unit_loops, + /*old_realize=*/old_realize, /*old_loop_domains=*/old_loop_domains, + /*kept_loop_domains=*/kept_loop_domains); SBlock new_scope_root = reconstructor->Mutate(scope_root).ValueOrUnchanged(scope_root).as_or_throw(); diff --git a/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py b/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py index f7c87daad3df..847a6b20e4f8 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_compute_at.py @@ -18,6 +18,7 @@ # ruff: noqa: E501, F401 from __future__ import annotations +import numpy as np import pytest import tvm @@ -1854,5 +1855,116 @@ def expect(): verify_trace_roundtrip(sch=sch, mod=before) +def _llvm_allclose(func, inputs, expected): + if not tvm_tirx.runtime.enabled("llvm"): + return + compiled = tvm.compile(func, target="llvm") + outputs = [] + args = [] + for array in inputs: + args.append(tvm_tirx.runtime.tensor(array)) + for array in expected: + buf = tvm_tirx.runtime.tensor(np.zeros_like(array)) + outputs.append(buf) + args.append(buf) + compiled(*args) + for got, ref in zip(outputs, expected): + tvm_tirx.testing.assert_allclose(got.numpy(), ref) + + +def test_compute_at_rfactor_keeps_tail_predicate(): + """`split` of a non-divisible reduction stores the tail on the block. + + `rfactor` keeps that guard, and `compute_at` must keep it too. Rebuilding + the rfactor block from its write region alone drops the guard and reads + past the input. `debug_mask` stays off: `rfactor` still records affine and + region-cover flags that a fresh analysis rejects for this predicate. + """ + + @Ts.function + def before(X: T.Tensor((4, 77), "float32"), Y: T.Tensor((4,), "float32")): + for r, k in T.grid(4, 77): + with Ts.sblock("Y"): + v_r, v_k = Ts.axis.remap("SR", [r, k]) + with Ts.init(): + Y[v_r] = T.float32(0) + Y[v_r] = Y[v_r] + X[v_r, v_k] + + @Ts.function + def expected(X: T.Tensor((4, 77), "float32"), Y: T.Tensor((4,), "float32")): + Y_rf = Ts.sblock_alloc_tensor((64, 4), "float32") + for r, k_1 in T.grid(4, 64): + for ax0 in range(2): + with Ts.sblock("Y_rf"): + vk_1 = Ts.axis.spatial(64, k_1) + v_r = Ts.axis.spatial(4, r) + vk_0 = Ts.axis.reduce(2, ax0) + Ts.where(ax0 * 64 + k_1 < 77) + with Ts.init(): + Y_rf[vk_1, v_r] = T.float32(0) + Y_rf[vk_1, v_r] = Y_rf[vk_1, v_r] + X[v_r, vk_0 * 64 + vk_1] + with Ts.sblock("Y"): + vk_1 = Ts.axis.reduce(64, k_1) + v_r = Ts.axis.spatial(4, r) + with Ts.init(): + Y[v_r] = T.float32(0) + Y[v_r] = Y[v_r] + Y_rf[vk_1, v_r] + + sch = tvm.s_tir.Schedule(before, debug_mask="none") + block = sch.get_sblock("Y") + _, k = sch.get_loops(block) + sch.split(k, factors=[None, 64]) + rf = sch.rfactor(sch.get_loops(block)[-1], factor_axis=0) + sch.compute_at(rf, sch.get_loops(block)[-1]) + assert_structural_equal_ignore_global_symbol(expected, sch.mod["main"]) + verify_trace_roundtrip(sch=sch, mod=before, debug_mask="none") + + x = (np.arange(4 * 77, dtype="float32") % 7).reshape(4, 77) + _llvm_allclose(sch.mod["main"], [x], [x.sum(axis=1)]) + + +def test_reverse_compute_at_keeps_tail_predicate(): + """`reverse_compute_at` rebuilds loops in the same helper as `compute_at`. + + A consumer whose split does not divide its domain must keep the tail guard + when that guard is not implied by the regenerated loop extent. + """ + + @Ts.function + def before(A: T.Tensor((128,), "float32"), C: T.Tensor((128,), "float32")): + B = Ts.sblock_alloc_tensor((128,), "float32") + for i in range(128): + with Ts.sblock("B"): + vi = Ts.axis.spatial(128, i) + B[vi] = A[vi] + for j0, j1 in T.grid(2, 64): + with Ts.sblock("C"): + vj = Ts.axis.spatial(128, j0 * 64 + j1) + Ts.where(j0 * 64 + j1 < 77) + C[vj] = B[vj] * T.float32(2) + + @Ts.function + def expected(A: T.Tensor((128,), "float32"), C: T.Tensor((128,), "float32")): + B = Ts.sblock_alloc_tensor((128,), "float32") + for i in range(128): + with Ts.sblock("B"): + vi = Ts.axis.spatial(128, i) + B[vi] = A[vi] + with Ts.sblock("C"): + vj = Ts.axis.spatial(128, i) + Ts.where(i < 77) + C[vj] = B[vj] * T.float32(2) + + sch = tvm.s_tir.Schedule(before, debug_mask="none") + sch.reverse_compute_at(sch.get_sblock("C"), sch.get_loops("B")[0]) + assert_structural_equal_ignore_global_symbol(expected, sch.mod["main"]) + verify_trace_roundtrip(sch=sch, mod=before, debug_mask="none") + + x = np.arange(128, dtype="float32") + ref = np.zeros(128, dtype="float32") + ref[:77] = x[:77] * np.float32(2) + _llvm_allclose(sch.mod["main"], [x], [ref]) + + if __name__ == "__main__": tvm_tirx.testing.main()