From 0f78d0a9f6ce1fd248af3377b8b109dc93c3922d Mon Sep 17 00:00:00 2001 From: yuxiyue Date: Tue, 16 Jul 2024 15:41:49 +0800 Subject: [PATCH] fix(schedule): Inline non-output consumers when padding in TileWithTensorIntrin In the TileWithTensorIntrin function, modified the inlining behavior of consumer blocks. Now, when padding is applied, the function inlines only non-output consumer blocks. This ensures that the padding and inlining process is correctly handled for both producers and consumers. Changes: - Added a check to ensure only non-output consumer blocks are inlined using tir::IsOutputBlock. - Updated the loop iterating over consumers to include the new check. This fix addresses issues where output blocks were being inappropriately inlined, maintaining the correct block shapes and dependencies. --- src/tir/schedule/transform.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/tir/schedule/transform.cc b/src/tir/schedule/transform.cc index 8f912c59ea16..fec214fa1fc7 100644 --- a/src/tir/schedule/transform.cc +++ b/src/tir/schedule/transform.cc @@ -340,7 +340,9 @@ Optional TileWithTensorIntrin(const tir::Schedule& sch, const tir::Block } auto consumers = sch->GetConsumers(block_rv); for (const auto& consumer : consumers) { - sch->ComputeInline(consumer); + auto sref = sch->GetSRef(consumer); + if (!tir::IsOutputBlock(sch->state(), sref, tir::GetScopeRoot(sch->state(), sref, true))) + sch->ComputeInline(consumer); } } // Construct a mapping from tir loops back to LoopRVs