[mlir][linalg] Fold init_tensor -> fill -> tensor_reshape chain

[antiagainst]

Summary

For such op chains, we can create new linalg.init_tensor ops
with the result type of the linalg.tensor_reshape op.

Depends On D99115

[mravishankar]

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

1673

I dont think you need to check that it is an InitTensor op. Its just needed for the shape. You can just take the shape of the outs operand of the FillOp.

[antiagainst]

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

1673

Not that this pattern will change the InitTensor op itself; so it need to check that.

[mravishankar]

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

1673

You just need to create a new InitTensor op. The old one will become dead. Stated differently, whether fill.output() is an InitTensorOp or not (it is always by construction today) is not needed for this pattern.

[antiagainst]

Address comments

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

1673

Ah, I see your points now. Changed, thanks.

[mravishankar]

LGTM. Just a few renaming nits.

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

1662

Nit: Fix stale comment

1664

Just say new FillOp

1666

Nit: Rename to FoldFillWithTensorReshape

mlir/test/Dialect/Linalg/canonicalize.mlir

808

super nit : rename to func @fold_fill_tensor_reshape()

[nicolasvasilache]

What is the benefit you see in duplicating the init_tensor?
Bufferization of reshape is still not in a good place but offhand, if the original init has multiple consumers, this will strictly increase the amount of memory.

I think I would have just swapped the reshape and the fill and punt decision about inplace vs alloc to later as I don't have enough data.

Maybe you already have a use case in mind?

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

LogicalResult matchAndRewrite(memref::DimOp dimOp,
     rewriter.replaceOp(dimOp,
                        getReshapeOutputDimFromInputShape(
                            rewriter, dimOp.getLoc(), *dimIndex, reshapeOp.src(),
                            reshapeOp.getResultType().getShape(),
                            reshapeOp.getReassociationMaps()));
     return success();
   }
 };
+
+/// Fold linalg.init_tensor -> linalg.fill -> linalg.tensor_reshape chain.
+///
+/// For such op chains, we can create new linalg.init_tensor ops with the result
+/// type of the linalg.tensor_reshape op.
+struct FoldInitTensorFillReshape : OpRewritePattern<TensorReshapeOp> {
+  using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
+                                PatternRewriter &rewriter) const override {
+    auto oldFill = reshapeOp.src().getDefiningOp<FillOp>();
+    if (!oldFill)
+      return failure();
+
+    auto newInit = rewriter.create<InitTensorOp>(
+        oldFill.getLoc(), reshapeOp.getResultType().getShape(),
+        reshapeOp.getResultType().getElementType());
+    rewriter.replaceOpWithNewOp<FillOp>(reshapeOp, newInit, oldFill.value());
+
+    return success();
+  }
+};
 } // namespace
 
 void TensorReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
-  results.add<CollapseReshapeOps<TensorReshapeOp>, FoldReshapeWithConstant,
-              ReplaceDimOfReshapeOpResult>(context);
+  results.add<CollapseReshapeOps<TensorReshapeOp>, FoldInitTensorFillReshape,
+              FoldReshapeWithConstant, ReplaceDimOfReshapeOpResult>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, linalg::YieldOp op) {
   p << op.getOperationName();

mlir/test/Dialect/Linalg/canonicalize.mlir

 func @self_copy(%arg0 : memref<2x3x?x4xf32>) {
 
 //   CHECK-NOT: linalg.copy
   linalg.copy(%arg0, %arg0): memref<2x3x?x4xf32>, memref<2x3x?x4xf32>
 
 //   CHECK: return
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @fold_init_fill_reshape()
+func @fold_init_fill_reshape() -> tensor<6x4xf32> {
+  %zero = constant 0.0 : f32
+  // CHECK: %[[INIT:.+]] = linalg.init_tensor [6, 4] : tensor<6x4xf32>
+  %init = linalg.init_tensor [1, 2, 3, 4] : tensor<1x2x3x4xf32>
+  // CHECK: %[[FILL:.+]] = linalg.fill(%[[INIT]], %cst) : tensor<6x4xf32>, f32 -> tensor<6x4xf32>
+  %fill = linalg.fill(%init, %zero) : tensor<1x2x3x4xf32>, f32 -> tensor<1x2x3x4xf32>
+  %reshape = linalg.tensor_reshape %fill [
+    affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>,
+    affine_map<(d0, d1, d2, d3) -> (d3)>] : tensor<1x2x3x4xf32> into tensor<6x4xf32>
+  // CHECK: return %[[FILL]] : tensor<6x4xf32>
+  return %reshape : tensor<6x4xf32>
+}

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

 
 
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>
 
 //       CHECK: func @fold_unit_dim_for_init_tensor
 
+
 //       CHECK: %[[INPUT_RESHAPE:.+]] = linalg.tensor_reshape %{{.+}} [#[[MAP0]]] : tensor<1x1000xf32> into tensor<1000xf32>
-//       CHECK: %[[INIT_RESHAPE:.+]] = linalg.tensor_reshape %{{.+}} [] : tensor<1xf32> into tensor<f32>
+//       CHECK: %[[INIT:.+]] = linalg.init_tensor [] : tensor<f32>
+//       CHECK: %[[FILL:.+]] = linalg.fill(%[[INIT]], %cst) : tensor<f32>, f32 -> tensor<f32>
 //       CHECK: %[[GENERIC:.+]] = linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP1]], #[[MAP2]]]
 //  CHECK-SAME:     iterator_types = ["reduction"]
 //  CHECK-SAME:   ins(%[[INPUT_RESHAPE]] : tensor<1000xf32>)
-//  CHECK-SAME:   outs(%[[INIT_RESHAPE]] : tensor<f32>)
+//  CHECK-SAME:   outs(%[[FILL]] : tensor<f32>)
 //       CHECK: %[[GENERIC_RESHAPE:.+]] = linalg.tensor_reshape %[[GENERIC]] [] : tensor<f32> into tensor<1xf32>
 //       CHECK: return %[[GENERIC_RESHAPE:.+]] : tensor<1xf32>