[Mlir-commits] [mlir] 40052b0 - [mlir][tensor] Add option to fold only tensor.empty with a single use
Matthias Springer
llvmlistbot at llvm.org
Fri Jun 9 03:37:08 PDT 2023
Author: Matthias Springer
Date: 2023-06-09T12:36:55+02:00
New Revision: 40052b08de6f73d7fb9b5e53c32ee21ba4d35fbc
URL: https://github.com/llvm/llvm-project/commit/40052b08de6f73d7fb9b5e53c32ee21ba4d35fbc
DIFF: https://github.com/llvm/llvm-project/commit/40052b08de6f73d7fb9b5e53c32ee21ba4d35fbc.diff
LOG: [mlir][tensor] Add option to fold only tensor.empty with a single use
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.
Differential Revision: https://reviews.llvm.org/D152455
Added:
Modified:
mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td
mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
mlir/test/Dialect/Tensor/fold-empty-op.mlir
mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td b/mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td
index 22367c98fdc72..9a569beaeb880 100644
--- a/mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td
@@ -31,11 +31,14 @@ 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,
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
index cadf3ad55f6a1..474c559a75220 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h
@@ -56,7 +56,11 @@ 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
diff --git a/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp b/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
index 18517cfa93655..8c85d18ada00d 100644
--- a/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
+++ b/mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp
@@ -90,7 +90,7 @@ void transform::ApplyDropRedundantInsertSliceRankExpansionPatternsOp::
void transform::ApplyFoldTensorEmptyPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
- tensor::populateFoldTensorEmptyPatterns(patterns);
+ tensor::populateFoldTensorEmptyPatterns(patterns, getFoldSingleUseOnly());
}
void transform::ApplyFoldIntoPackAndUnpackPatternsOp::populatePatterns(
diff --git a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
index 99679bc9b1378..7a707e749e69b 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/EmptyOpPatterns.cpp
@@ -18,17 +18,30 @@ 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());
@@ -40,21 +53,34 @@ struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
}
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());
@@ -62,14 +88,17 @@ struct FoldEmptyTensorWithExtractSliceOp
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);
}
diff --git a/mlir/test/Dialect/Tensor/fold-empty-op.mlir b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
index 799c69154ce3b..c3461d250e8ba 100644
--- a/mlir/test/Dialect/Tensor/fold-empty-op.mlir
+++ b/mlir/test/Dialect/Tensor/fold-empty-op.mlir
@@ -1,4 +1,14 @@
-// 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>
@@ -6,34 +16,28 @@ func.func @empty_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
: 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>
{
@@ -42,14 +46,12 @@ func.func @fold_empty_tensor_with_slice
: 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>
@@ -59,3 +61,28 @@ func.func @rank_reducing_empty_tensor_extract(%sz : index, %idx : index) -> tens
%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
diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
index c6aad2c8fdb35..baa19a980146e 100644
--- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
@@ -72,10 +72,6 @@ struct TestTensorTransforms
llvm::cl::desc("Test folding of expand_shape/collapse_shape"),
llvm::cl::init(false)};
- Option<bool> testEmptyOpFolding{
- *this, "test-empty-op-folding",
- llvm::cl::desc("Test folding of tensor.empty"), llvm::cl::init(false)};
-
Option<bool> testFoldIntoPackAndUnpack{
*this, "test-fold-into-pack-and-unpack",
llvm::cl::desc("Test folding ops into tensor.pack and tensor.unpack"),
@@ -106,12 +102,6 @@ static void applyReassociativeReshapeFoldingPatterns(Operation *rootOp) {
(void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
}
-static void applyEmptyOpFoldingPatterns(Operation *rootOp) {
- RewritePatternSet patterns(rootOp->getContext());
- tensor::populateFoldTensorEmptyPatterns(patterns);
- (void)applyPatternsAndFoldGreedily(rootOp, std::move(patterns));
-}
-
static void applyFoldIntoPackAndUnpackPatterns(Operation *rootOp) {
RewritePatternSet patterns(rootOp->getContext());
tensor::populateFoldIntoPackAndUnpackPatterns(patterns);
@@ -386,8 +376,6 @@ void TestTensorTransforms::runOnOperation() {
applyDropRedundantInsertSliceRankExpansionPatterns(rootOp);
if (testReassociativeReshapeFolding)
applyReassociativeReshapeFoldingPatterns(rootOp);
- if (testEmptyOpFolding)
- applyEmptyOpFoldingPatterns(rootOp);
if (testFoldIntoPackAndUnpack)
applyFoldIntoPackAndUnpackPatterns(rootOp);
if (testRewriteExtractSliceWithTiledCollapseShape) {
More information about the Mlir-commits
mailing list