[Mlir-commits] [mlir] [SCF][Transform] Add support for scf.for in LoopFuseSibling op (PR #81495)
Oleksandr Alex Zinenko
llvmlistbot at llvm.org
Tue Feb 13 03:56:53 PST 2024
================
@@ -970,3 +970,69 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
return fusedLoop;
}
+
+scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
+ scf::ForOp source,
+ RewriterBase &rewriter) {
+ // Create fused init_args.
+ auto targetInitArgs = target.getInitArgs();
+ auto sourceInitArgs = source.getInitArgs();
+ SmallVector<Value> fusedInitArgs;
+ fusedInitArgs.reserve(targetInitArgs.size() + sourceInitArgs.size());
+ fusedInitArgs.append(sourceInitArgs.begin(), sourceInitArgs.end());
+ fusedInitArgs.append(targetInitArgs.begin(), targetInitArgs.end());
+
+ // Create a new scf::for op after the source loop.
+ rewriter.setInsertionPointAfter(source);
+ scf::ForOp fusedLoop = rewriter.create<scf::ForOp>(
+ source.getLoc(), source.getLowerBound(), source.getUpperBound(),
+ source.getStep(), fusedInitArgs);
+
+ SmallVector<Value> yieldResults;
+
+ // First merge source loop into the new (fused) for loop and then target loop.
+ rewriter.setInsertionPointToEnd(fusedLoop.getBody());
+ for (auto loopAndInitArgsBegin :
+ {std::pair(source, (unsigned int)0),
+ std::pair(target, source.getNumRegionIterArgs())}) {
+ auto origLoop = loopAndInitArgsBegin.first;
----------------
ftynse wrote:
Nit: you can unpack tuples now `for (auto &&[loop, position] : {...})`.
https://github.com/llvm/llvm-project/pull/81495
More information about the Mlir-commits
mailing list