This is an automated email from the ASF dual-hosted git repository.
syfeng 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 f1ba5ed5f3 [BugFix][TIR] Schedule support reverse-inline with
reduction blocks (#17838)
f1ba5ed5f3 is described below
commit f1ba5ed5f38198947648a00bfeef8f8da86e9f7c
Author: Ruihang Lai <[email protected]>
AuthorDate: Tue Apr 15 04:34:59 2025 -0400
[BugFix][TIR] Schedule support reverse-inline with reduction blocks (#17838)
This PR fixes a bug in reverse-compute-inline of tir Schedule, which
generates incorrect TIR after inlining a transpose block into a
reduction block.
---
src/tir/schedule/primitive/compute_inline.cc | 28 ++++++++++++-
.../test_tir_schedule_compute_inline.py | 49 ++++++++++++++++++++++
2 files changed, 76 insertions(+), 1 deletion(-)
diff --git a/src/tir/schedule/primitive/compute_inline.cc
b/src/tir/schedule/primitive/compute_inline.cc
index df74497b4a..85f3a0f82f 100644
--- a/src/tir/schedule/primitive/compute_inline.cc
+++ b/src/tir/schedule/primitive/compute_inline.cc
@@ -586,6 +586,30 @@ class ReverseComputeInliner : public BaseInliner {
ReverseComputeInliner* self_;
};
+ class RecursionResolver : public StmtExprMutator {
+ public:
+ explicit RecursionResolver(ReverseComputeInliner* self) : self_(self) {}
+
+ private:
+ PrimExpr VisitExpr_(const VarNode* var) final {
+ auto it = self_->idx_sub_.find(var);
+ if (it == self_->idx_sub_.end()) {
+ return GetRef<Var>(var);
+ }
+ return (*it).second;
+ }
+
+ PrimExpr VisitExpr_(const BufferLoadNode* _load) final {
+ BufferLoad load =
Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(_load));
+ return load->buffer.same_as(self_->inlined_buffer_)
+ ? StmtExprMutator::VisitExpr(
+ BufferLoad(self_->inlined_store_->buffer,
self_->inlined_store_->indices))
+ : load;
+ }
+
+ ReverseComputeInliner* self_;
+ };
+
public:
explicit ReverseComputeInliner(const Buffer& inlined_buffer, const
BlockNode* producer_block,
const BlockRealize& consumer_block_realize,
@@ -784,7 +808,9 @@ class ReverseComputeInliner : public BaseInliner {
}
Stmt ReplaceInlinedBuffer(BufferStore producer) {
- producer_rhs_ = producer->value;
+ // "producer->value" may contain the buffer that is inlined in cases of
reduction,
+ // so we need to resolve the recursion first
+ producer_rhs_ = RecursionResolver(this)(producer->value);
return Substituter(this)(GetRef<BufferStore>(inlined_store_));
}
diff --git a/tests/python/tir-schedule/test_tir_schedule_compute_inline.py
b/tests/python/tir-schedule/test_tir_schedule_compute_inline.py
index 2f779612a7..066070f763 100644
--- a/tests/python/tir-schedule/test_tir_schedule_compute_inline.py
+++ b/tests/python/tir-schedule/test_tir_schedule_compute_inline.py
@@ -1529,5 +1529,54 @@ def
test_reverse_compute_inline_slicing_then_cachewrite():
assert_structural_equal_ignore_global_symbol(after, sch.mod["main"])
+def test_inline_with_reduction():
+ @T.prim_func
+ def before(
+ T_softmax_norm: T.Buffer((T.int64(6), T.int64(1), T.int64(1)),
"float32"),
+ T_reshape_2: T.Buffer((T.int64(6), T.int64(1), T.int64(64)),
"float32"),
+ T_transpose: T.Buffer((T.int64(1), T.int64(1), T.int64(6),
T.int64(64)), "float32"),
+ ):
+ T_batch_matmul_NN = T.alloc_buffer((T.int64(6), T.int64(1),
T.int64(64)))
+ for ax0, ax1 in T.grid(T.int64(6), T.int64(64)):
+ with T.block("bmm"):
+ v0, v1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(T_softmax_norm[v0, T.int64(0), T.int64(0)],
T_reshape_2[v0, T.int64(0), v1])
+ T.writes(T_batch_matmul_NN[v0, T.int64(0), v1])
+ with T.init():
+ T_batch_matmul_NN[v0, T.int64(0), v1] = T.float32(0.0)
+ T_batch_matmul_NN[v0, T.int64(0), v1] = (
+ T_batch_matmul_NN[v0, T.int64(0), v1]
+ + T_softmax_norm[v0, T.int64(0), T.int64(0)] *
T_reshape_2[v0, T.int64(0), v1]
+ )
+ for ax0, ax1 in T.grid(T.int64(6), T.int64(64)):
+ with T.block("transpose"):
+ v0, v1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(T_batch_matmul_NN[v0, T.int64(0), v1])
+ T.writes(T_transpose[T.int64(0), T.int64(0), v0, v1])
+ T_transpose[T.int64(0), T.int64(0), v0, v1] =
T_batch_matmul_NN[v0, T.int64(0), v1]
+
+ @T.prim_func
+ def after(
+ T_softmax_norm: T.Buffer((T.int64(6), T.int64(1), T.int64(1)),
"float32"),
+ T_reshape_2: T.Buffer((T.int64(6), T.int64(1), T.int64(64)),
"float32"),
+ T_transpose: T.Buffer((T.int64(1), T.int64(1), T.int64(6),
T.int64(64)), "float32"),
+ ):
+ for ax0, ax1 in T.grid(T.int64(6), T.int64(64)):
+ with T.block("bmm"):
+ v0, v1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(T_softmax_norm[v0, T.int64(0), T.int64(0)],
T_reshape_2[v0, T.int64(0), v1])
+ T.writes(T_transpose[T.int64(0), T.int64(0), v0, v1])
+ with T.init():
+ T_transpose[T.int64(0), T.int64(0), v0, v1] =
T.float32(0.0)
+ T_transpose[T.int64(0), T.int64(0), v0, v1] = (
+ T_transpose[T.int64(0), T.int64(0), v0, v1]
+ + T_softmax_norm[v0, T.int64(0), T.int64(0)] *
T_reshape_2[v0, T.int64(0), v1]
+ )
+
+ sch = tir.Schedule(before)
+ sch.reverse_compute_inline(sch.get_block("transpose"))
+ assert_structural_equal_ignore_global_symbol(after, sch.mod["main"])
+
+
if __name__ == "__main__":
tvm.testing.main()
