[mlir][Linalg] Avoid unnecessary propagating producer result to fused op result.

[mravishankar]

Summary

Elementwise op fusion conserves the result of the producer in the
fused op, relying on later clean up patterns to drop unused results of
the fused op. Instead, if the producer result has no other use apart
from the consumer op, avoid making the producer result available in
the fused node. This saves some unnecessary IR manipulations.

[hanchung]

I was wondering why there are no tests, and found that corresponding lit tests failed. The code itself looks good to me, please also fix the lit test. Thanks!

[mravishankar]

Rebase and add lit test.

[mravishankar]

In D138096#3934433, @hanchung wrote:

I was wondering why there are no tests, and found that corresponding lit tests failed. The code itself looks good to me, please also fix the lit test. Thanks!

Actually the lit test failure is just mid-air collision. Should be fixed now. With the pattern to remove unused results this change is mostly NFC. Added a test anyway.

[hw-yuexy]

Hi Mahesh, I have a question regarding Linalg fusion in general.

It's common that user wants to fuse a matmul with a unary elementwise op (such as relu). Since matmul usually have 2 parallel + 1 reduction iterator, while the elementwise op usually only have 2 parallel iterators, it's hard to combine them into a single linalg.generic op like the examples in fusion-elementwise-ops.mlir.

So I'm wondering how is this type of fusion supported (or planned) currently? I'm trying to fuse a matmul op with a relu (contained in a generic op), both have memref input and output, so It will be really good if there are some examples of this type of fusion in the repo.

Thank you!

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

 bool linalg::detail::canOpOperandsBeDroppedImpl(
     linalg::LinalgOp linalgOp, ArrayRef<OpOperand *> droppedOperands) {
   SmallVector<AffineMap> indexingMaps;
   for (auto &opOperand : linalgOp->getOpOperands()) {
     if (llvm::is_contained(droppedOperands, &opOperand))
       continue;
     indexingMaps.push_back(linalgOp.getMatchingIndexingMap(&opOperand));
   }
+  if (indexingMaps.empty()) {
+    // If there are no indexing maps, the operand can only be dropped
+    // if the op has no loops.
+    return linalgOp.getNumLoops() == 0;
+  }
   return inversePermutation(concatAffineMaps(indexingMaps)) != AffineMap();
 }
 
 //===----------------------------------------------------------------------===//
 // ContractionOpInterface implementation
 //===----------------------------------------------------------------------===//
 
 /// Return true if the use-def chain from `v` to `from` consists of 0 or more

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

if (!coveredDims.all())
       return false;
   }
 
   return true;
 }
 
 /// Generate the region of the fused tensor operation. The region of the fused
 /// op must be empty.
-static void
-generateFusedElementwiseOpRegion(RewriterBase &rewriter, GenericOp fusedOp,
-                                 AffineMap consumerToProducerLoopsMap,
-                                 OpOperand *fusedOperand, unsigned nloops) {
+static void generateFusedElementwiseOpRegion(
+    RewriterBase &rewriter, GenericOp fusedOp,
+    AffineMap consumerToProducerLoopsMap, OpOperand *fusedOperand,
+    unsigned nloops, llvm::SmallDenseSet<int> &preservedProducerResults) {
   auto producer = cast<GenericOp>(fusedOperand->get().getDefiningOp());
   auto consumer = cast<GenericOp>(fusedOperand->getOwner());
   // Build the region of the fused op.
   Block &producerBlock = producer->getRegion(0).front();
   Block &consumerBlock = consumer->getRegion(0).front();
   Block *fusedBlock = new Block();
   fusedOp.getRegion().push_back(fusedBlock);
   BlockAndValueMapping mapper;
static void generateFusedElementwiseOpRegion(
   // 5. Remaining consumer's input operands (drop past index `consumerIdx`).
   for (BlockArgument bbArg :
        consumerBlock.getArguments()
            .take_front(consumer.getNumDpsInputs())
            .drop_front(fusedOperand->getOperandNumber() + 1))
     mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
 
   // 6. All of the producer's output operands
-  for (BlockArgument bbArg :
-       producerBlock.getArguments().take_back(producer.getNumDpsInits()))
-    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
+  for (auto bbArg : llvm::enumerate(
+           producerBlock.getArguments().take_back(producer.getNumDpsInits()))) {
+    if (!preservedProducerResults.count(bbArg.index()))
+      continue;
+    mapper.map(bbArg.value(), fusedBlock->addArgument(bbArg.value().getType(),
+                                                      bbArg.value().getLoc()));
+  }
 
   // 7. All of consumer's output operands.
   for (BlockArgument bbArg :
        consumerBlock.getArguments().take_back(consumer.getNumDpsInits()))
     mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
 
   // 8. Clone all producer operations except for the yield and index operations
   // to the fused operation.
for (auto &op : consumerBlock.without_terminator())
     rewriter.clone(op, mapper);
 
   // 11. Include the final yield (which is the remapped values for all the
   // yield)
   auto consumerYieldOp = cast<linalg::YieldOp>(consumerBlock.getTerminator());
   SmallVector<Value> fusedYieldValues;
   fusedYieldValues.reserve(producerYieldOp.getNumOperands() +
                            consumerYieldOp.getNumOperands());
-  for (auto producerYieldVal : producerYieldOp.getOperands())
-    fusedYieldValues.push_back(mapper.lookupOrDefault(producerYieldVal));
+  for (auto producerYieldVal : llvm::enumerate(producerYieldOp.getOperands())) {
+    if (preservedProducerResults.count(producerYieldVal.index()))
+      fusedYieldValues.push_back(
+          mapper.lookupOrDefault(producerYieldVal.value()));
+  }
   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 *>
 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(consumer.isDpsInput(fusedOperand) &&
          "expected producer of input operand");
+  /// Find the results of the producer that have uses outside of the consumer.
+  llvm::SmallDenseSet<int> preservedProducerResults;
+  for (auto producerResult : llvm::enumerate(producer->getResults())) {
+    auto outputOperand = producer.getDpsInitOperand(producerResult.index());
+    if (producer.payloadUsesValueFromOperand(outputOperand) ||
+        !producer.canOpOperandsBeDropped(outputOperand) ||
+        llvm::any_of(producerResult.value().getUsers(), [&](Operation *user) {
+          return user != consumer.getOperation();
+        })) {
+      preservedProducerResults.insert(producerResult.index());
+    }
+  }
 
   // Compute the fused operands list and indexing maps.
   SmallVector<Value> fusedInputOperands, fusedOutputOperands;
   SmallVector<Type> fusedResultTypes;
   SmallVector<AffineMap> fusedIndexMaps;
   fusedInputOperands.reserve(producer.getNumDpsInputs() +
                              consumer.getNumDpsInputs());
-  fusedOutputOperands.reserve(producer.getNumDpsInits() +
+  fusedOutputOperands.reserve(preservedProducerResults.size() +
                               consumer.getNumDpsInits());
-  fusedResultTypes.reserve(producer.getNumDpsInits() +
+  fusedResultTypes.reserve(preservedProducerResults.size() +
                            consumer.getNumDpsInits());
   fusedIndexMaps.reserve(producer->getNumOperands() +
                          consumer->getNumOperands());
   // In the following, numbering matches that of `generateFusedTensorOpRegion`.
   // 3. Consumer input operands/maps up to consumerIdx (exclusive).
   auto consumerInputs = consumer.getDpsInputOperands();
   auto *it = llvm::find_if(consumerInputs, [&](OpOperand *operand) {
     return operand == fusedOperand;
mlir::linalg::fuseElementwiseOps(RewriterBase &rewriter,
   // `consumerIdx`).
   for (OpOperand *opOperand :
        llvm::make_range(std::next(it), consumerInputs.end())) {
     fusedInputOperands.push_back(opOperand->get());
     fusedIndexMaps.push_back(consumer.getMatchingIndexingMap(opOperand));
   }
 
   // 6. Collect all of the producer outputs.
-  for (OpOperand *opOperand : producer.getDpsInitOperands()) {
-    fusedOutputOperands.push_back(opOperand->get());
+  for (auto opOperand : llvm::enumerate(producer.getDpsInitOperands())) {
+    if (!preservedProducerResults.count(opOperand.index()))
+      continue;
+
+    fusedOutputOperands.push_back(opOperand.value()->get());
     AffineMap map = getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp(
-        opOperand, producerResultIndexMap,
+        opOperand.value(), producerResultIndexMap,
         consumer.getMatchingIndexingMap(fusedOperand));
     fusedIndexMaps.push_back(map);
-    fusedResultTypes.push_back(opOperand->get().getType());
+    fusedResultTypes.push_back(opOperand.value()->get().getType());
   }
 
   // 7. All of consumer's output operands (skip operands: added by the builder).
   for (OpOperand *opOperand : consumer.getDpsInitOperands()) {
     fusedOutputOperands.push_back(opOperand->get());
     fusedIndexMaps.push_back(consumer.getMatchingIndexingMap(opOperand));
     fusedResultTypes.push_back(opOperand->get().getType());
   }
mlir::linalg::fuseElementwiseOps(RewriterBase &rewriter,
   AffineMap invProducerResultIndexMap =
       inversePermutation(producerResultIndexMap);
   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());
+  generateFusedElementwiseOpRegion(
+      rewriter, fusedOp, consumerToProducerLoopsMap, fusedOperand,
+      consumer.getNumLoops(), preservedProducerResults);
   return fusedOp.getOperation();
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
 public:
   FuseElementwiseOps(MLIRContext *context, ControlFusionFn fun,

mlir/test/Dialect/Linalg/fusion-elementwise-options.mlir

-// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-generic-ops -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-generic-ops-control -split-input-file | FileCheck %s
 
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #binary2Dpointwise = {
   indexing_maps = [#map0, #map0, #map0],
   iterator_types = ["parallel", "parallel"]
 }
 #ternary2Dpointwise = {
   indexing_maps = [#map0, #map0, #map0, #map0],

mlir/test/Dialect/Linalg/fusion-elementwise.mlir

+// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-generic-ops-control -split-input-file | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @drop_unused_producer_result(%arg0 : tensor<?x?xf32>,
+    %arg1 : 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(%arg0, %arg0  : tensor<?x?xf32>, tensor<?x?xf32>) {
+    ^bb0(%b0: f32, %b1: f32, %b2: f32):
+      %1 = arith.addf %b0, %b0 : f32
+      %2 = arith.mulf %b0, %b0 : f32
+      linalg.yield %1, %2 : f32, f32
+    } -> (tensor<?x?xf32>, tensor<?x?xf32>)
+  %3 = linalg.generic {
+      indexing_maps = [#map, #map, #map],
+      iterator_types = ["parallel", "parallel"]}
+      ins(%0#0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%arg0  : tensor<?x?xf32>) {
+    ^bb0(%b0: f32, %b1: f32, %b2: f32):
+      %4 = arith.subf %b0, %b1 : f32
+      linalg.yield %4 : f32
+    } -> tensor<?x?xf32>
+  return %3 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func @drop_unused_producer_result
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+//       CHECK:   %[[FUSED_OP:.+]] = linalg.generic
+//  CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]] :
+//       CHECK:   return %[[FUSED_OP]]

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

StringRef getDescription() const final {
     return "Test Linalg element wise operation fusion patterns";
   }
 
   Option<bool> fuseGenericOps{
       *this, "fuse-generic-ops",
       llvm::cl::desc("Test fusion of generic operations."),
       llvm::cl::init(false)};
 
+  Option<bool> fuseGenericOpsControl{
+      *this, "fuse-generic-ops-control",
+      llvm::cl::desc(
+          "Test fusion of generic operations with a control function."),
+      llvm::cl::init(false)};
+
   Option<bool> fuseWithReshapeByExpansion{
       *this, "fuse-with-reshape-by-expansion",
       llvm::cl::desc(
           "Test fusion of generic operations with reshape by expansion"),
       llvm::cl::init(false)};
 
   Option<bool> controlFuseByExpansion{
       *this, "control-fusion-by-expansion",
ListOption<int64_t> collapseDimensions{
       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));
+      return;
+    }
+
+    if (fuseGenericOpsControl) {
+      RewritePatternSet fusionPatterns(context);
       linalg::populateElementwiseOpsFusionPatterns(fusionPatterns,
                                                    setFusedOpOperandLimit<4>);
 
       (void)applyPatternsAndFoldGreedily(funcOp.getBody(),
                                          std::move(fusionPatterns));
       return;
     }