[Mlir-commits] [mlir] c5501a7 - [mlir][linalg][bufferize] Eliminate InitTensorOps of InsertSliceOp sources
Matthias Springer
llvmlistbot at llvm.org
Thu Oct 21 05:43:40 PDT 2021
Author: Matthias Springer
Date: 2021-10-21T21:33:45+09:00
New Revision: c5501a7a5c582168d1b12335c092bce53b32022d
URL: https://github.com/llvm/llvm-project/commit/c5501a7a5c582168d1b12335c092bce53b32022d
DIFF: https://github.com/llvm/llvm-project/commit/c5501a7a5c582168d1b12335c092bce53b32022d.diff
LOG: [mlir][linalg][bufferize] Eliminate InitTensorOps of InsertSliceOp sources
An InitTensorOp is replaced with an ExtractSliceOp on the InsertSliceOp's destination. This optimization is applied after analysis and only to InsertSliceOps that were decided to bufferize inplace. Another analysis on the new ExtractSliceOp is needed after the rewrite.
Differential Revision: https://reviews.llvm.org/D111955
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
index 0655d97b4ccb9..eb54ceeadea0a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
@@ -1791,11 +1791,15 @@ static LogicalResult bufferize(OpBuilder &b, FuncOp funcOp,
return success();
}
-/// InitTensor always allocates.
+/// InitTensor always allocates (unless it was eliminated).
/// TODO: consider hoisting across function boundaries prior to bufferization.
static LogicalResult bufferize(OpBuilder &b, InitTensorOp initTensorOp,
BlockAndValueMapping &bvm,
BufferizationAliasInfo &aliasInfo) {
+ // The InitTensorOp may have been eliminated.
+ if (initTensorOp->getUses().empty())
+ return success();
+
// Take a guard before anything else.
OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(initTensorOp);
@@ -2761,6 +2765,89 @@ static void layoutPostProcessing(ModuleOp moduleOp) {
}
}
+/// Try to eliminate InitTensorOps inside funcOp. An InitTensorOp can be
+/// eliminated if it is eventually inserted into another tensor (and some other
+/// conditions are met).
+///
+/// E.g.:
+/// %0 = linalg.init_tensor
+/// %1 = linalg.fill(%cst, %0) {inplace = [true]}
+/// %2 = tensor.insert_slice %1 into %t[10][20][1]
+///
+/// InitTensorOp elimination will try to fill %t inplace instead of filling a
+/// new allocation %0 and inserting it into %t. This is done by replacing the
+/// InitTensorOp with:
+///
+/// %0 = tensor.extract_slice %t[10][20][1]
+///
+/// The analysis looks for matching ExtractSliceOp/InsertSliceOp pairs and lets
+/// those bufferize inplace in the absence of other conflicts.
+///
+/// Starting from an InsertSliceOp, an InitTensorOp at the end of the insert
+/// source's reverse use-def chain is eliminated if:
+/// * The InsertSliceOp was decided to bufferize inplace.
+/// * On the reverse use-def chain path from the InsertSliceOp to the
+/// InitTensorOp, all ops were decided to bufferize inplace and the buffer
+/// relation is "equivalent" (TODO: can be relaxed if needed).
+/// * The reverse use-def chain has exactly one end, which is the InitTensorOp.
+///
+/// Note that the newly inserted ExtractSliceOp may have to bufferize
+/// out-of-place due to RaW conflicts.
+static LogicalResult runInitTensorElimination(FuncOp funcOp,
+ BufferizationAliasInfo &aliasInfo,
+ DominanceInfo &domInfo) {
+ OpBuilder b(funcOp->getContext());
+
+ WalkResult status = funcOp->walk([&](tensor::InsertSliceOp insertOp) {
+ // Only inplace bufferized InsertSliceOps are eligible.
+ if (getInPlace(insertOp->getOpResult(0)) != InPlaceSpec::True)
+ return WalkResult::skip();
+
+ SetVector<Value> maybeInitTensor =
+ findValueInReverseUseDefChain(insertOp.source(), [](Value val) {
+ // Continue traversal until this function returns true.
+ OpResult opResult = val.dyn_cast<OpResult>();
+ if (!opResult)
+ return true;
+ if (getInPlace(opResult) != InPlaceSpec::True)
+ return true;
+ // Only equivalent tensors are supported at the moment. E.g., when
+ // taking a tensor.extract_slice of an init_tensor, we can currently
+ // not eliminate the init_tensor.
+ SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
+ if (!llvm::all_of(opOperands, [](OpOperand *operand) {
+ return bufferRelation(*operand) == BufferRelation::Equivalent;
+ }))
+ return true;
+ return false;
+ });
+ // Replace only if the InsertSliceOp source originates from exactly one
+ // InitTensorOp.
+ if (maybeInitTensor.size() != 1 ||
+ !maybeInitTensor.front().getDefiningOp<InitTensorOp>())
+ return WalkResult::skip();
+ Value initTensor = maybeInitTensor.front();
+
+ b.setInsertionPoint(initTensor.getDefiningOp());
+ auto extractOp = b.create<tensor::ExtractSliceOp>(
+ initTensor.getLoc(), insertOp.dest(), insertOp.getMixedOffsets(),
+ insertOp.getMixedSizes(), insertOp.getMixedStrides());
+ // Uses of the InitTensorOp are replaced here, but the op is not deleted.
+ // InitTensorOps without uses are ignored by the bufferization.
+ initTensor.replaceAllUsesWith(extractOp.result());
+ aliasInfo.createAliasInfoEntry(extractOp.result());
+
+ // Run analysis on the ExtractSliceOp.
+ if (failed(bufferizableInPlaceAnalysis(extractOp, aliasInfo, domInfo)))
+ return WalkResult::interrupt();
+
+ // Advance to the next operation.
+ return WalkResult::advance();
+ });
+
+ return failure(status.wasInterrupted());
+}
+
void LinalgComprehensiveModuleBufferize::runOnOperation() {
ModuleOp moduleOp = getOperation();
applyEnablingTransformations(moduleOp);
@@ -2800,6 +2887,13 @@ void LinalgComprehensiveModuleBufferize::runOnOperation() {
return;
}
+ // Try to eliminate InitTensorOps to avoid new allocations during the
+ // bufferization phase.
+ if (failed(runInitTensorElimination(funcOp, aliasInfo, domInfo))) {
+ signalPassFailure();
+ return;
+ }
+
// Bufferization phase.
if (!testAnalysisOnly) {
BlockAndValueMapping tensorToBufferMap;
diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
index ffcabaf4cf54c..1283525ae33cc 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir
@@ -1039,3 +1039,51 @@ func @non_reading_scf_for(%t1: tensor<?xf32> {linalg.inplaceable = true},
return %o, %v3 : tensor<?xf32>, vector<5xf32>
}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// InitTensorOp elimination
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @buffer_forwarding_conflict
+func @buffer_forwarding_conflict(%arg0: tensor<?xf32> {linalg.inplaceable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) {
+ %cst = arith.constant 0.000000e+00 : f32
+ // CHECK: tensor.extract_slice
+ // CHECK-SAME: {__inplace_results_attr__ = ["false"]
+ // Instead of allocating, share buffer with some inplace bufferization?
+ %0 = linalg.init_tensor [%arg1] : tensor<?xf32>
+
+ // CHECK: linalg.fill
+ // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+ %1 = linalg.fill(%cst, %0) : f32, tensor<?xf32> -> tensor<?xf32>
+
+ // CHECK: tensor.insert_slice
+ // CHECK-SAME: {__inplace_results_attr__ = ["false"]
+ %2 = tensor.insert_slice %1 into %arg0[0] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
+
+ // CHECK: tensor.insert_slice
+ // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+ %3 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
+ return %2, %3 : tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @buffer_forwarding_no_conflict
+func @buffer_forwarding_no_conflict(%arg0: tensor<?xf32> {linalg.inplaceable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) {
+ %cst = arith.constant 0.000000e+00 : f32
+ // CHECK: tensor.extract_slice
+ // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+ // Instead of allocating, share buffer with some inplace bufferization?
+ %0 = linalg.init_tensor [%arg1] : tensor<?xf32>
+
+ // CHECK: linalg.fill
+ // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+ %1 = linalg.fill(%cst, %0) : f32, tensor<?xf32> -> tensor<?xf32>
+
+ // CHECK: tensor.insert_slice
+ // CHECK-SAME: {__inplace_results_attr__ = ["true"]
+ %2 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
+ return %2, %2 : tensor<?xf32>, tensor<?xf32>
+}
diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
index 3df671251b196..693974cf9296c 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
@@ -798,3 +798,70 @@ func @dominance_violation_bug_1(%A : tensor<?x?xf32>, %idx : index) -> tensor<?x
return %rA : tensor<?x?xf32>
}
+
+// -----
+
+// CHECK: func @buffer_forwarding_conflict(
+// CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
+// CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index
+func @buffer_forwarding_conflict(
+ %t: tensor<?xf32> {linalg.buffer_layout = affine_map<(d0) -> (d0)>, linalg.inplaceable = true},
+ %sz: index)
+ -> (tensor<?xf32>, tensor<?xf32>)
+{
+ %f0 = arith.constant 0.0: f32
+ // Alloc is needed for the **first** insert_slice (due to backward traversal during analysis).
+ // CHECK: %[[DIM:.*]] = memref.dim %[[FUNC_ARG]]
+ // This allocs the whole dim to allow for a full clone of t.
+ // CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
+
+ // init_tensor itself does not alloc but forwards to the **second**
+ // insert_slice. InitTensorOp replaces the init_tensor with an out-of-place
+ // extract_slice.
+ // CHECK: %[[EXTRACT_SLICE_ALLOC:.*]] = memref.alloc(%[[sz]])
+ // CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
+ // TODO: This copy can be avoided because the copied data is never read.
+ // CHECK: linalg.copy(%[[T_SUBVIEW]], %[[EXTRACT_SLICE_ALLOC]])
+ %a = linalg.init_tensor[%sz] : tensor<?xf32>
+
+ // CHECK: linalg.fill({{.*}}, %[[EXTRACT_SLICE_ALLOC]]) : f32, memref<?xf32>
+ %f = linalg.fill(%f0, %a) : f32, tensor<?xf32> -> tensor<?xf32>
+
+ // CHECK: linalg.copy(%[[FUNC_ARG]], %[[ALLOC]]) : memref<?xf32>, memref<?xf32>
+ // CHECK: %[[SV0_ALLOC:.*]] = memref.subview %[[ALLOC]][0] [%[[sz]]] [1] : memref<?xf32> to memref<?xf32>
+ // CHECK: linalg.copy(%[[EXTRACT_SLICE_ALLOC]], %[[SV0_ALLOC]]) : memref<?xf32>, memref<?xf32>
+ %r0 = tensor.insert_slice %f into %t[0][%sz][1]: tensor<?xf32> into tensor<?xf32>
+
+ // CHECK: linalg.copy(%[[EXTRACT_SLICE_ALLOC]], %[[T_SUBVIEW]])
+ %r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
+
+ return %r0, %r1: tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK: func @buffer_forwarding_no_conflict(
+// CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
+// CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index
+func @buffer_forwarding_no_conflict(
+ %t: tensor<?xf32> {linalg.buffer_layout = affine_map<(d0) -> (d0)>, linalg.inplaceable = true},
+ %sz: index)
+ -> (tensor<?xf32>)
+{
+ %f0 = arith.constant 0.0: f32
+
+ // init_tensor itself does not alloc but forwards to the insert_slice.
+ // InitTensorOp replaces the init_tensor with an inplace extract_slice.
+ // CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
+ %a = linalg.init_tensor[%sz] : tensor<?xf32>
+
+ // CHECK: linalg.fill({{.*}}, %[[T_SUBVIEW]]) : f32, memref<?xf32
+ %f = linalg.fill(%f0, %a) : f32, tensor<?xf32> -> tensor<?xf32>
+
+ // Self-copy canonicalizes away later.
+ // CHECK: linalg.copy(%[[T_SUBVIEW]], %[[T_SUBVIEW]])
+ %r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
+
+ return %r1: tensor<?xf32>
+}
+
More information about the Mlir-commits
mailing list