add checkFusionStructuralLegality

Author: srcarroll

mlir/include/mlir/Dialect/SCF/Utils/Utils.h

@@ -156,6 +156,13 @@ Loops tilePerfectlyNested(scf::ForOp rootForOp, ArrayRef<Value> sizes);
 void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
                              scf::ForOp root);
 
+//===----------------------------------------------------------------------===//
+// Fusion related helpers
+//===----------------------------------------------------------------------===//
+
+template <typename LoopTy>
+bool checkFusionStructuralLegality(Operation *target, Operation *source);
+
 /// Prepends operations of firstPloop's body into secondPloop's body.
 /// Updates secondPloop with new loop.
 void fuseIfLegal(scf::ParallelOp firstPloop, scf::ParallelOp &secondPloop,

mlir/lib/Dialect/SCF/Utils/Utils.cpp

@@ -1184,13 +1184,39 @@ static bool equalIterationSpaces(scf::ParallelOp firstPloop,
          matchOperands(firstPloop.getStep(), secondPloop.getStep());
 }
 
+//===----------------------------------------------------------------------===//
+// Fusion related helpers
+//===----------------------------------------------------------------------===//
+
 /// Verify there are no nested ParallelOps.
 static bool hasNestedParallelOp(scf::ParallelOp ploop) {
   auto walkResult = ploop.getBody()->walk(
       [](scf::ParallelOp) { return WalkResult::interrupt(); });
   return walkResult.wasInterrupted();
 }
 
+template <typename LoopTy>
+static bool checkFusionStructuralLegality(Operation *target,
+                                             Operation *source) {
+  static_assert(llvm::is_one_of<LoopTy, scf::ForallOp, scf::ForOp,
+                                scf::ParallelOp>::value,
+                "applies to only `forall`, `for` and `parallel`");
+  auto targetOp = dyn_cast<LoopTy>(target);
+  auto sourceOp = dyn_cast<LoopTy>(source);
+  if (!targetOp || !sourceOp)
+    return false;
+
+  if constexpr (std::is_same_v<LoopTy, scf::ForallOp>)
+    return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
+           targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
+           targetOp.getMixedStep() == sourceOp.getMixedStep() &&
+           targetOp.getMapping() == sourceOp.getMapping();
+  else
+    return targetOp.getLowerBound() == sourceOp.getLowerBound() &&
+           targetOp.getUpperBound() == sourceOp.getUpperBound() &&
+           targetOp.getStep() == sourceOp.getStep();
+}
+
 static bool isFusionLegal(scf::ParallelOp firstPloop,
                           scf::ParallelOp secondPloop,
                           const IRMapping &firstToSecondPloopIndices,