This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e03e9784ef [Fix][S-TIR] Preserve general reduction predicates (#20242)
e03e9784ef is described below
commit e03e9784efd9ab11e754502779131fd65541b830
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Sep 1 19:05:06 2026 -0700
[Fix][S-TIR] Preserve general reduction predicates (#20242)
Rebuild decomposed reduction predicates recursively across conjunctions.
Preserve clauses independent of discarded loops and drop clauses that
reference discarded loops. This supports comparisons beyond `<` and
detects loop variables on either operand.
Tail predicates generated by `split` and `rfactor`, such as `outer *
factor + inner < extent`, are omitted from initialization when they
reference a discarded reduction loop, while remaining on the update
block.
---
src/s_tir/schedule/primitive/reduction.cc | 34 +++++-----
.../s_tir/schedule/test_tir_schedule_reduction.py | 78 ++++++++++++++++++++++
2 files changed, 93 insertions(+), 19 deletions(-)
diff --git a/src/s_tir/schedule/primitive/reduction.cc
b/src/s_tir/schedule/primitive/reduction.cc
index 99ff247f5d..beb05afbd5 100644
--- a/src/s_tir/schedule/primitive/reduction.cc
+++ b/src/s_tir/schedule/primitive/reduction.cc
@@ -159,23 +159,17 @@ class LoopHeightError : public ScheduleError {
SBlock block_;
};
-PrimExpr RemakePredicate(PrimExpr pred, const std::unordered_set<const
VarNode*>& discarded_loops) {
+PrimExpr RewriteInitPredicate(PrimExpr pred,
+ const std::unordered_set<const VarNode*>&
discarded_loops) {
if (is_one(pred)) return IntImm::Bool(true);
- PrimExpr new_pred = IntImm::Bool(true);
- auto f = [&](const VarNode* var) { return discarded_loops.count(var); };
- arith::PVar<PrimExpr> lhs, rhs, rest;
- for (;;) {
- if ((rest && (lhs < rhs)).Match(pred)) {
- if (!UsesVar(lhs.Eval(), f)) new_pred = new_pred && (lhs.Eval() <
rhs.Eval());
- pred = rest.Eval();
- } else if ((lhs < rhs).Match(pred)) {
- if (!UsesVar(lhs.Eval(), f)) new_pred = new_pred && (lhs.Eval() <
rhs.Eval());
- break;
- } else {
- TVM_FFI_ICHECK(false) << "Unexpected predicate for reduction block";
- }
+ if (const auto* and_node = pred.as<AndNode>()) {
+ return RewriteInitPredicate(and_node->a, discarded_loops) &&
+ RewriteInitPredicate(and_node->b, discarded_loops);
}
- return new_pred;
+ auto uses_discarded_loop = [&discarded_loops](const VarNode* var) {
+ return discarded_loops.count(var);
+ };
+ return UsesVar(pred, uses_discarded_loop) ? IntImm::Bool(true) : pred;
}
StmtSRef DecomposeReduction(ScheduleState self, const StmtSRef& block_sref,
@@ -259,15 +253,17 @@ StmtSRef DecomposeReduction(ScheduleState self, const
StmtSRef& block_sref,
discarded = false;
break;
}
- if (discarded) discarded_loops.insert(loop_var);
+ if (discarded) {
+ discarded_loops.insert(loop_var);
+ }
// Only scan loops not higher than the given loop
if (loops[i].same_as(loop_sref)) {
break;
}
}
- // Step 4. After scanning loops, make a new predicate in the init block
realize
- // We discard predicate that is related to discarded loops
- init_realize->predicate = RemakePredicate(realize->predicate,
discarded_loops);
+ // Step 4. Derive the predicate for the init block realize. Omit
conjunction clauses that
+ // depend on discarded loops.
+ init_realize->predicate = RewriteInitPredicate(realize->predicate,
discarded_loops);
// Step 5. Create new loops above init block
std::unordered_map<Var, Var> loop_var_map;
Stmt body = SBlockRealize(init_realize);
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
b/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
index 1643b13df0..598f68b1f9 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_reduction.py
@@ -391,5 +391,83 @@ def test_decompose_reduction_with_thread_binding():
tvm.ir.assert_structural_equal(After, Expected)
+def test_decompose_reduction_preserves_general_spatial_predicates():
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((8, 8), "float32"), B: T.Buffer((8,), "float32")):
+ for i, k in T.grid(10, 10):
+ with T.sblock("B"):
+ T.where(1 <= i and i < 9 and 1 <= k and k < 9)
+ vi = T.axis.spatial(8, i - 1)
+ vk = T.axis.reduce(8, k - 1)
+ with T.init():
+ B[vi] = T.float32(0)
+ B[vi] += A[vi, vk]
+
+ @I.ir_module(s_tir=True)
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((8, 8), "float32"), B: T.Buffer((8,), "float32")):
+ for i_init in range(10):
+ with T.sblock("B_init"):
+ T.where(1 <= i_init and i_init < 9)
+ vi = T.axis.spatial(8, i_init - 1)
+ B[vi] = T.float32(0)
+ for i, k in T.grid(10, 10):
+ with T.sblock("B_update"):
+ T.where(1 <= i and i < 9 and 1 <= k and k < 9)
+ vi = T.axis.spatial(8, i - 1)
+ vk = T.axis.reduce(8, k - 1)
+ B[vi] += A[vi, vk]
+
+ sch = tvm.s_tir.Schedule(Before)
+ i, _ = sch.get_loops("B")
+ sch.decompose_reduction("B", i)
+ tvm.ir.assert_structural_equal(sch.mod, Expected)
+
+
+def test_decompose_reduction_drops_mixed_rfactor_bound():
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((20,), "float32"), B: T.Buffer((), "float32")):
+ for k in range(20):
+ with T.sblock("B"):
+ vk = T.axis.reduce(20, k)
+ with T.init():
+ B[()] = T.float32(0)
+ B[()] += A[vk]
+
+ @I.ir_module(s_tir=True)
+ class Expected:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((20,), "float32"), B: T.Buffer((), "float32")):
+ B_rf = T.sblock_alloc_buffer((16,), elem_offset=T.int64(0))
+ for k_1_init in range(16):
+ with T.sblock("B_rf_init"):
+ vk_1 = T.axis.spatial(16, k_1_init)
+ B_rf[vk_1] = T.float32(0)
+ for k_0, k_1 in T.grid(2, 16):
+ with T.sblock("B_rf_update"):
+ vk_1, vk_0 = T.axis.remap("SR", [k_1, k_0])
+ T.where(k_0 * 16 + k_1 < 20)
+ B_rf[vk_1] += A[vk_0 * 16 + vk_1]
+ for k_1 in range(16):
+ with T.sblock("B"):
+ vk_1 = T.axis.reduce(16, k_1)
+ with T.init():
+ B[()] = T.float32(0)
+ B[()] += B_rf[vk_1]
+
+ sch = tvm.s_tir.Schedule(Before)
+ (k,) = sch.get_loops("B")
+ _, k_1 = sch.split(k, factors=[None, 16])
+ rf = sch.rfactor(k_1, 0)
+ k_0, _ = sch.get_loops(rf)
+ sch.decompose_reduction(rf, k_0)
+ tvm.ir.assert_structural_equal(sch.mod, Expected)
+
+
if __name__ == "__main__":
tvm.testing.main()
