Skip to content

Commit

Permalink
fix(schedule): Inline non-output consumers when padding in TileWithTe…
Browse files Browse the repository at this point in the history
…nsorIntrin

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.
  • Loading branch information
yuxiyue authored and yuxiyue committed Jul 16, 2024
1 parent b654852 commit 0f78d0a
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion src/tir/schedule/transform.cc
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,9 @@ Optional<LoopRV> 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
Expand Down

0 comments on commit 0f78d0a

Please sign in to comment.