[mlir][linalg] Move linalg.fill folding into linalg.generic pattern to elementwise op fusion

[nirvedhmeshram]

Summary

Move linalg.fill folding into linalg.generic pattern from canonicalization to elementwise op fusion

[hanchung]

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

2217–2218

nit: adding a blank line before comments usually helps readability.

2233

[optional] I'd prefer not having this comment. The below line already describes it.

2287–2288

Can merge this into one line?

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

1

Why do we need canoincalize here?

I found that they are copied from canonicalize.mlir. So I assume that there are at least two patterns kicked in to get the result IR. Can we split tests to test two patterns separately?

[nirvedhmeshram]

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

1

Yes, so the other pattern it needs is this one
https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp#L859
And that pattern has its own tests here that are not being moved by this change
https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/canonicalize.mlir#L405-L443

Coming back to the tests and the pattern I am moving here, without the deadarg removing pattern it will have dead args and fill ops with no consumers. I can account for that in the pattern matching checks. However, it is also important to check if the combination works as that is the end goal here. Perhaps one solution is to add a separate test file for checking this combination. What do you think?

[hanchung]

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

1

In this case, we can put the output of this pattern to canonicalize.mlir and test it there. I don't have strong opinion about having canonicalize in this file or not. I was afraid that we're not testing what the pattern does. I was also afraid if the test is robust, i.e., if adding other patterns would break this.

After reading patterns deeply, I think it's a fine change to me.

1

I'll let you make the call.

[nirvedhmeshram]

Addressing reviewer comments

[nirvedhmeshram]

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

1

Sounds good, I think we can leave the canonicalize in then.

[mravishankar]

Minor change to not use canonicalize in lit test. Otherwise LGTM! THanks!

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

1

I think MLIR convention is to not rely on canonicalize to clean up lit tests. Its pretty disruptive cause changes to canonicalize then hits many lit tests. So It is probably better to remove the canonicalize and test the pass by itself....

[nirvedhmeshram]

Addressing reviewer comments and rebase

[nirvedhmeshram]

There was a bug that was found by @benvanik in the folding pattern where it would trigger infinitely if the fill already did not have a use in the generic op.
Fixing that in this revision. With this change also dead args get dropped automatically with the transform pattern because it calls canonicalization within the pass. The fact that this was not happening initially was due to this bug.

[mravishankar]

Minot nits, but thanks for moving it!

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

2231

Nit: LLVM prefers early exit style.

So

auto fillOp = ...
if (!fillOp) continue;
``

2232

You could just use

if (!genericOp.payloadUsesValueFromOperand(opOperand))
  continue;

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

for (NamedAttribute kv : genericOp->getAttrs()) {
       }
     }
     rewriter.inlineRegionBefore(genericOp.region(), newOp.region(),
                                 newOp.region().begin());
     rewriter.replaceOp(genericOp, newOp->getResults());
     return success();
   }
 };
-
-/// Fold linalg.fill into linalg.generic
-struct FoldFillWithGenericOp : public OpRewritePattern<GenericOp> {
-  using OpRewritePattern<GenericOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(GenericOp genericOp,
-                                PatternRewriter &rewriter) const override {
-    if (!genericOp.hasTensorSemantics())
-      return failure();
-    bool fillFound = false;
-    Block &payload = genericOp.region().front();
-    for (OpOperand *opOperand : genericOp.getInputOperands()) {
-      FillOp fillOp = opOperand->get().getDefiningOp<FillOp>();
-      if (fillOp) {
-        fillFound = true;
-        payload.getArgument(opOperand->getOperandNumber())
-            .replaceAllUsesWith(fillOp.value());
-      }
-    }
-    // fail if there are no FillOps to fold.
-    return success(fillFound);
-  }
-};
 } // namespace
 
 void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results.add<DeduplicateGenericOpInputs, EraseIdentityGenericOp,
-              DeadArgsGenericOpInputs, FoldFillWithGenericOp>(context);
+              DeadArgsGenericOpInputs>(context);
 }
 
 LogicalResult GenericOp::fold(ArrayRef<Attribute>,
                               SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//

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

LogicalResult matchAndRewrite(GenericOp op,
     }
     if (!modifiedOutput) {
       rewriter.cancelRootUpdate(op);
       return failure();
     }
     rewriter.finalizeRootUpdate(op);
     return success();
   }
 };
-} // namespace
 
+/// Fold linalg.fill into linalg.generic
+struct FoldFillWithGenericOp : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    if (!genericOp.hasTensorSemantics())
+      return failure();
+    bool fillFound = false;
+    Block &payload = genericOp.region().front();
+    for (OpOperand *opOperand : genericOp.getInputOperands()) {
+      if (!genericOp.payloadUsesValueFromOperand(opOperand))
+        continue;
+      FillOp fillOp = opOperand->get().getDefiningOp<FillOp>();
+      if (!fillOp)
+        continue;
+      fillFound = true;
+      payload.getArgument(opOperand->getOperandNumber())
+          .replaceAllUsesWith(fillOp.value());
+    }
+    return success(fillFound);
+  }
+};
+} // namespace
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 
 void mlir::linalg::populateFoldReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns) {
   patterns.add<
       FoldProducerReshapeOpByLinearization<false, tensor::CollapseShapeOp>,

 }
 
 void mlir::linalg::populateElementwiseOpsFusionPatterns(
     RewritePatternSet &patterns, LinalgElementwiseFusionOptions options) {
   auto *context = patterns.getContext();
   patterns.add<FuseElementwiseOps, FoldScalarOrSplatConstant,
                FoldConstantTranspose>(context,
                                       options.controlElementwiseOpsFusionFn);
-  patterns.add<RemoveOutsDependency>(context);
+  patterns.add<RemoveOutsDependency, FoldFillWithGenericOp>(context);
   populateSparseTensorRewriting(patterns);
   populateFoldReshapeOpsByExpansionPatterns(patterns,
                                             options.controlFoldingReshapesFn);
   AffineApplyOp::getCanonicalizationPatterns(patterns, context);
   GenericOp::getCanonicalizationPatterns(patterns, context);
   tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
   tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
   context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
       patterns);

mlir/test/Dialect/Linalg/canonicalize.mlir

//   CHECK-NOT: memref.copy
   memref.copy %arg0, %arg0 : memref<2x3x?x4xf32> to memref<2x3x?x4xf32>
 
 //   CHECK: return
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @fold_fill_generic_basic
-//  CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> { 
-//   CHECK-NOT: linalg.fill
-//       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
-//  CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)
-//  CHECK-SAME: outs({{.*}} : tensor<?xf32>) {
-#map0 = affine_map<(d0) -> (d0)>
-func @fold_fill_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
-  %c0 = arith.constant 0 : index
-  %cst = arith.constant 7.0 : f32
-  %0 = tensor.dim %arg0, %c0 : tensor<?xf32>
-  %1 = linalg.init_tensor [%0] : tensor<?xf32>
-  %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf32>) -> tensor<?xf32>
-  %3 = linalg.init_tensor [%0] : tensor<?xf32>
-  %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %2 : tensor<?xf32>, tensor<?xf32>) outs (%3:tensor<?xf32>) {
-  ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
-    %5 = arith.addf  %arg1, %arg2 : f32
-  linalg.yield %5 : f32
-  } -> tensor<?xf32>
-  return %4 : tensor<?xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @fold_fill_generic_mixedaccess
-//   CHECK-NOT: linalg.fill
-//       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
-//   CHECK-NOT: ins
-//  CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) {
-#map0 = affine_map<(d0, d1) -> (d0, d1)>
-#map1 = affine_map<(d0, d1) -> (d1, d0)>
-func @fold_fill_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) {
-  %c0 = arith.constant 0 : index
-  %c1 = arith.constant 0 : index
-  %cst1 = arith.constant 7.0 : f32
-  %cst2 = arith.constant 6.0 : f32
-  %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
-  %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
-  %2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32>
-  %3 = linalg.fill ins(%cst1 : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
-  %4 = linalg.init_tensor [%1, %0] : tensor<?x?xf32>
-  %5 = linalg.fill ins(%cst2 : f32) outs(%4 : tensor<?x?xf32>) -> tensor<?x?xf32>
-  %6 = linalg.init_tensor [%0, %1] : tensor<?x?xf32>
-  %7 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%3, %5 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%6:tensor<?x?xf32>) {
-  ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
-    %8 = arith.divf  %arg1, %arg2 : f32
-  linalg.yield %8 : f32
-  } -> tensor<?x?xf32>
-  return %7 : tensor<?x?xf32>
-}
-
-// -----
-
 // 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 @remove_deadargs_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
   %c0 = arith.constant 0 : index

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

 // RUN: mlir-opt %s -linalg-fuse-elementwise-ops -split-input-file | FileCheck %s
 
 // CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 
 // CHECK-LABEL: @add_mul_fusion
 func @add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
 {
   %c0 = arith.constant 0 : index
%2 = linalg.generic {
           linalg.yield %23 : i32
         } -> tensor<5000xi32>
   return %2 : tensor<5000xi32>
 }
 // CHECK-LABEL: func @illegal_fusion(
 //       CHECK:   %[[PRODUCER:.+]] = linalg.generic
 //       CHECK:    linalg.generic
 //   CHECK-SAME:       ins(%[[PRODUCER]]
+
+// -----
+
+// CHECK-LABEL: func @fold_fill_generic_basic
+//  CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> { 
+//   CHECK-NOT: linalg.fill
+//       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
+//  CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)
+//  CHECK-SAME: outs({{.*}} : tensor<?xf32>) {
+#map0 = affine_map<(d0) -> (d0)>
+func @fold_fill_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 7.0 : f32
+  %0 = tensor.dim %arg0, %c0 : tensor<?xf32>
+  %1 = linalg.init_tensor [%0] : tensor<?xf32>
+  %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf32>) -> tensor<?xf32>
+  %3 = linalg.init_tensor [%0] : tensor<?xf32>
+  %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %2 : tensor<?xf32>, tensor<?xf32>) outs (%3:tensor<?xf32>) {
+  ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
+    %5 = arith.addf  %arg1, %arg2 : f32
+  linalg.yield %5 : f32
+  } -> tensor<?xf32>
+  return %4 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_fill_generic_mixedaccess
+//   CHECK-NOT: linalg.fill
+//       CHECK: %[[GENERIC_OP:.*]] = linalg.generic
+//   CHECK-NOT: ins
+//  CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) {
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1, d0)>
+func @fold_fill_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 0 : index
+  %cst1 = arith.constant 7.0 : f32
+  %cst2 = arith.constant 6.0 : f32
+  %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
+  %2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32>
+  %3 = linalg.fill ins(%cst1 : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %4 = linalg.init_tensor [%1, %0] : tensor<?x?xf32>
+  %5 = linalg.fill ins(%cst2 : f32) outs(%4 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %6 = linalg.init_tensor [%0, %1] : tensor<?x?xf32>
+  %7 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%3, %5 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%6:tensor<?x?xf32>) {
+  ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
+    %8 = arith.divf  %arg1, %arg2 : f32
+  linalg.yield %8 : f32
+  } -> tensor<?x?xf32>
+  return %7 : tensor<?x?xf32>
+}