[Mlir-commits] [mlir] [mlir][TilingInterface] Use `LoopLikeOpInterface` in tiling using SCF to unify tiling with `scf.for` and `scf.forall`. (PR #77874)
Matthias Springer
llvmlistbot at llvm.org
Wed Jan 17 01:02:47 PST 2024
================
@@ -622,6 +626,47 @@ LogicalResult scf::ForallOp::promoteIfSingleIteration(RewriterBase &rewriter) {
return success();
}
+Block::BlockArgListType ForallOp::getRegionIterArgs() {
+ return getBody()->getArguments().drop_front(getRank());
+}
+
+MutableArrayRef<OpOperand> ForallOp::getInitsMutable() {
+ return getOutputsMutable();
+}
+
+FailureOr<LoopLikeOpInterface>
+ForallOp::replaceWithAdditionalYields(RewriterBase &rewriter,
+ ValueRange newInitOperands,
+ bool replaceInitOperandUsesInLoop,
+ const NewYieldValuesFn &newYieldValueFn) {
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(getOperation());
+ auto inits = llvm::to_vector(getOutputs());
+ inits.append(newInitOperands.begin(), newInitOperands.end());
+ auto newLoop = rewriter.create<scf::ForallOp>(
+ getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
+ inits, getMapping(), [](OpBuilder &, Location, ValueRange) {});
+
+ // Move the region of the current block to the newly created op.
+ Block *newLoopBody = newLoop.getBody();
+ rewriter.mergeBlocks(
+ getBody(), newLoopBody,
+ newLoopBody->getArguments().take_front(getBody()->getNumArguments()));
+
+ // Update the terminator.
+ {
+ OpBuilder::InsertionGuard g(rewriter);
+ auto terminator = cast<scf::InParallelOp>(newLoopBody->getTerminator());
+ rewriter.setInsertionPointToEnd(terminator.getBody());
+ newYieldValueFn(
----------------
matthias-springer wrote:
The implementations of `yieldTiledValuesAndReplace` make sense to me. But I wouldn't put it in `LoopLikeOpInterface`.
`LoopLikeOpInterface` was meant to expose basic loop properties and low-level transformations of loop-like ops. The new `yieldTiledValuesAndReplace` interface method is tiling-specific and does not belong in `LoopLikeOpInterface` imo. How about putting `yieldTiledValuesAndReplace` in a new interface in `TilingInterface.td`? Then we have all tiling-related stuff in one place.
https://github.com/llvm/llvm-project/pull/77874
More information about the Mlir-commits
mailing list