check if isOpSibling in checkFusionStructuralLegality

Author: srcarroll

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

@@ -160,9 +160,11 @@ void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
 // Fusion related helpers
 //===----------------------------------------------------------------------===//
 
-/// Check structural compatibility between two loops such as iteration space.
+/// Check structural compatibility between two loops such as iteration space
+/// and dominance.
 bool checkFusionStructuralLegality(LoopLikeOpInterface target,
-                                   LoopLikeOpInterface source);
+                                   LoopLikeOpInterface source,
+                                   Diagnostic &diag);
 
 /// Given two scf.forall loops, `target` and `source`, fuses `target` into
 /// `source`. Assumes that the given loops are siblings and are independent of

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

@@ -425,78 +425,6 @@ void transform::TakeAssumedBranchOp::getEffects(
 // LoopFuseSiblingOp
 //===----------------------------------------------------------------------===//
 
-/// Check if `target` and `source` are siblings, in the context that `target`
-/// is being fused into `source`.
-///
-/// This is a simple check that just checks if both operations are in the same
-/// block and some checks to ensure that the fused IR does not violate
-/// dominance.
-static DiagnosedSilenceableFailure isOpSibling(Operation *target,
-                                               Operation *source) {
-  // Check if both operations are same.
-  if (target == source)
-    return emitSilenceableFailure(source)
-           << "target and source need to be different loops";
-
-  // Check if both operations are in the same block.
-  if (target->getBlock() != source->getBlock())
-    return emitSilenceableFailure(source)
-           << "target and source are not in the same block";
-
-  // Check if fusion will violate dominance.
-  DominanceInfo domInfo(source);
-  if (target->isBeforeInBlock(source)) {
-    // Since `target` is before `source`, all users of results of `target`
-    // need to be dominated by `source`.
-    for (Operation *user : target->getUsers()) {
-      if (!domInfo.properlyDominates(source, user, /*enclosingOpOk=*/false)) {
-        return emitSilenceableFailure(target)
-               << "user of results of target should be properly dominated by "
-                  "source";
-      }
-    }
-  } else {
-    // Since `target` is after `source`, all values used by `target` need
-    // to dominate `source`.
-
-    // Check if operands of `target` are dominated by `source`.
-    for (Value operand : target->getOperands()) {
-      Operation *operandOp = operand.getDefiningOp();
-      // Operands without defining operations are block arguments. When `target`
-      // and `source` occur in the same block, these operands dominate `source`.
-      if (!operandOp)
-        continue;
-
-      // Operand's defining operation should properly dominate `source`.
-      if (!domInfo.properlyDominates(operandOp, source,
-                                     /*enclosingOpOk=*/false))
-        return emitSilenceableFailure(target)
-               << "operands of target should be properly dominated by source";
-    }
-
-    // Check if values used by `target` are dominated by `source`.
-    bool failed = false;
-    OpOperand *failedValue = nullptr;
-    visitUsedValuesDefinedAbove(target->getRegions(), [&](OpOperand *operand) {
-      Operation *operandOp = operand->get().getDefiningOp();
-      if (operandOp && !domInfo.properlyDominates(operandOp, source,
-                                                  /*enclosingOpOk=*/false)) {
-        // `operand` is not an argument of an enclosing block and the defining
-        // op of `operand` is outside `target` but does not dominate `source`.
-        failed = true;
-        failedValue = operand;
-      }
-    });
-
-    if (failed)
-      return emitSilenceableFailure(failedValue->getOwner())
-             << "values used inside regions of target should be properly "
-                "dominated by source";
-  }
-
-  return DiagnosedSilenceableFailure::success();
-}
-
 DiagnosedSilenceableFailure
 transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
                                     transform::TransformResults &results,
@@ -520,14 +448,10 @@ transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
     return emitSilenceableFailure(target->getLoc())
            << "target or source is not a loop op";
 
-  // Check if the target and source are siblings.
-  DiagnosedSilenceableFailure diag = isOpSibling(target, source);
-  if (!diag.succeeded())
-    return diag;
-
-  if (!mlir::checkFusionStructuralLegality(target, source))
-    return emitSilenceableFailure(target->getLoc())
-           << "operations cannot be fused";
+  // Check if loops can be fused
+  Diagnostic diag(target.getLoc(), DiagnosticSeverity::Error);
+  if (!mlir::checkFusionStructuralLegality(target, source, diag))
+    return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
 
   Operation *fusedLoop;
   // TODO: Support fusion for loop-like ops besides scf.for, scf.forall

mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp

@@ -136,9 +136,10 @@ verifyDependencies(ParallelOp firstPloop, ParallelOp secondPloop,
 static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
                           const IRMapping &firstToSecondPloopIndices,
                           llvm::function_ref<bool(Value, Value)> mayAlias) {
+  Diagnostic diag(firstPloop.getLoc(), DiagnosticSeverity::Remark);
   return !hasNestedParallelOp(firstPloop) &&
          !hasNestedParallelOp(secondPloop) &&
-         checkFusionStructuralLegality(firstPloop, secondPloop) &&
+         checkFusionStructuralLegality(firstPloop, secondPloop, diag) &&
          succeeded(verifyDependencies(firstPloop, secondPloop,
                                       firstToSecondPloopIndices, mayAlias));
 }

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

@@ -17,6 +17,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -1074,8 +1075,86 @@ TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
 // Fusion related helpers
 //===----------------------------------------------------------------------===//
 
+/// Check if `target` and `source` are siblings, in the context that `target`
+/// is being fused into `source`.
+///
+/// This is a simple check that just checks if both operations are in the same
+/// block and some checks to ensure that the fused IR does not violate
+/// dominance.
+static bool isOpSibling(Operation *target, Operation *source,
+                        Diagnostic &diag) {
+  // Check if both operations are same.
+  if (target == source) {
+    diag << "target and source need to be different loops";
+    return false;
+  }
+
+  // Check if both operations are in the same block.
+  if (target->getBlock() != source->getBlock()) {
+    diag << "target and source are not in the same block";
+    return false;
+  }
+
+  // Check if fusion will violate dominance.
+  DominanceInfo domInfo(source);
+  if (target->isBeforeInBlock(source)) {
+    // Since `target` is before `source`, all users of results of `target`
+    // need to be dominated by `source`.
+    for (Operation *user : target->getUsers()) {
+      if (!domInfo.properlyDominates(source, user, /*enclosingOpOk=*/false)) {
+        diag << "user of results of target should "
+                "be properly dominated by source";
+        return false;
+      }
+    }
+  } else {
+    // Since `target` is after `source`, all values used by `target` need
+    // to dominate `source`.
+
+    // Check if operands of `target` are dominated by `source`.
+    for (Value operand : target->getOperands()) {
+      Operation *operandOp = operand.getDefiningOp();
+      // Operands without defining operations are block arguments. When `target`
+      // and `source` occur in the same block, these operands dominate `source`.
+      if (!operandOp)
+        continue;
+
+      // Operand's defining operation should properly dominate `source`.
+      if (!domInfo.properlyDominates(operandOp, source,
+                                     /*enclosingOpOk=*/false)) {
+        diag << "operands of target should be properly dominated by source";
+        return false;
+      }
+    }
+
+    // Check if values used by `target` are dominated by `source`.
+    bool failed = false;
+    OpOperand *failedValue = nullptr;
+    visitUsedValuesDefinedAbove(target->getRegions(), [&](OpOperand *operand) {
+      Operation *operandOp = operand->get().getDefiningOp();
+      if (operandOp && !domInfo.properlyDominates(operandOp, source,
+                                                  /*enclosingOpOk=*/false)) {
+        // `operand` is not an argument of an enclosing block and the defining
+        // op of `operand` is outside `target` but does not dominate `source`.
+        failed = true;
+        failedValue = operand;
+      }
+    });
+
+    if (failed) {
+      diag << "values used inside regions of target should be properly "
+              "dominated by source";
+      diag.attachNote(failedValue->getOwner()->getLoc()) << "see operation";
+      return false;
+    }
+  }
+
+  return true;
+}
+
 bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
-                                         LoopLikeOpInterface source) {
+                                         LoopLikeOpInterface source,
+                                         Diagnostic &diag) {
   bool iterSpaceEq =
       target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
       target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
@@ -1085,9 +1164,13 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
   // 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;
+    iterSpaceEq =
+        iterSpaceEq && forAllTarget.getMapping() == forAllSource.getMapping();
+  if (!iterSpaceEq) {
+    diag << "target and source iteration spaces must be equal";
+    return false;
+  }
+  return isOpSibling(target, source, diag);
 }
 
 scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,

mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir

@@ -335,8 +335,9 @@ func.func @target_for_region_uses_result_of_source_for_err(%A: tensor<128xf32>,
     %6 = vector.transfer_write %5, %arg4[%arg3] {in_bounds = [true]} : vector<16xf32>, tensor<128xf32>
     scf.yield %6 : tensor<128xf32>
   }
-  %dup1 = scf.for %arg3 = %c0 to %c128 step %c16 iter_args(%arg4 = %B) -> (tensor<128xf32>) {
   // expected-error @below {{values used inside regions of target should be properly dominated by source}}
+  %dup1 = scf.for %arg3 = %c0 to %c128 step %c16 iter_args(%arg4 = %B) -> (tensor<128xf32>) {
+    // expected-note @below {{see operation}}
     %dup2 = vector.transfer_read %1[%arg3], %cst {in_bounds = [true]} : tensor<128xf32>, vector<16xf32>
     %dup3 = vector.transfer_read %arg4[%arg3], %cst {in_bounds = [true]} : tensor<128xf32>, vector<16xf32>
     %dup5 = arith.addf %dup3, %dup2 : vector<16xf32>