[mlir][linalg] Fold duplicate and unused inputs in linalg.generic

[springerm]

Summary

If an input bbArg is not used, its corresponding input operand is removed. If there are duplicate input operands or input operands that are also used as output operands, the duplicate input operands are removed. Output operands are never modified.

[nicolasvasilache]

Nice!

[nicolasvasilache]

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

401

for (auto [x, y] : replacements) ?

[springerm]

address comments

[mravishankar]

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

373

This seems to be a duplicate of what the deduplicate operands does. I am not sure why it is done separately and not added to the existing patterns. Unless I am missing something could we revert this patch (or do a follow up fix) and see why the existing patterns doesn't cover this case. It should.

[mravishankar]

Want some clarification on went this is done as a separate pattern

[springerm]

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

373

DeduplicateAndRemoveDeadOperandsAndResults removes duplicate inputs and duplicate outputs separately. A value that is both an input and an output is not considered a duplicate by this pattern. The way the pattern is written it seems difficult to support this.

We needed a specific way of de-duplicating operands: Outputs should not be modified (because they are bufferization DPS destinations), but inputs that are also outputs can be removed. (Removing the outputs instead would not be desirable.)

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

 
 /// Patterns to bubble up or down data layout ops across other operations.
 void populateDataLayoutPropagationPatterns(RewritePatternSet &patterns);
 
 /// Pattern to remove dead operands and results of `linalg.generic` operations.
 /// This is effectively DCE for a linalg op.
 void populateEraseUnusedOperandsAndResultsPatterns(RewritePatternSet &patterns);
 
+/// Patterns to promote inputs to outputs and remove unused inputs of
+/// `linalg.generic` ops.
+void populateEraseUnnecessaryInputsPatterns(RewritePatternSet &patterns);
+
 /// Function type to control generic op dimension collapsing. It is expected
 /// to return an array of `ReassociationIndices` representing dimensions that
 /// should be merged.
 using GetCollapsableDimensionsFn =
     std::function<SmallVector<ReassociationIndices>(linalg::GenericOp)>;
 
 /// Pattern to collapse dimensions in a linalg.generic op. This will collapse
 /// tensor operands when needed and expand back the result tensors.

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

static bool isResultValueDead(linalg::GenericOp genericOp, OpResult result) {
 
   return true;
 }
 
 namespace {
 
 struct DeduplicateAndRemoveDeadOperandsAndResults
     : public OpRewritePattern<GenericOp> {
-  using OpRewritePattern<GenericOp>::OpRewritePattern;
+  DeduplicateAndRemoveDeadOperandsAndResults(MLIRContext *ctx,
+                                             bool removeOutputs)
+      : OpRewritePattern<GenericOp>(ctx), removeOutputs(removeOutputs) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     // Create a map from argument position in the original op to the argument
     // position in the new op. If the argument is dropped it wont have an entry.
     SmallVector<OpOperand *> droppedOpOperands;
 
     // Information needed to build the new op.
for (const auto &result : llvm::enumerate(genericOp.getResults())) {
         continue;
       replacementsVals[result.index()] = newOp.getResult(it->second);
     }
     rewriter.replaceOp(genericOp, replacementsVals);
     return success();
   }
 
 private:
+  /// If unset, outputs are not modified by this pattern.
+  bool removeOutputs;
+
   // Deduplicate input operands, and return the
   // - Mapping from operand position in the original op, to operand position in
   // the canonicalized op.
   // - The preserved input operands list (by reference).
   llvm::SmallDenseMap<unsigned, unsigned>
   deduplicateInputOperands(GenericOp genericOp,
                            SmallVector<OpOperand *> &droppedOpOperands,
                            SmallVector<Value> &newInputOperands,
private:
   llvm::SmallDenseMap<unsigned, unsigned>
   deduplicateOutputOperands(GenericOp genericOp,
                             SmallVector<OpOperand *> &droppedOpOperands,
                             SmallVector<Value> &newOutputOperands,
                             SmallVector<AffineMap> &newIndexingMaps) const {
     llvm::SmallDenseMap<unsigned, unsigned> origToNewPos;
     llvm::SmallDenseMap<std::tuple<Value, AffineMap, Value>, unsigned>
         dedupedOutpts;
-    // If the op doesnt have tensor semantics, keep all the outputs as
-    // preserved.
-    if (!genericOp.hasTensorSemantics()) {
+    // If the op doesn't have tensor semantics or outputs should not be removed,
+    // keep all the outputs as preserved.
+    if (!genericOp.hasTensorSemantics() || !removeOutputs) {
       for (const auto &en : llvm::enumerate(genericOp.getDpsInitOperands())) {
         origToNewPos[en.index()] = newOutputOperands.size();
         newOutputOperands.push_back(en.value()->get());
         newIndexingMaps.push_back(genericOp.getMatchingIndexingMap(en.value()));
       }
       return origToNewPos;
     }
     // Output argument can be dropped if the result has
LogicalResult matchAndRewrite(GenericOp genericOp,
 
     if (hasRemovedCycles) {
       return success();
     }
 
     return failure();
   }
 };
+
+/// Fold uses of duplicate inputs in the body of a linalg.generic. E.g.:
+/// ```
+/// linalg.generic ins(%a, %b, %a, %b) outs(%a)
+/// ^bb0(%in0, %in1, %in2, %in3, %out1)
+/// ```
+/// Assuming that all %a and %b have the same index map:
+/// * All uses of %in0 and %in2 are replaced with %out1
+/// * All uses of %in1 are replaced with %in3
+/// This pattern can enable additional canonicalizations: In the above example,
+/// %in0, %in1 and %in3 have no uses anymore and their corresponding operands
+/// can be folded away. This pattern does not modify uses of output block args.
+struct FoldDuplicateInputBbArgs : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    // Find replacement bbArgs for all input bbArg.
+    DenseMap<int, int> replacements;
+    for (int i = 0; i < genericOp.getNumDpsInputs(); ++i) {
+      // Skip bbArgs that have no uses.
+      if (genericOp.getBody()->getArgument(i).getUses().empty())
+        continue;
+      // Find replacement bbArg. This can be an input or an output bbArg.
+      for (int j = genericOp->getNumOperands() - 1; j > i; --j) {
+        if (genericOp->getOperand(i) == genericOp->getOperand(j) &&
+            genericOp.getIndexingMapsArray()[i] ==
+                genericOp.getIndexingMapsArray()[j]) {
+          replacements[i] = j;
+          break;
+        }
+      }
+    }
+
+    // Stop here if no replacements were found.
+    if (replacements.empty())
+      return failure();
+
+    // Rewrite the op.
+    rewriter.updateRootInPlace(genericOp, [&]() {
+      for (const auto it : replacements) {
+        BlockArgument bbArg = genericOp.getBody()->getArgument(it.first);
+        BlockArgument replacement = genericOp.getBody()->getArgument(it.second);
+        rewriter.replaceAllUsesWith(bbArg, replacement);
+      }
+    });
+
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::linalg::populateEraseUnusedOperandsAndResultsPatterns(
     RewritePatternSet &patterns) {
-  patterns.insert<DeduplicateAndRemoveDeadOperandsAndResults,
-                  RemoveUnusedCycleInGenericOp>(patterns.getContext());
+  patterns.insert<DeduplicateAndRemoveDeadOperandsAndResults>(
+      patterns.getContext(), /*removeOutputs=*/true);
+  patterns.insert<RemoveUnusedCycleInGenericOp>(patterns.getContext());
+}
+
+void mlir::linalg::populateEraseUnnecessaryInputsPatterns(
+    RewritePatternSet &patterns) {
+  patterns.insert<DeduplicateAndRemoveDeadOperandsAndResults>(
+      patterns.getContext(), /*removeOutputs=*/false);
+  patterns.insert<FoldDuplicateInputBbArgs>(patterns.getContext());
 }
+

mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir

 // RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-erase-unused-operands-and-results | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -test-linalg-transform-patterns=test-erase-unnecessary-inputs | FileCheck %s --check-prefix=CHECK-INPUT
 
 // CHECK-LABEL: func @remove_deadargs_generic_basic
 //  CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> {
 //       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
 //  CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)
 //  CHECK-SAME: outs({{.*}} : tensor<?xf32>) {
 #map0 = affine_map<(d0) -> (d0)>
 func.func @remove_deadargs_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {

 //  CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
 //      CHECK: func @drop_only_the_cycles_not_used_by_others(
 // CHECK-SAME:     %[[ARG0:.+]]: tensor<?x?x?xf32>)
 //      CHECK:   %[[INIT:.+]] = tensor.empty
 //      CHECK:   %[[GENERIC:.+]]:2 = linalg.generic
 // CHECK-SAME:       indexing_maps = [#[[MAP1]], #[[MAP2]]]
 // CHECK-SAME:       outs(%[[ARG0]], %[[INIT]] :
 //      CHECK:   return %[[GENERIC]]#0
+
+
+// -----
+
+// CHECK-INPUT-LABEL: func @remove_unnecessary_input(
+//  CHECK-INPUT-SAME:     %[[a:.*]]: tensor<?xf32>, %[[b:.*]]: tensor<?xf32>
+#map = affine_map<(d0) -> (d0)>
+func.func @remove_unnecessary_input(%a: tensor<?xf32>, %b: tensor<?xf32>)
+    -> tensor<?xf32>
+{
+  //      CHECK-INPUT: %[[result:.*]] = linalg.generic {indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel"]}
+  // CHECK-INPUT-SAME:     ins(%[[a]] : tensor<?xf32>) outs(%[[b]] : tensor<?xf32>) {
+  //      CHECK-INPUT: ^bb0(%[[in:.*]]: f32, %[[out:.*]]: f32):
+  //      CHECK-INPUT:   %[[add:.*]] = arith.addf %[[in]], %[[out]]
+  //      CHECK-INPUT:   linalg.yield %[[add]]
+  //      CHECK-INPUT: } -> tensor<?xf32>
+  //      CHECK-INPUT: return %[[result]]
+  %0 = linalg.generic
+    {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]}
+    ins(%a, %b : tensor<?xf32>, tensor<?xf32>) outs(%b : tensor<?xf32>) {
+  ^bb0(%in: f32, %in_2: f32, %out: f32):
+    %16 = arith.addf %in, %in_2 : f32
+    linalg.yield %16 : f32
+  } -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}

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

Option<bool> testSwapExtractSliceWithFill{
       *this, "test-swap-extract-slice-with-fill-pattern",
       llvm::cl::desc(
           "Test patterns to swap tensor.extract_slice(linalg.fill())"),
       llvm::cl::init(false)};
   Option<bool> testEraseUnusedOperandsAndResults{
       *this, "test-erase-unused-operands-and-results",
       llvm::cl::desc("Test patterns to erase unused operands and results"),
       llvm::cl::init(false)};
+  Option<bool> testEraseUnnecessaryInputs{
+      *this, "test-erase-unnecessary-inputs",
+      llvm::cl::desc("Test patterns to erase unnecessary inputs"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applyPatterns(func::FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//

 }
 
 static void applyEraseUnusedOperandsAndResultsPatterns(func::FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   populateEraseUnusedOperandsAndResultsPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+static void applyEraseUnnecessaryInputs(func::FuncOp funcOp) {
+  RewritePatternSet patterns(funcOp.getContext());
+  populateEraseUnnecessaryInputsPatterns(patterns);
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnOperation() {
   if (testPatterns)
     return applyPatterns(getOperation());
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getOperation());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getOperation());
   if (testTransformPadTensor)
     return applyPadTensorToGenericPatterns(getOperation());
   if (testGeneralizePadTensor)
     return applyGeneralizePadTensorPatterns(getOperation());
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getOperation());
   if (testBubbleUpExtractSliceOpPattern)
     return applyBubbleUpExtractSliceOpPattern(getOperation());
   if (testSwapExtractSliceWithFill)
     return applySwapExtractSliceWithFillPattern(getOperation());
   if (testEraseUnusedOperandsAndResults)
     return applyEraseUnusedOperandsAndResultsPatterns(getOperation());
+  if (testEraseUnnecessaryInputs)
+    return applyEraseUnnecessaryInputs(getOperation());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir