move createFused to LoopLikeInterface.h

llvm · srcarroll · Jul 3, 2024 · Jun 5, 2024 · Jun 4, 2024 · Jun 5, 2024
commit cc95d75d2cc09f8a33850f3867c8313e374a0dfd
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -57,4 +57,24 @@ class HasParallelRegion : public TraitBase<ConcreteType, HasParallelRegion> {
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/LoopLikeInterface.h.inc"
 
+namespace mlir {
+/// A function that rewrites `target`'s terminator as a teminator obtained by
+/// fusing `source` into `target`.
+using FuseTerminatorFn =
+    std::function<void(RewriterBase &rewriter, LoopLikeOpInterface source,
+                       LoopLikeOpInterface &target, IRMapping mapping)>;
+
+/// Returns a fused `LoopLikeOpInterface` created by fusing `source` to
+/// `target`.  The `NewYieldValuesFn` callback is used to pass to the
+/// `replaceWithAdditionalYields` interface method to replace the loop with a
+/// new loop with (possibly) additional yields, while the `FuseTerminatorFn`
+/// callback is repsonsible for updating the fused loop terminator.
+LoopLikeOpInterface createFused(LoopLikeOpInterface target,
+                                LoopLikeOpInterface source,
+   
10000
                             RewriterBase &rewriter,
+                                NewYieldValuesFn newYieldValuesFn,
+                                FuseTerminatorFn fuseTerminatorFn);
+
+} // namespace mlir
+
 #endif // MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1082,93 +1082,14 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
       target.getLoopSteps() == source.getLoopSteps();
   auto forAllTarget = dyn_cast<scf::ForallOp>(*target);
   auto forAllSource = dyn_cast<scf::ForallOp>(*source);
+  // TODO: Decouple checks on concrete loop types and move this function
+  // somewhere for general utility for `LoopLikeOpInterface`
   if (forAllTarget && forAllSource)
     return iterSpaceEq &&
            forAllTarget.getMapping() == forAllSource.getMapping();
   return iterSpaceEq;
 }
 
-template <typename LoopTy>
-void fuseTerminator(RewriterBase &rewriter, LoopTy source, LoopTy &fused,
-                    IRMapping &mapping) {}
-
-template <>
-void fuseTerminator(RewriterBase &rewriter, scf::ForallOp source,
-                    scf::ForallOp &fused, IRMapping &mapping) {
-  // Fuse the old terminator in_parallel ops into the new one.
-  scf::InParallelOp fusedTerm = fused.getTerminator();
-  rewriter.setInsertionPointToEnd(fusedTerm.getBody());
-  for (Operation &op : source.getTerminator().getYieldingOps())
-    rewriter.clone(op, mapping);
-}
-
-template <>
-void fuseTerminator(RewriterBase &rewriter, scf::ForOp source,
-                    scf::ForOp &fused, IRMapping &mapping) {
-  // Build fused yield results by appropriately mapping original yield operands.
-  auto newTerm = rewriter.clone(*fused.getBody()->getTerminator(), mapping);
-  rewriter.replaceOp(fused.getBody()->getTerminator(), newTerm);
-}
-
-// TODO: We should maybe add a method to LoopLikeOpInterface that will
-// facilitate this transformation. For now, this acts as a placeholder.
-template <>
-void fuseTerminator(RewriterBase &rewriter, LoopLikeOpInterface source,
-                    LoopLikeOpInterface &fused, IRMapping &mapping) {
-  if (isa<scf::ForOp>(source) && isa<scf::ForOp>(fused)) {
-    fuseTerminator(rewriter, cast<scf::ForOp>(source), cast<scf::ForOp>(fused),
-                   mapping);
-  } else if (isa<scf::ForallOp>(source) && isa<scf::ForallOp>(fused)) {
-    fuseTerminator(rewriter, cast<scf::ForallOp>(source),
-                   cast<scf::ForallOp>(fused), mapping);
-  } else if (isa<scf::ParallelOp>(source) && isa<scf::ParallelOp>(fused)) {
-    fuseTerminator(rewriter, cast<scf::ParallelOp>(source),
-                   cast<scf::ParallelOp>(fused), mapping);
-  } else {
-    llvm_unreachable("unsupported loop types.");
-    return;
-  }
-}
-
-LoopLikeOpInterface createFused(LoopLikeOpInterface target,
-                                LoopLikeOpInterface source,
-                                RewriterBase &rewriter,
-                                NewYieldValuesFn newYieldValuesFn) {
-  auto targetIterArgs = target.getRegionIterArgs();
-  auto targetInductionVar = *target.getLoopInductionVars();
-  SmallVector<Value> targetYieldOperands(target.getYieldedValues());
-  auto sourceIterArgs = source.getRegionIterArgs();
-  auto sourceInductionVar = *source.getLoopInductionVars();
-  SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
-  auto sourceRegion = source.getLoopRegions().front();
-  LoopLikeOpInterface fusedLoop = *target.replaceWithAdditionalYields(
-      rewriter, source.getInits(), /*replaceInitOperandUsesInLoop=*/false,
-      newYieldValuesFn);
-
-  // Map control operands.
-  IRMapping mapping;
-  mapping.map(targetInductionVar, *fusedLoop.getLoopInductionVars());
-  mapping.map(targetIterArgs,
-              fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
-  mapping.map(targetYieldOperands,
-              fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
-  mapping.map(sourceInductionVar, *fusedLoop.getLoopInductionVars());
-  mapping.map(sourceIterArgs,
-              fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
-  mapping.map(sourceYieldOpera
8000
nds,
-              fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
-  // Append everything except the terminator into the fused operation.
-  rewriter.setInsertionPoint(
-      fusedLoop.getLoopRegions().front()->front().getTerminator());
-  for (Operation &op : sourceRegion->front().without_terminator())
-    rewriter.clone(op, mapping);
-
-  // TODO: Replace with corresponding interface method if added
-  fuseTerminator(rewriter, source, fusedLoop, mapping);
-
-  return fusedLoop;
-}
-
 scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
                                                       scf::ForallOp source,
                                                       RewriterBase &rewriter) {
@@ -1177,6 +1098,15 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
       [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
         // `ForallOp` does not have yields, rather an `InParallelOp` terminator.
         return ValueRange{};
+      },
+      [&](RewriterBase &b, LoopLikeOpInterface source,
+          LoopLikeOpInterface &target, IRMapping mapping) {
+        auto sourceForall = cast<scf::ForallOp>(source);
+        auto targetForall = cast<scf::ForallOp>(target);
+        scf::InParallelOp fusedTerm = targetForall.getTerminator();
+        b.setInsertionPointToEnd(fusedTerm.getBody());
+        for (Operation &op : sourceForall.getTerminator().getYieldingOps())
+          b.clone(op, mapping);
       }));
   rewriter.replaceOp(source,
                      fusedLoop.getResults().take_back(source.getNumResults()));
@@ -1191,12 +1121,21 @@ scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
       target, source, rewriter,
       [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
         return source.getYieldedValues();
+      },
+      [&](RewriterBase &b, LoopLikeOpInterface source,
+          LoopLikeOpInterface &target, IRMapping mapping) {
+        auto sourceFor = cast<scf::ForOp>(source);
+        auto targetFor = cast<scf::ForOp>(target);
+        auto newTerm = b.clone(*targetFor.getBody()->getTerminator(), mapping);
+        b.replaceOp(targetFor.getBody()->getTerminator(), newTerm);
       }));
   rewriter.replaceOp(source,
                      fusedLoop.getResults().take_back(source.getNumResults()));
   return fusedLoop;
 }
 
+// TODO: Finish refactoring this a la the above, but likely requires additional
+// interface methods.
 scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
     scf::ParallelOp target, scf::ParallelOp source, RewriterBase &rewriter) {
   Block *block1 = target.getBody();

diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -8,6 +8,8 @@
 
 #include "mlir/Interfaces/LoopLikeInterface.h"
 
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "llvm/ADT/DenseSet.h"
 
@@ -113,3 +115,43 @@ LogicalResult detail::verifyLoopLikeOpInterface(Operation *op) {
 
   return success();
 }
+
+LoopLikeOpInterface mlir::createFused(LoopLikeOpInterface target,
+                                      LoopLikeOpInterface source,
+                                      RewriterBase &rewriter,
+                                      NewYieldValuesFn newYieldValuesFn,
+                                      FuseTerminatorFn fuseTerminatorFn) {
+  auto targetIterArgs = target.getRegionIterArgs();
+  auto targetInductionVar = *target.getLoopInductionVars();
+  SmallVector<Value> targetYieldOperands(target.getYieldedValues());
+  auto sourceIterArgs = source.getRegionIterArgs();
+  auto sourceInductionVar = *source.getLoopInductionVars();
+  SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
+  auto sourceRegion = source.getLoopRegions().front();
+  LoopLikeOpInterface fusedLoop = *target.replaceWithAdditionalYields(
+      rewriter, source.getInits(), /*replaceInitOperandUsesInLoop=*/false,
+      newYieldValuesFn);
+
+  // Map control operands.
+  IRMapping mapping;
+  mapping.map(targetInductionVar, *fusedLoop.getLoopInductionVars());
+  mapping.map(targetIterArgs,
+              fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
+  mapping.map(targetYieldOperands,
+              fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
+  mapping.map(sourceInductionVar, *fusedLoop.getLoopInductionVars());
+  mapping.map(sourceIterArgs,
+              fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
+  mapping.map(sourceYieldOperands,
+              fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
+  // Append everything except the terminator into the fused operation.
+  rewriter.setInsertionPoint(
+      fusedLoop.getLoopRegions().front()->front().getTerminator());
+  for (Operation &op : sourceRegion->front().without_terminator())
+    rewriter.clone(op, mapping);
+
+  // TODO: Replace with corresponding interface method if added
+  fuseTerminatorFn(rewriter, source, fusedLoop, mapping);
+
+  return fusedLoop;
+}