[DispatchCreation] Rework horizontal fusion to not create concats.

MaheshRavishankar · MaheshRavishankar · commit 15c9ceab2fdb · 2025-02-01T16:28:54.000-08:00
Signed-off-by: MaheshRavishankar &lt;mahesh.ravishankar@gmail.com&gt;
diff --git a/compiler/src/iree/compiler/DispatchCreation/FuseHorizontalContractions.cpp b/compiler/src/iree/compiler/DispatchCreation/FuseHorizontalContractions.cpp
@@ -7,6 +7,7 @@
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
+#include "iree/compiler/DispatchCreation/FusionUtils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Analysis/TopologicalSortUtils.h"
@@ -50,9 +51,9 @@ struct FuseHorizontalContractionsPass final
 /// Structs that captures the ops that are to be fused
 struct HorizontalFusionGroup {
   // Contractions op that are to be fused.
-  SmallVector<linalg::LinalgOp> contractionOps;
+  SmallVector<Operation *> contractionOps;
   // Optional truncate operations that could be following the contraction op.
-  std::optional<SmallVector<linalg::GenericOp>> truncateOps;
+  std::optional<SmallVector<Operation *>> truncateOps;
 };
 
 /// Helper method to check operations equivalence
@@ -128,11 +129,11 @@ static bool isHorizontalToGroup(Operation *op,
 }
 
 /// Get user of operation that is a truncate operation.
-static std::optional<linalg::GenericOp>
+static std::optional<linalg::LinalgOp>
 getTruncateOp(Operation *op,
               const llvm::SetVector<Operation *> &groupedOperations,
               const DominanceInfo &dominanceInfo,
-              std::optional<linalg::GenericOp> seedTruncateOp = std::nullopt) {
+              std::optional<linalg::LinalgOp> seedTruncateOp = std::nullopt) {
   if (!op->hasOneUse()) {
     return std::nullopt;
   }
@@ -177,7 +178,7 @@ getTruncateOp(Operation *op,
 /// the `truncf` on the result.
 static std::optional<HorizontalFusionGroup> getHorizontalFusionGroupMembers(
     linalg::LinalgOp seedOp,
-    const llvm::SmallDenseSet<linalg::LinalgOp> &groupedOperations,
+    const llvm::SmallDenseSet<Operation *> &groupedOperations,
     const DominanceInfo &dominanceInfo, int fusionLimit) {
 
   Value lhs = seedOp->getOperand(0);
@@ -188,10 +189,10 @@ static std::optional<HorizontalFusionGroup> getHorizontalFusionGroupMembers(
   auto outType = cast<RankedTensorType>(out.getType());
 
   SetVector<Operation *> allOps;
-  SmallVector<linalg::LinalgOp> contractionOps = {seedOp};
-  std::optional<linalg::GenericOp> seedTruncOp =
+  SmallVector<Operation *> contractionOps = {seedOp};
+  std::optional<linalg::LinalgOp> seedTruncOp =
       getTruncateOp(seedOp, allOps, dominanceInfo);
-  std::optional<SmallVector<linalg::GenericOp>> truncateOps;
+  std::optional<SmallVector<Operation *>> truncateOps;
   if (seedTruncOp) {
     truncateOps = {seedTruncOp.value()};
   }
@@ -254,7 +255,7 @@ static std::optional<HorizontalFusionGroup> getHorizontalFusionGroupMembers(
       continue;
     }
 
-    std::optional<linalg::GenericOp> userTruncOp =
+    std::optional<linalg::LinalgOp> userTruncOp =
         getTruncateOp(linalgUser, allOps, dominanceInfo, seedTruncOp);
 
     // If there are truncate ops to fuse and current contraction op
@@ -356,38 +357,140 @@ static AffineMap getConcatenatedIndexingMap(RewriterBase &rewriter,
   return newIndexingMap.insertResult(rewriter.getAffineDimExpr(0), 0);
 }
 
-/// During horizontal fusion, there might be operands of the fused operations
-/// whose definitions are interspersed between the fused operations. For groups
-/// chosen to fuse horizontally, such operations can be moved before the
-/// seed contraction operation (where the fused operation is generated).
-template <typename T>
-static LogicalResult
-moveOperandDefs(RewriterBase &rewriter, ArrayRef<T> operations,
-                Operation *insertionPoint, DominanceInfo &dominanceInfo,
-                ArrayRef<linalg::LinalgOp> ignoreOperations = {}) {
-  BackwardSliceOptions options;
-  llvm::DenseSet<Operation *> ignoreOperationsSet;
-  ignoreOperationsSet.insert(ignoreOperations.begin(), ignoreOperations.end());
-  options.filter = [&](Operation *op) {
-    return !dominanceInfo.properlyDominates(op, insertionPoint) &&
-           !ignoreOperationsSet.contains(op);
-  };
-  // Set inclusive to true cause the slice is computed from the operand, and
-  // we want to include the defining op (which is the point here)
-  options.inclusive = true;
+template <typename V, typename R>
+static void appendRange(V &vector, R &&range) {
+  vector.append(range.begin(), range.end());
+}
 
-  llvm::SetVector<Operation *> slice;
-  for (auto op : operations) {
-    for (auto operand : op->getOperands()) {
-      getBackwardSlice(operand, &slice, options);
+static FailureOr<linalg::GenericOp>
+fuseHorizontally(RewriterBase &rewriter, Location loc,
+                 ArrayRef<Operation *> ops) {
+
+  SmallVector<linalg::LinalgOp> linalgOps;
+  for (auto op : ops) {
+    auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+    if (!linalgOp) {
+      return op->emitOpError("expected linalg op for horizontal fusion");
     }
+    linalgOps.push_back(linalgOp);
   }
 
-  mlir::topologicalSort(slice);
-  for (auto op : slice) {
-    rewriter.moveOpBefore(op, insertionPoint);
+  SmallVector<Value> fusedIns;
+  SmallVector<Value> fusedOuts;
+  SmallVector<Type> fusedResultTypes;
+  SmallVector<AffineMap> fusedInsIndexingMaps;
+  SmallVector<AffineMap> fusedOutsIndexingMaps;
+
+  for (auto linalgOp : linalgOps) {
+    fusedIns.append(linalgOp.getDpsInputs());
+    appendRange(fusedOuts, llvm::map_range(linalgOp.getDpsInitsMutable(),
+                                           [](OpOperand &operand) {
+                                             return operand.get();
+                                           }));
+    fusedResultTypes.append(linalgOp->result_type_begin(),
+                            linalgOp->result_type_end());
+    appendRange(
+        fusedInsIndexingMaps,
+        llvm::map_range(linalgOp.getIndexingMaps().getValue().take_front(
+                            linalgOp.getNumDpsInputs()),
+                        [](Attribute attr) {
+                          return cast<AffineMapAttr>(attr).getValue();
+                        }));
+    appendRange(
+        fusedOutsIndexingMaps,
+        llvm::map_range(linalgOp.getIndexingMaps().getValue().drop_front(
+                            linalgOp.getNumDpsInputs()),
+                        [](Attribute attr) {
+                          return cast<AffineMapAttr>(attr).getValue();
+                        }));
   }
-  return success();
+
+  SmallVector<utils::IteratorType> fusedIteratorTypes =
+      linalgOps.front().getIteratorTypesArray();
+  SmallVector<AffineMap> fusedIndexingMaps = std::move(fusedInsIndexingMaps);
+  fusedIndexingMaps.append(fusedOutsIndexingMaps);
+  auto fusedOp = rewriter.create<linalg::GenericOp>(
+      loc, fusedResultTypes, fusedIns, fusedOuts, fusedIndexingMaps,
+      fusedIteratorTypes, [](OpBuilder &, Location, ValueRange) {});
+
+  Block *fusedBody = fusedOp.getBlock();
+  auto insIndex = 0;
+  auto outsIndex = fusedOp.getNumDpsInputs();
+  SmallVector<Value> yieldVals;
+  for (auto op : linalgOps) {
+    auto linalgOp = cast<linalg::LinalgOp>(op);
+    Block *body = linalgOp.getBlock();
+    SmallVector<Value> replacements = llvm::map_to_vector(
+        fusedBody->getArguments().slice(insIndex, linalgOp.getNumDpsInputs()),
+        [](BlockArgument arg) -> Value { return arg; });
+    appendRange(
+        replacements,
+        llvm::map_range(fusedBody->getArguments().slice(
+                            outsIndex, linalgOp.getNumDpsInits()),
+                        [](BlockArgument arg) -> Value { return arg; }));
+
+    rewriter.mergeBlocks(body, fusedBody, replacements);
+    insIndex += linalgOp.getNumDpsInputs();
+    outsIndex += linalgOp.getNumDpsInits();
+
+    auto yieldOp = cast<linalg::YieldOp>(fusedBody->getTerminator());
+    yieldVals.append(yieldOp->operand_begin(), yieldOp->operand_end());
+    rewriter.eraseOp(yieldOp);
+  }
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPointToEnd(fusedBody);
+  rewriter.create<linalg::YieldOp>(loc, yieldVals);
+
+  auto resultsIndex = 0;
+  for (auto linalgOp : linalgOps) {
+    rewriter.replaceOp(linalgOp, fusedOp->getResults().slice(
+                                     resultsIndex, linalgOp->getNumResults()));
+    resultsIndex += linalgOp->getNumResults();
+  }
+
+  return fusedOp;
+}
+
+static FailureOr<SmallVector<linalg::GenericOp>>
+fuseGroup2(RewriterBase &rewriter, HorizontalFusionGroup &fusionGroup,
+           DominanceInfo &dominanceInfo) {
+  if (!llvm::all_of(fusionGroup.contractionOps,
+                    [](Operation *op) { return isa<linalg::LinalgOp>(op); })) {
+    return failure();
+  }
+  linalg::LinalgOp baseContractOp =
+      cast<linalg::LinalgOp>(fusionGroup.contractionOps.front());
+  Location loc = baseContractOp.getLoc();
+  OpBuilder::InsertionGuard g(rewriter);
+  rewriter.setInsertionPoint(baseContractOp);
+  SmallVector<linalg::GenericOp> fusedOperations;
+
+  if (failed(moveOperandDefs(rewriter, fusionGroup.contractionOps,
+                             baseContractOp, dominanceInfo))) {
+    return baseContractOp.emitOpError("failed to re-order operand definitions");
+  }
+
+  FailureOr<linalg::GenericOp> fusedContractionOp =
+      fuseHorizontally(rewriter, loc, fusionGroup.contractionOps);
+  if (failed(fusedContractionOp)) {
+    return baseContractOp.emitOpError(
+        "failed to fuse contraction ops horizontally");
+  }
+  fusedOperations.push_back(fusedContractionOp.value());
+
+  // if (!fusionGroup.truncateOps) {
+  //   return fusedOperations;
+  // }
+
+  // rewriter.setInsertionPoint(fusionGroup.truncateOps->front());
+  // FailureOr<linalg::GenericOp> fusedTruncateOp =
+  //     fuseHorizontally(rewriter, loc, fusionGroup.truncateOps.value());
+  // if (failed(fusedTruncateOp)) {
+  //   return baseContractOp.emitOpError(
+  //       "failed to fuse truncation ops horizontally");
+  // }
+  // fusedOperations.push_back(fusedTruncateOp.value());
+  return fusedOperations;
 }
 
 /// On finding this pattern
@@ -428,22 +531,23 @@ moveOperandDefs(RewriterBase &rewriter, ArrayRef<T> operations,
 static LogicalResult fuseGroup(RewriterBase &rewriter,
                                HorizontalFusionGroup &fusionGroup,
                                DominanceInfo &dominanceInfo) {
-  linalg::LinalgOp baseContractOp = fusionGroup.contractionOps.front();
+  linalg::LinalgOp baseContractOp =
+      cast<linalg::LinalgOp>(fusionGroup.contractionOps.front());
   Location loc = baseContractOp.getLoc();
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(baseContractOp);
 
-  if (failed(moveOperandDefs(
-          rewriter, ArrayRef<linalg::LinalgOp>(fusionGroup.contractionOps),
-          baseContractOp, dominanceInfo))) {
+  if (failed(moveOperandDefs(rewriter, fusionGroup.contractionOps,
+                             baseContractOp, dominanceInfo))) {
     return baseContractOp.emitOpError("failed to re-order operand definitions");
   }
 
   SmallVector<Value> rhsValues;
   SmallVector<Value> initValues;
   for (auto op : fusionGroup.contractionOps) {
-    Value rhs = op.getDpsInputOperand(1)->get();
-    Value init = op.getDpsInitOperand(0)->get();
+    auto linalgOp = cast<linalg::LinalgOp>(op);
+    Value rhs = linalgOp.getDpsInputOperand(1)->get();
+    Value init = linalgOp.getDpsInitOperand(0)->get();
     rhsValues.push_back(rhs);
     initValues.push_back(init);
   }
@@ -485,15 +589,15 @@ static LogicalResult fuseGroup(RewriterBase &rewriter,
   if (fusionGroup.truncateOps) {
     SmallVector<Value> newTruncOperands;
     SmallVector<AffineMap> newTruncIndexingMaps;
-    linalg::GenericOp baseTruncOp = fusionGroup.truncateOps->front();
+    linalg::LinalgOp baseTruncOp =
+        cast<linalg::LinalgOp>(fusionGroup.truncateOps->front());
     SmallVector<AffineMap> baseTruncOpIndexingMaps =
         baseTruncOp.getIndexingMapsArray();
 
     rewriter.setInsertionPoint(baseTruncOp);
-    if (failed(moveOperandDefs(
-            rewriter,
-            ArrayRef<linalg::GenericOp>(fusionGroup.truncateOps.value()),
-            baseTruncOp, dominanceInfo, fusionGroup.contractionOps))) {
+    if (failed(moveOperandDefs(rewriter, fusionGroup.truncateOps.value(),
+                               baseTruncOp, dominanceInfo,
+                               fusionGroup.contractionOps))) {
       return baseTruncOp.emitOpError(
           "failed to move operand defs for truncate operations");
     }
@@ -503,7 +607,7 @@ static LogicalResult fuseGroup(RewriterBase &rewriter,
       // Collect all the operands for the trunc operation.
       SmallVector<Value> truncOperands;
       for (auto truncOp : fusionGroup.truncateOps.value()) {
-        truncOperands.push_back(truncOp.getOperand(operandIndex));
+        truncOperands.push_back(truncOp->getOperand(operandIndex));
       }
 
       // Three cases to handle here.
@@ -539,7 +643,7 @@ static LogicalResult fuseGroup(RewriterBase &rewriter,
 
     // Insert truncate operator.
     auto baseTruncType =
-        cast<RankedTensorType>(baseTruncOp.getResult(0).getType());
+        cast<RankedTensorType>(baseTruncOp->getResult(0).getType());
     SmallVector<int64_t> newTruncShape = {
         static_cast<int64_t>(rhsValues.size())};
     newTruncShape.append(baseTruncType.getShape().begin(),
@@ -610,7 +714,7 @@ void FuseHorizontalContractionsPass::runOnOperation() {
   DominanceInfo dominanceInfo(getOperation());
 
   SmallVector<HorizontalFusionGroup> horizontalFusionGroups;
-  llvm::SmallDenseSet<linalg::LinalgOp> groupedOperations;
+  llvm::SmallDenseSet<Operation *> groupedOperations;
 
   getOperation()->walk([&](linalg::LinalgOp linalgOp) {
     if (!isEmptyFillContractionDAGRootOp(linalgOp)) {
@@ -653,38 +757,19 @@ void FuseHorizontalContractionsPass::runOnOperation() {
 
   IRRewriter rewriter(context);
   for (auto &fusionGroup : horizontalFusionGroups) {
-    if (failed(fuseGroup(rewriter, fusionGroup, dominanceInfo))) {
+    FailureOr<SmallVector<linalg::GenericOp>> fusedOperations =
+        fuseGroup2(rewriter, fusionGroup, dominanceInfo);
+    if (failed(fusedOperations)) {
       return signalPassFailure();
     }
-  }
-
-  {
-    RewritePatternSet foldReshapePatterns(context);
-    tensor::populateFoldTensorEmptyPatterns(foldReshapePatterns);
-    linalg::FillOp::getCanonicalizationPatterns(foldReshapePatterns, context);
-    if (failed(applyPatternsGreedily(getOperation(),
-                                     std::move(foldReshapePatterns)))) {
-      getOperation()->emitOpError("failed during reshape folding patterns");
-      return signalPassFailure();
-    }
-
-    RewritePatternSet foldPatterns(context);
-    tensor::populateFoldTensorEmptyPatterns(foldPatterns);
-    linalg::FillOp::getCanonicalizationPatterns(foldPatterns, context);
-    if (failed(
-            applyPatternsGreedily(getOperation(), std::move(foldPatterns)))) {
-      getOperation()->emitOpError("failed to fold empty/fill with concats");
-      return signalPassFailure();
+    for (auto fusedOp : fusedOperations.value()) {
+      rewriter.setInsertionPoint(fusedOp);
+      if (failed(linalg::deduplicateOperandsAndRemoveDeadResults(
+              rewriter, fusedOp, /*removeOutputs=*/false))) {
+        fusedOp->emitOpError("failed to remove duplicate operands");
+        return signalPassFailure();
+      }
     }
   }
-
-  // Note: Currently these patterns are required due to early lowering of
-  // tensor.concat. When we choose  to move the lowering of tensor.concat later,
-  // these patterns should be dropped.
-  RewritePatternSet patterns(context);
-  tensor::populateDecomposeTensorConcatPatterns(patterns);
-  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
-    return signalPassFailure();
-  }
 }
 } // namespace mlir::iree_compiler::DispatchCreation
