Add linalg.fill canonicalization pattern for tensor.extract

[suryajasper]

Added the initial canonicalization pattern (mozphab issues with first submission)

[antiagainst]

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

639

Can we add some tests in test/Dialect/Linalg/canonicalize.mlir for this pattern?

647

No wrapping { } per LLVM coding style.

655

We don't need to perform extension here as fill then extract shouldn't change the element type anyway. Also, it's not guaranteed that the element types would be floating point.

[antiagainst]

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

637–638

We need to register this pattern to getCanonicalizationPatterns.

637–638

Could you capitalize each sentence and close with a .? Also for other comments.

mlir/test/Dialect/Linalg/canonicalize.mlir

341

Can we make this as an input to the function?

353

We need to check that we are returning the scalar value.

[antiagainst]

mlir/test/Dialect/Linalg/canonicalize.mlir

354–355

Have you tried to run the tests? I don't think we will see this in the transformed IR and so this won't pass?

354–355

For this line too.

[suryajasper]

Requesting review after addressing revision comments

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

637–638

It has been registered already. That change was in the previous diff, so it's included in the full expanded source here

[suryajasper]

Fixed FileCheck prompts for test mlir

[antiagainst]

Ah, I think you only uploaded the last commit?? Can we squash all your local commits into one and reupload it so I can see the full patch? Just make sure the commit message contains "Differential Revision: https://reviews.llvm.org/D156008" when you upload

[suryajasper]

Squashed prior commits for review

[suryajasper]

Squashed prior commits for review (fixed commit message)

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

LogicalResult matchAndRewrite(tensor::PadOp padOp,
     if (replacement.getType() != padOp.getResultType()) {
       replacement = rewriter.create<tensor::CastOp>(
           fillOp.getLoc(), padOp.getResultType(), replacement);
     }
     rewriter.replaceOp(padOp, replacement);
     return success();
   }
 };
 
 /// Fold tensor.insert_slice(tensor.pad(<input>), linalg.fill) into
 /// tensor.insert_slice(<input>, linalg.fill) if the padding value and the
 /// filling value are the same.
 struct FoldInsertPadIntoFill : public OpRewritePattern<tensor::InsertSliceOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
                                 PatternRewriter &rewriter) const override {
     auto srcPadOp = insertOp.getSource().getDefiningOp<tensor::PadOp>();
     if (!srcPadOp)
       return failure();
 
     if (insertOp.getType().getRank() != insertOp.getSourceType().getRank())
       return failure();
 
     // Walk back the tensor.insert_slice chain and find the first destination
     // value at the start of the chain.
     Value firstDest = insertOp.getDest();
     while (auto prevOp = firstDest.getDefiningOp<tensor::InsertSliceOp>()) {
       if (prevOp.getType().getRank() != prevOp.getSourceType().getRank())
         return failure();
 
       // Make sure the range of values accessed are disjoint. Without this, we
       // cannot fold tensor.pad away.
       bool disjoint = false;
       for (int i = 0, e = prevOp.getType().getRank(); i < e; ++i) {
LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
 
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
         insertOp, srcPadOp.getSource(), insertOp.getDest(), newOffsets,
         newSizes, insertOp.getMixedStrides());
     return success();
   }
 };
 
+/// Fold tensor.extract(linalg.fill(<input>)) into <input>
+struct FoldFillWithTensorExtract : public OpRewritePattern<tensor::ExtractOp> {
+public:
+  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    // See if tensor input of tensor.extract op is the result of a linalg.fill op.
+    auto fillOp = extractOp.getTensor().getDefiningOp<linalg::FillOp>();
+    if (!fillOp)
+      return failure();
+
+    // Get scalar input operand of linalg.fill op.
+    Value extractedScalar = fillOp.getInputs()[0];
+
+    // Replace tensor.extract op with scalar value used to fill the tensor.
+    rewriter.replaceOp(extractOp, extractedScalar);
+    return success();
+  }
+};
+
 } // namespace
 
 void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
   results
-      .add<FoldFillWithPad, FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
+      .add<FoldFillWithTensorExtract,
+           FoldFillWithPad, FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
            FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
            FoldInsertPadIntoFill>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // GenericOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Linalg/canonicalize.mlir

func.func @fold_fill_reshape_dynamic(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32> {
   // CHECK: %[[RESULT:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[RESHAPE]]
   %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]]
       : tensor<?x?x?x?x?xf32> into tensor<?x?xf32>
   // CHECK: return %[[RESULT]]
   return %1 : tensor<?x?xf32>
 }
 
 // -----
+//       CHECK: func @fold_fill_extract
+//  CHECK-SAME:   %[[ARG0:.+]]: i1
+func.func @fold_fill_extract(%arg0 : i1) -> i1 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  %empty_dynamic = tensor.empty(%c1) : tensor<1x2x3x?xi1>
+  %filled = linalg.fill ins(%arg0 : i1) outs(%empty_dynamic : tensor<1x2x3x?xi1>) -> tensor<1x2x3x?xi1>
+
+  %extracted = tensor.extract %filled[%c0, %c0, %c0, %c0] : tensor<1x2x3x?xi1>
+
+  //  CHECK:   return %[[ARG0]]
+  return %extracted : i1
+}
+
+// -----
 
 // CHECK: func @fold_self_copy
 func.func @fold_self_copy(%0 : memref<4x16xf32>) {
 // CHECK-NEXT: return
   linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                                    affine_map<(d0, d1) -> (d0, d1)>],
                   iterator_types = ["parallel", "parallel"]}
     ins(%0 : memref<4x16xf32>)
     outs(%0 : memref<4x16xf32>) {