[mlir][tensor] Add option to fold only tensor.empty with a single use

[springerm]

Summary

This is useful for transformations such as bufferization, which is looking for tensor.extract_slice/insert_slice pairs.

Also fix the documentation of the corresponding tranform op.

[springerm]

add test case

[springerm]

cleanup

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

def ApplyDropRedundantInsertSliceRankExpansionPatternsOp : Op<Transform_Dialect,
 
   let assemblyFormat = "attr-dict";
 }
 
 def ApplyFoldTensorEmptyPatternsOp : Op<Transform_Dialect,
     "apply_patterns.tensor.fold_tensor_empty",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
-    Indicates that reassociative reshapes (tensor.collapse_shape /
-    tensor.expand_shape) should be folded with inverse rank expansions / rank
-    reductions (via tensor.insert_slice / tensor.extract_slice).
+    Indicates that tensor.extract_slice and reassociative reshapes should be
+    folded into tensor.empty.
+
+    If `fold_single_use_only` is set to "true", only tensor.empty that have a
+    single use are folded.
   }];
 
+  let arguments = (ins DefaultValuedAttr<BoolAttr, "false">:$fold_single_use_only);
   let assemblyFormat = "attr-dict";
 }
 def ApplyFoldIntoPackAndUnpackPatternsOp : Op<Transform_Dialect,
     "apply_patterns.tensor.fold_into_pack_and_unpack",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Indicates that operations like tensor.pad and tensor.extract_slice should
     be folded into tensor.pack and tensor.unpack operations, respectively.

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

void populateDropRedundantInsertSliceRankExpansionPatterns(
     RewritePatternSet &patterns);
 
 /// Populates `patterns` with patterns that fold `tensor.expand_shape` and
 /// `tensor.collapse_shape` into other ops.
 void populateReassociativeReshapeFoldingPatterns(RewritePatternSet &patterns);
 
 /// Populates `patterns` with patterns that fold tensor.empty with
 /// tensor.[extract_slice|expand_shape|collapse_shape].
-void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns);
+///
+/// If `singleUseOnly` is set to "true", only tensor.empty ops with a single
+/// use are folded.
+void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns,
+                                     bool foldSingleUseOnly = false);
 
 /// Populates `patterns` with patterns that fold operations like `tensor.pad`
 /// and `tensor.extract_slice` into `tensor.pack` and `tensor.unpack` operations
 /// respectively.
 void populateFoldIntoPackAndUnpackPatterns(RewritePatternSet &patterns);
 
 //===----------------------------------------------------------------------===//
 // Transform helpers

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

 
 void transform::ApplyDropRedundantInsertSliceRankExpansionPatternsOp::
     populatePatterns(RewritePatternSet &patterns) {
   tensor::populateDropRedundantInsertSliceRankExpansionPatterns(patterns);
 }
 
 void transform::ApplyFoldTensorEmptyPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
-  tensor::populateFoldTensorEmptyPatterns(patterns);
+  tensor::populateFoldTensorEmptyPatterns(patterns, getFoldSingleUseOnly());
 }
 
 void transform::ApplyFoldIntoPackAndUnpackPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   tensor::populateFoldIntoPackAndUnpackPatterns(patterns);
 }
 
 void transform::ApplyFoldTensorSubsetOpsPatternsOp::populatePatterns(

mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp

 
 using namespace mlir;
 using namespace mlir::tensor;
 
 namespace {
 
 template <typename ReshapeOp>
 struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
-  using OpRewritePattern<ReshapeOp>::OpRewritePattern;
+  FoldEmptyTensorWithReshapeOp(MLIRContext *ctx, PatternBenefit benefit = 1,
+                               bool foldSingleUseOnly = false)
+      : OpRewritePattern<ReshapeOp>(ctx, benefit),
+        foldSingleUseOnly(foldSingleUseOnly) {}
 
   LogicalResult matchAndRewrite(ReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
-    if (!reshapeOp.getSrc().template getDefiningOp<EmptyOp>())
+    // Check for tensor.empty source.
+    auto emptyOp = reshapeOp.getSrc().template getDefiningOp<EmptyOp>();
+    if (!emptyOp)
       return failure();
+
+    // Check for single use.
+    if (foldSingleUseOnly && !llvm::hasSingleElement(emptyOp->getUses()))
+      return failure();
+
+    // Reify result shape.
     Location loc = reshapeOp.getLoc();
     ReifiedRankedShapedTypeDims resultShapes;
     if (failed(reifyResultShapes(rewriter, reshapeOp, resultShapes)) ||
         !llvm::hasSingleElement(resultShapes))
       return failure();
+
+    // Create new tensor.empty op.
     // TODO: Do not drop tensor type encoding.
     Value emptyTensor = rewriter.create<EmptyOp>(
         loc, resultShapes[0], reshapeOp.getResultType().getElementType());
     if (emptyTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), emptyTensor);
     } else {
       rewriter.replaceOp(reshapeOp, emptyTensor);
     }
     return success();
   }
+
+private:
+  bool foldSingleUseOnly = false;
 };
 
-/// `tensor.empty` does not define any tensor contents, so a slice of a
-/// `tensor.empty` can be canonicalized to a smaller `tensor.empty`.
+/// tensor.empty does not define any tensor contents, so a slice of a
+/// tensor.empty can be folded to a smaller tensor.empty.
 struct FoldEmptyTensorWithExtractSliceOp
     : public OpRewritePattern<ExtractSliceOp> {
-  using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
+  FoldEmptyTensorWithExtractSliceOp(MLIRContext *ctx,
+                                    PatternBenefit benefit = 1,
+                                    bool foldSingleUseOnly = false)
+      : OpRewritePattern<ExtractSliceOp>(ctx, benefit),
+        foldSingleUseOnly(foldSingleUseOnly) {}
 
   LogicalResult matchAndRewrite(ExtractSliceOp sliceOp,
                                 PatternRewriter &rewriter) const override {
-    if (!sliceOp.getSource().getDefiningOp<EmptyOp>())
+    // Check for tensor.empty source.
+    auto emptyOp = sliceOp.getSource().template getDefiningOp<EmptyOp>();
+    if (!emptyOp)
+      return failure();
+
+    // Check for single use.
+    if (foldSingleUseOnly && !llvm::hasSingleElement(emptyOp->getUses()))
       return failure();
 
-    // ExtractSliceOp may be rank-reducing; its dynamic sizes must be
-    // preserved as well as its result type.
+    // Create new tensor.empty op. tensor.extract_slice may be rank-reducing;
+    // its dynamic sizes must be preserved as well as its result type.
     auto tensorType = RankedTensorType::get(sliceOp.getType().getShape(),
                                             sliceOp.getType().getElementType(),
                                             sliceOp.getType().getEncoding());
     rewriter.replaceOpWithNewOp<EmptyOp>(sliceOp, tensorType,
                                          sliceOp.getSizes());
     return success();
   }
+
+private:
+  bool foldSingleUseOnly = false;
 };
 
 } // namespace
 
-void mlir::tensor::populateFoldTensorEmptyPatterns(
-    RewritePatternSet &patterns) {
+void mlir::tensor::populateFoldTensorEmptyPatterns(RewritePatternSet &patterns,
+                                                   bool foldSingleUseOnly) {
   patterns.add<FoldEmptyTensorWithExtractSliceOp,
                FoldEmptyTensorWithReshapeOp<tensor::ExpandShapeOp>,
                FoldEmptyTensorWithReshapeOp<tensor::CollapseShapeOp>>(
-      patterns.getContext());
+      patterns.getContext(), /*benefit=*/1, foldSingleUseOnly);
 }

mlir/test/Dialect/Tensor/fold-empty-op.mlir

-// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-empty-op-folding  %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -test-transform-dialect-interpreter %s | FileCheck %s
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !transform.any_op):
+  transform.apply_patterns to %module_op {
+    transform.apply_patterns.tensor.fold_tensor_empty
+  } : !transform.any_op
+}
+
+// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
+// CHECK: #[[$MAP2:.+]] = affine_map<()[s0] -> (s0 * 28)>
 
 func.func @empty_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
   %0 = tensor.empty(%arg0) : tensor<6x5x?xf32>
   %1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]]
       : tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
   return %1 : tensor<2x3x5x4x?x7xf32>
 }
-//      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
-//      CHECK: func @empty_reshape_expansion
+// CHECK-LABEL: func @empty_reshape_expansion
 // CHECK-SAME:     %[[ARG0:.+]]: index
 // CHECK:        %[[OLD_INIT:.+]] = tensor.empty(%{{.*}}) : tensor<6x5x?xf32>
 // CHECK-NEXT:   %[[DIM:.*]] = tensor.dim %[[OLD_INIT]]
-// CHECK-NEXT:   %[[D:.+]] = affine.apply #[[MAP]]()[%[[DIM]]]
+// CHECK-NEXT:   %[[D:.+]] = affine.apply #[[$MAP]]()[%[[DIM]]]
 // CHECK-NEXT:   %[[INIT:.+]] = tensor.empty(%[[D]])
 // CHECK-NEXT:   return %[[INIT]]
 
-// -----
-
 func.func @empty_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
   %0 = tensor.empty(%arg0) : tensor<2x3x5x4x?x7xf32>
   %1 = tensor.collapse_shape %0 [[0, 1], [2], [3, 4, 5]]
       : tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>
   return %1 : tensor<6x5x?xf32>
 }
-//      CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 * 28)>
-//      CHECK: func @empty_reshape_collapse
+// CHECK-LABEL: func @empty_reshape_collapse
 // CHECK-SAME:     %[[ARG0:.+]]: index
 // CHECK:        %[[OLD_INIT:.+]] = tensor.empty(%{{.*}}) : tensor<2x3x5x4x?x7xf32>
 // CHECK-NEXT:   %[[DIM:.*]] = tensor.dim %[[OLD_INIT]]
-// CHECK-NEXT:   %[[D:.+]] = affine.apply #[[MAP]]()[%[[DIM]]]
+// CHECK-NEXT:   %[[D:.+]] = affine.apply #[[$MAP2]]()[%[[DIM]]]
 // CHECK-NEXT:   %[[INIT:.+]] = tensor.empty(%[[D]])
 // CHECK-NEXT:   return %[[INIT]]
 
-// -----
-
 func.func @fold_empty_tensor_with_slice
   (%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>
 {
   %0 = tensor.empty(%arg0) : tensor<?x10x40xf32>
   %1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]
     : tensor<?x10x40xf32> to tensor<5x?x20xf32>
   return %1 : tensor<5x?x20xf32>
 }
-//      CHECK: func @fold_empty_tensor_with_slice
+// CHECK-LABEL: func @fold_empty_tensor_with_slice
 // CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
 //      CHECK:   %[[T0:.+]] = tensor.empty(%[[ARG1]])
 //      CHECK:   return %[[T0]]
 
-// -----
-
 // CHECK-LABEL: func @rank_reducing_empty_tensor_extract
 func.func @rank_reducing_empty_tensor_extract(%sz : index, %idx : index) -> tensor<2xf32> {
   // CHECK: tensor.empty() : tensor<2xf32>
   %a = tensor.empty(%sz) : tensor<?x2xf32>
 
   // CHECK-NOT: extract
   %r = tensor.extract_slice %a[%idx, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
   return %r: tensor<2xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !transform.any_op):
+  transform.apply_patterns to %module_op {
+    transform.apply_patterns.tensor.fold_tensor_empty
+        {fold_single_use_only = true}
+  } : !transform.any_op
+}
+
+func.func @double_use_of_tensor_empty(%arg0: index, %arg1: index)
+    -> (tensor<5x?x20xf32>, tensor<5x?x20xf32>)
+{
+  %0 = tensor.empty(%arg0) : tensor<?x10x40xf32>
+  %1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]
+    : tensor<?x10x40xf32> to tensor<5x?x20xf32>
+  %2 = tensor.extract_slice %0[1, 1, 1] [5, %arg1, 20] [1, 1, 1]
+    : tensor<?x10x40xf32> to tensor<5x?x20xf32>
+  return %1, %2 : tensor<5x?x20xf32>, tensor<5x?x20xf32>
+}
+// CHECK-LABEL: func @double_use_of_tensor_empty(
+//       CHECK:   tensor.empty{{.*}} : tensor<?x10x40xf32>
+//       CHECK:   tensor.extract_slice
+//       CHECK:   tensor.extract_slice