[mlir][Linalg] Make Elementwise op fusion return a map from existing values to values in the fused op.

[mravishankar]

Summary

This replacement can be used to eliminate all uses of the
producer/consumer for case where producer/consumer has other uses
outside of the producer/consumer pair. This makes the
producer/consumer dead.

Add test and minor fixup to the test harness.

[hanchung]

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

165–171

The comment above is for the fuseElementwiseOps method, maybe we can move the struct above the comment?

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

427–428

It is assuming that all replacements are appending to the result of fused op, which is fragile. Should we take replacements from fusionResult?

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

69–76

style nit: we dont need braces for single if statements.

[mravishankar]

Address comments.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

165–171

I think its just easier to keep it here. Really its part of the function ABI.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// Return true if two `linalg.generic` operations with producer/consumer
 /// relationship through `fusedOperand` can be fused using elementwise op
 /// fusion.
 bool areElementwiseOpsFusable(OpOperand *fusedOperand);
 
 /// Fuse two `linalg.generic` operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
-FailureOr<Operation *> fuseElementwiseOps(RewriterBase &rewriter,
-                                          OpOperand *fusedOperand);
+struct ElementwiseOpFusionResult {
+  Operation *fusedOp;
+  llvm::DenseMap<Value, Value> replacements;
+};
+FailureOr<ElementwiseOpFusionResult>
+fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand);
 
 /// Split the given `op` into two parts along the given iteration space
 /// `dimension` at the specified `splitPoint`, and return the two parts.
 /// If the second part is statically known to be empty, do not create it
 /// and return nullptr instead. Error state is signalled by returning
 /// a pair of nullptrs.
 ///
 /// For example, the following op:
 ///

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include <utility>
 #include <optional>
+#include <utility>
 
 namespace mlir {
 #define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS
 #define GEN_PASS_DEF_LINALGELEMENTWISEOPFUSION
 #include "mlir/Dialect/Linalg/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
static AffineMap getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(
 
   // Compose t1 with fusedConsumerArgIndexMap gives an indexing map from
   // consumer loop/ fused loop -> producer arg tensor index.
   return t1.compose(fusedConsumerArgIndexMap);
 }
 
 /// Conditions for elementwise fusion of generic operations.
 bool mlir::linalg::areElementwiseOpsFusable(OpOperand *fusedOperand) {
+  if (!fusedOperand)
+    return false;
+
   auto producer = fusedOperand->get().getDefiningOp<GenericOp>();
   auto consumer = dyn_cast<GenericOp>(fusedOperand->getOwner());
 
   // Check producer and consumer are generic ops.
   if (!producer || !consumer)
     return false;
 
   // Consumer can have mixed semantics, just check operand itself has tensor
for (auto consumerYieldVal : consumerYieldOp.getOperands())
     fusedYieldValues.push_back(mapper.lookupOrDefault(consumerYieldVal));
   rewriter.create<YieldOp>(fusedOp.getLoc(), fusedYieldValues);
 
   // Sanity checks.
   assert(fusedBlock->getNumArguments() == fusedOp.getNumOperands() &&
          "Ill-formed GenericOp region");
 }
 
-FailureOr<Operation *>
+FailureOr<mlir::linalg::ElementwiseOpFusionResult>
 mlir::linalg::fuseElementwiseOps(RewriterBase &rewriter,
                                  OpOperand *fusedOperand) {
   assert(areElementwiseOpsFusable(fusedOperand) &&
          "expected elementwise operation pre-conditions to pass");
   auto producerResult = fusedOperand->get().cast<OpResult>();
   auto producer = cast<GenericOp>(producerResult.getOwner());
   auto consumer = cast<GenericOp>(fusedOperand->getOwner());
   // TODO: allow fusing the producer of an output operand.
assert(invProducerResultIndexMap &&
          "expected producer result indexig map to be invertible");
   // consumer loop -> producer loop
   AffineMap consumerToProducerLoopsMap =
       invProducerResultIndexMap.compose(consumerResultIndexMap);
 
   generateFusedElementwiseOpRegion(
       rewriter, fusedOp, consumerToProducerLoopsMap, fusedOperand,
       consumer.getNumLoops(), preservedProducerResults);
-  return fusedOp.getOperation();
+  ElementwiseOpFusionResult result;
+  result.fusedOp = fusedOp;
+  int resultNum = 0;
+  for (auto [index, producerResult] : llvm::enumerate(producer->getResults()))
+    if (preservedProducerResults.count(index))
+      result.replacements[producerResult] = fusedOp->getResult(resultNum++);
+  for (auto consumerResult : consumer->getResults())
+    result.replacements[consumerResult] = fusedOp->getResult(resultNum++);
+  return result;
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
 public:
   FuseElementwiseOps(MLIRContext *context, ControlFusionFn fun,
                      PatternBenefit benefit = 1)
       : OpRewritePattern<GenericOp>(context, benefit),
         controlFn(std::move(fun)) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     // Find the first operand that is defined by another generic op on tensors.
     for (OpOperand &opOperand : genericOp->getOpOperands()) {
       if (!areElementwiseOpsFusable(&opOperand))
         continue;
       if (!controlFn(&opOperand))
         continue;
 
-      FailureOr<Operation *> fusedOp = fuseElementwiseOps(rewriter, &opOperand);
-      if (succeeded(fusedOp)) {
-        auto replacements =
-            (*fusedOp)->getResults().take_back(genericOp.getNumResults());
-        rewriter.replaceOp(genericOp, replacements);
-        return success();
+      FailureOr<ElementwiseOpFusionResult> fusionResult =
+          fuseElementwiseOps(rewriter, &opOperand);
+      if (failed(fusionResult))
+        rewriter.notifyMatchFailure(genericOp, "fusion failed");
+      Operation *producer = opOperand.get().getDefiningOp();
+      for (auto [origVal, replacement] : fusionResult->replacements) {
+        Value origValCopy = origVal;
+        rewriter.replaceUseIf(origVal, replacement, [&](OpOperand &use) {
+          // Only replace consumer uses.
+          return use.get().getDefiningOp() != producer;
+        });
       }
+      rewriter.eraseOp(genericOp);
+      return success();
     }
     return failure();
   }
 
 private:
   ControlFusionFn controlFn;
 };
 } // namespace

mlir/test/Dialect/Linalg/fusion-multiuse-producer.mlir

+// RUN: mlir-opt -test-linalg-elementwise-fusion-patterns=fuse-multiuse-producer -split-input-file %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @multi_use_producer(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
+    %arg2 : tensor<?x?xf32>, %arg3 : tensor<?x?xf32>, %arg4 : tensor<?x?xf32>)
+    -> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
+  %0:2 = linalg.generic {
+      indexing_maps = [#map, #map, #map],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%arg0 : tensor<?x?xf32>)
+      outs(%arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>) {
+  ^bb0(%b0: f32, %b1 : f32, %b2 : f32):
+    %1 = arith.addf %b0, %b1 : f32
+    linalg.yield %1, %1 : f32, f32
+  } -> (tensor<?x?xf32>, tensor<?x?xf32>)
+  %2 = linalg.generic {
+      indexing_maps = [#map, #map, #map],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%0#1, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg4 : tensor<?x?xf32>) {
+  ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
+    %3 = arith.mulf %b0, %b1 : f32
+    linalg.yield %3 : f32
+  } -> tensor<?x?xf32>
+  return %0#0, %0#1, %2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
+}
+//      CHECK: func @multi_use_producer(
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
+//      CHECK:   %[[RESULT:.+]]:3 = linalg.generic
+//      CHECK:   return %[[RESULT]]#0, %[[RESULT]]#1, %[[RESULT]]#2

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

if (producer->getNumResults() != 1)
     return false;
   addOperands(consumer, fusedOpOperands);
   fusedOpOperands.remove(producer->getResult(0));
   addOperands(producer, fusedOpOperands);
   return fusedOpOperands.size() <= limit;
 }
 
 namespace {
+
+/// Pattern to test fusion of producer with consumer, even if producer has
+/// multiple uses.
+struct TestMultiUseProducerFusion : public OpRewritePattern<linalg::GenericOp> {
+  using OpRewritePattern<linalg::GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    OpOperand *fusableOperand = nullptr;
+    for (OpOperand &operand : genericOp->getOpOperands()) {
+      if (linalg::areElementwiseOpsFusable(&operand)) {
+        fusableOperand = &operand;
+        break;
+      }
+    }
+    if (!fusableOperand) {
+      return rewriter.notifyMatchFailure(genericOp, "no fusable operand found");
+    }
+    std::optional<linalg::ElementwiseOpFusionResult> fusionResult =
+        linalg::fuseElementwiseOps(rewriter, fusableOperand);
+    if (!fusionResult)
+      rewriter.notifyMatchFailure(genericOp, "fusion failed");
+    for (auto [origValue, replacement] : fusionResult->replacements) {
+      rewriter.replaceUseIf(origValue, replacement, [&](OpOperand &use) {
+        return use.getOwner() != genericOp.getOperation();
+      });
+    }
+    rewriter.eraseOp(genericOp);
+    return success();
+  }
+};
+
 struct TestLinalgElementwiseFusion
     : public PassWrapper<TestLinalgElementwiseFusion,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgElementwiseFusion)
 
   TestLinalgElementwiseFusion() = default;
   TestLinalgElementwiseFusion(const TestLinalgElementwiseFusion &pass)
       : PassWrapper(pass) {}
Option<bool> fuseWithReshapeByCollapsing{
       llvm::cl::init(false)};
 
   Option<bool> fuseWithReshapeByCollapsingWithControlFn{
       *this, "fuse-with-reshape-by-collapsing-control",
       llvm::cl::desc("Test controlling the linalg expand_shape -> generic "
                      "fusion patterns that "
                      "collapse the iteration space of the consumer"),
       llvm::cl::init(false)};
+
+  Option<bool> fuseMultiUseProducer{
+      *this, "fuse-multiuse-producer",
+      llvm::cl::desc("Test fusion of producer ops with multiple uses"),
+      llvm::cl::init(false)};
+
   ListOption<int64_t> collapseDimensions{
       *this, "collapse-dimensions-control",
       llvm::cl::desc("Test controlling dimension collapse pattern")};
 
   void runOnOperation() override {
     MLIRContext *context = &this->getContext();
     func::FuncOp funcOp = this->getOperation();
 
     if (fuseGenericOps) {
       RewritePatternSet fusionPatterns(context);
       auto controlFn = [](OpOperand *operand) { return true; };
       linalg::populateElementwiseOpsFusionPatterns(fusionPatterns, controlFn);
-      (void)applyPatternsAndFoldGreedily(funcOp.getBody(),
-                                         std::move(fusionPatterns));
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(fusionPatterns))))
+        return signalPassFailure();
       return;
     }
 
     if (fuseGenericOpsControl) {
       RewritePatternSet fusionPatterns(context);
       linalg::populateElementwiseOpsFusionPatterns(fusionPatterns,
                                                    setFusedOpOperandLimit<4>);
 
-      (void)applyPatternsAndFoldGreedily(funcOp.getBody(),
-                                         std::move(fusionPatterns));
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(fusionPatterns))))
+        return signalPassFailure();
       return;
     }
 
     if (fuseWithReshapeByExpansion) {
       RewritePatternSet fusionPatterns(context);
       linalg::populateFoldReshapeOpsByExpansionPatterns(
           fusionPatterns, [](OpOperand * /*fusedOperand*/) { return true; });
       if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
if (controlFuseByExpansion) {
               }
               return false;
             }
             return true;
           };
 
       linalg::populateFoldReshapeOpsByExpansionPatterns(fusionPatterns,
                                                         controlReshapeFusionFn);
-      (void)applyPatternsAndFoldGreedily(funcOp.getBody(),
-                                         std::move(fusionPatterns));
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(fusionPatterns))))
+        return signalPassFailure();
       return;
     }
 
     if (fuseWithReshapeByCollapsing) {
       RewritePatternSet patterns(context);
       linalg::populateFoldReshapeOpsByCollapsingPatterns(
           patterns, [](OpOperand * /*fusedOperand */) { return true; });
-      (void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(patterns))))
+        return signalPassFailure();
+      return;
     }
 
     if (fuseWithReshapeByCollapsingWithControlFn) {
       RewritePatternSet patterns(context);
       linalg::ControlFusionFn controlFn = [](OpOperand *fusedOperand) -> bool {
         Operation *producer = fusedOperand->get().getDefiningOp();
         if (isa<tensor::ExpandShapeOp>(producer)) {
           // Skip fusing the first operand.
           return fusedOperand->getOperandNumber();
         }
         return true;
       };
       linalg::populateFoldReshapeOpsByCollapsingPatterns(patterns, controlFn);
-      (void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(patterns))))
+        return signalPassFailure();
+      return;
+    }
+
+    if (fuseMultiUseProducer) {
+      RewritePatternSet patterns(context);
+      patterns.insert<TestMultiUseProducerFusion>(context);
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(patterns))))
+        return signalPassFailure();
+      return;
     }
 
     if (!collapseDimensions.empty()) {
       SmallVector<int64_t, 2> dims(collapseDimensions.begin(),
                                    collapseDimensions.end());
       linalg::GetCollapsableDimensionsFn collapseFn =
           [&dims](linalg::GenericOp op) {
             SmallVector<ReassociationIndices> reassociations;
             reassociations.emplace_back(dims);
             return reassociations;
           };
       RewritePatternSet patterns(context);
       linalg::populateCollapseDimensions(patterns, collapseFn);
-      (void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));
+      if (failed(applyPatternsAndFoldGreedily(funcOp.getBody(),
+                                              std::move(patterns))))
+        return signalPassFailure();
+      return;
     }
   }
 };
 
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestLinalgElementwiseFusion() {
   PassRegistration<TestLinalgElementwiseFusion>();
 }
 } // namespace test
 } // namespace mlir