[Dispatch Creation] Don't add unfusable consumers to fusion group (#22461)

There are cases where two ops can't be fused as it would result in
use-def violation (see the diagram below). Currently, they are still
placed in the same fusion group. This means that the consumer gets
marked by the analysis as being in a fusion group causing it to
potentially miss out on fusion opportunities. This change moves this
check so that it occurs during the analysis and modifies the old check
to error out.


```
    A (in fusion group)
    |  \
    |   \
    |    v
    |    B (unfusable consumer of A)
    |   /
    |  /
    v v
    C (trying to fuse with A)
```

---------

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp

@@ -20,6 +20,7 @@
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
+#include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -43,6 +44,7 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/RegionUtils.h"
 
 #define DEBUG_TYPE "iree-dispatch-creation-form-dispatch-regions"
 
@@ -145,6 +147,50 @@ class FusionGroup {
   // Insert `op` into the fusion group.
   void insert(Operation *op);
 
+  /// Returns true if `consumerOp` has a transitive dependency on the fusion
+  /// group. This means that some transitive dependency of `consumerOp` (not in
+  /// the fusion group) itself uses an operation in the fusion group. This is
+  /// required for fusion because it must be legal to take a program slice that
+  /// contains only the ops in the fusion group.
+  bool
+  hasTransitiveDependencyOnFusionGroup(Operation *consumerOp,
+                                       DominanceInfo const &dominance) const {
+    BackwardSliceOptions options;
+    options.inclusive = true;
+    options.omitUsesFromAbove = false;
+    options.omitBlockArguments = true;
+    options.filter = [&](Operation *sliceBoundaryOp) {
+      return !llvm::all_of(
+          loopMaps.getArrayRef(), [&](std::pair<Operation *, AffineMap> pair) {
+            return dominance.properlyDominates(sliceBoundaryOp, pair.first);
+          });
+    };
+
+    llvm::SetVector<Operation *> slice;
+    auto populateSlice = [&](OpOperand *operand) {
+      // It's okay if the consumer directly uses an operation in the fusion
+      // group.
+      if (loopMaps.contains(operand->get().getDefiningOp())) {
+        return;
+      }
+      LogicalResult result = getBackwardSlice(operand->get(), &slice, options);
+      assert(result.succeeded() && "expected a backward slice");
+      (void)result;
+    };
+
+    // Search all of the operands op `consumerOp` as well as all the values used
+    // in its regions.
+    mlir::visitUsedValuesDefinedAbove(consumerOp->getRegions(), populateSlice);
+    for (OpOperand &operand : consumerOp->getOpOperands()) {
+      populateSlice(&operand);
+    }
+
+    return llvm::any_of(loopMaps.getArrayRef(),
+                        [&](std::pair<Operation *, AffineMap> pair) {
+                          return slice.contains(pair.first);
+                        });
+  }
+
 private:
   Operation *rootOp;
   // All operations to be fused with the root op. This does not include
@@ -435,6 +481,9 @@ getFusableUses(MLIRContext *context, Operation *op,
     if (isa<tensor::DimOp>(user)) {
       continue;
     }
+    if (op->getBlock() != user->getBlock()) {
+      continue;
+    }
     fusableUses.insert(&use);
   }
 
@@ -667,6 +716,13 @@ fuseRootsWithConsumers(MLIRContext *context, ArrayRef<Operation *> roots,
           continue;
         }
 
+        // Ensure that fusing the consumer would not cause use-def violations.
+        if (tracker.getFusionGroup(currRoot)
+                .hasTransitiveDependencyOnFusionGroup(fusableUse->getOwner(),
+                                                      dominanceInfo)) {
+          continue;
+        }
+
         if (isFusableWithConsumer(*fusableUse, tracker, options)) {
           tracker.appendToFusionGroup(consumerOp, fusionGroup);
           workList.push_back(consumerOp);
@@ -957,7 +1013,8 @@ createFusionGroups(TensorDimTrackingRewriter &rewriter,
       auto newRegionOp =
           movePrecedingOpsIntoDispatchRegion(rewriter, producer, regionOp);
       if (failed(newRegionOp)) {
-        return producer->emitOpError("failed to move producer into region");
+        producer->emitWarning("failed to move producer into region");
+        continue;
       }
       regionOp = *newRegionOp;
     }
@@ -974,7 +1031,7 @@ createFusionGroups(TensorDimTrackingRewriter &rewriter,
       auto newRegionOp = IREE::Flow::moveFollowingOpIntoDispatchRegion(
           rewriter, consumer, regionOp);
       if (failed(newRegionOp)) {
-        continue;
+        return consumer->emitOpError("failed to move consumer into region");
       }
       regionOp = *newRegionOp;
     }

compiler/src/iree/compiler/DispatchCreation/test/form_dispatch_regions.mlir

@@ -1391,15 +1391,16 @@ util.func public @avoid_illegal_consumer_fusion(%arg0: tensor<75600x5120xbf16>)
   util.return %6 : tensor<75600x1x5120xbf16>
 }
 // CHECK-LABEL: @avoid_illegal_consumer_fusion(
-//       CHECK:   %[[DISPATCH:.+]]:2 = flow.dispatch.region
+//       CHECK:   %[[DISPATCH0:.+]]:2 = flow.dispatch.region
 //       CHECK:     %[[GENERIC0:.+]] = linalg.generic
 //       CHECK:     %[[GENERIC1:.+]] = linalg.generic
 //  CHECK-SAME:         ins(%[[GENERIC0]] :
 //       CHECK:     flow.return %[[GENERIC1]], %[[GENERIC0]]
-//       CHECK:   %[[EXPAND_SHAPE:.+]] = tensor.expand_shape %[[DISPATCH]]#1
+//       CHECK:   %[[EXPAND_SHAPE:.+]] = tensor.expand_shape %[[DISPATCH0]]#1
+//       CHECK:   %[[DISPATCH1:.+]] = flow.dispatch.region
 //       CHECK:   %[[GENERIC2:.+]] = linalg.generic
-//  CHECK-SAME:       ins(%[[EXPAND_SHAPE]], %[[DISPATCH]]#0 :
-//       CHECK:   util.return %[[GENERIC2]]
+//  CHECK-SAME:       ins(%[[EXPAND_SHAPE]], %[[DISPATCH0]]#0 :
+//       CHECK:   util.return %[[DISPATCH1]]
 
 // -----
 
@@ -1791,3 +1792,48 @@ util.func public @dont_fuse_producer_matmuls(%arg0 : tensor<4x4x7xf32>, %arg1 :
 //       CHECK:     %[[OP:.+]] = linalg.generic
 //  CHECK-SAME:       ins(%[[DISPATCH0]]
 //       CHECK:     flow.return %[[OP]]
+
+// -----
+
+util.func public @no_fusion_across_blocks(%arg0: tensor<3x2xf32>) -> tensor<f32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<f32>
+  %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<f32>) -> tensor<f32>
+  %2 = tensor.empty() : tensor<3x2xf32>
+  %4 = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> ()>],
+      iterator_types = ["reduction", "reduction"]}
+      ins(%arg0 : tensor<3x2xf32>) outs(%1 : tensor<f32>) {
+  ^bb0(%in: f32, %out: f32):
+    %9 = arith.addf %in, %out : f32
+    linalg.yield %9 : f32
+  } -> tensor<f32>
+  // Dispatch region uses the reduction result
+  %5 = flow.dispatch.region -> (tensor<f32>) {
+    %9 = linalg.generic {
+        indexing_maps = [affine_map<() -> ()>,
+                         affine_map<() -> ()>,
+                         affine_map<() -> ()>],
+        iterator_types = []}
+        ins(%4, %1 : tensor<f32>, tensor<f32>) outs(%0 : tensor<f32>) {
+    ^bb0(%in: f32, %in_0: f32, %out: f32):
+      %10 = arith.divf %in, %in_0 : f32
+      linalg.yield %10 : f32
+    } -> tensor<f32>
+    flow.return %9 : tensor<f32>
+  }
+  util.return %5 : tensor<f32>
+}
+// CHECK-LABEL: util.func public @no_fusion_across_blocks
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<3x2xf32>
+//       CHECK:   %[[FILL:.+]] = linalg.fill
+//       CHECK:   %[[DISPATCH0:.+]] = flow.dispatch.region
+//       CHECK:     %[[REDUCTION:.+]] = linalg.generic
+//  CHECK-SAME:       ins(%[[ARG0]]
+//       CHECK:     flow.return %[[REDUCTION]]
+//       CHECK:   %[[DISPATCH1:.+]] = flow.dispatch.region
+//       CHECK:     %[[DIV:.+]] = linalg.generic
+//  CHECK-SAME:       ins(%[[DISPATCH0]], %[[FILL]]
+//       CHECK:     flow.return %[[DIV]]
+//       CHECK:   util.return %[[DISPATCH1]]