[mlir][bufferization] Buffer deallocation: Make op preconditions stricter #75127
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-bufferization Author: Matthias Springer (matthias-springer) ChangesMemory side effects of unregistered ops are unknown. In particular, we do not know whether an unregistered op allocates memory or not. Therefore, unregistered ops cannot be handled safely in the buffer deallocation pass. In test case, Depends on #75126. Review only the top commit. Full diff: https://github.com/llvm/llvm-project/pull/75127.diff 7 Files Affected:
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 38ffae68a43de2..12b02d0f96313a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -48,6 +48,19 @@ static Value buildBoolValue(OpBuilder &builder, Location loc, bool value) {
static bool isMemref(Value v) { return v.getType().isa<BaseMemRefType>(); }
+/// Return "true" if the given op has buffer semantics. I.e., it has buffer
+/// operands, buffer results and/or buffer region entry block arguments.
+static bool hasBufferSemantics(Operation *op) {
+ if (llvm::any_of(op->getOperands(), isMemref) ||
+ llvm::any_of(op->getResults(), isMemref))
+ return true;
+ for (Region ®ion : op->getRegions())
+ if (!region.empty())
+ if (llvm::any_of(region.front().getArguments(), isMemref))
+ return true;
+ return false;
+}
+
//===----------------------------------------------------------------------===//
// Backedges analysis
//===----------------------------------------------------------------------===//
@@ -462,31 +475,6 @@ BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,
return state.getMemrefWithUniqueOwnership(builder, memref, block);
}
-static bool regionOperatesOnMemrefValues(Region ®ion) {
- auto checkBlock = [](Block *block) {
- if (llvm::any_of(block->getArguments(), isMemref))
- return WalkResult::interrupt();
- for (Operation &op : *block) {
- if (llvm::any_of(op.getOperands(), isMemref))
- return WalkResult::interrupt();
- if (llvm::any_of(op.getResults(), isMemref))
- return WalkResult::interrupt();
- }
- return WalkResult::advance();
- };
- WalkResult result = region.walk(checkBlock);
- if (result.wasInterrupted())
- return true;
-
- // Note: Block::walk/Region::walk visits only blocks that are nested under
- // nested operations, but not direct children.
- for (Block &block : region)
- if (checkBlock(&block).wasInterrupted())
- return true;
-
- return false;
-}
-
LogicalResult
BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) {
// (1) Ensure that there are supported loops only (no explicit control flow
@@ -501,6 +489,11 @@ BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) {
}
LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {
+ // (0) Memory side effects of unregistered ops are unknown. In particular, we
+ // do not know whether an unregistered op allocates memory or not.
+ if (!op->isRegistered())
+ return op->emitError("Unregistered ops are not supported");
+
// (1) Check that the control flow structures are supported.
auto regions = op->getRegions();
// Check that if the operation has at
@@ -512,9 +505,8 @@ LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {
size_t size = regions.size();
if (((size == 1 && !op->getResults().empty()) || size > 1) &&
!dyn_cast<RegionBranchOpInterface>(op)) {
- if (llvm::any_of(regions, regionOperatesOnMemrefValues))
- return op->emitError("All operations with attached regions need to "
- "implement the RegionBranchOpInterface.");
+ return op->emitError("All operations with attached regions need to "
+ "implement the RegionBranchOpInterface.");
}
// (2) The pass does not work properly when deallocations are already present.
@@ -648,6 +640,11 @@ LogicalResult BufferDeallocation::deallocate(Block *block) {
// For each operation in the block, handle the interfaces that affect aliasing
// and ownership of memrefs.
for (Operation &op : llvm::make_early_inc_range(*block)) {
+ // Skip ops that do not operate on buffers and are not terminators.
+ // (bufferization.dealloc ops are inserted in front of terminators, so
+ // terminators cannot be skipped.)
+ if (!op.hasTrait<OpTrait::IsTerminator>() && !hasBufferSemantics(&op))
+ continue;
FailureOr<Operation *> result = handleAllInterfaces(&op);
if (failed(result))
return failure();
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
index 3ae0529ab7d746..6d65502bcf9be7 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-branchop-interface.mlir
@@ -570,7 +570,7 @@ func.func @select_captured_in_next_block(%arg0: index, %arg1: memref<?xi8>, %arg
func.func @blocks_not_preordered_by_dominance() {
cf.br ^bb1
^bb2:
- "test.memref_user"(%alloc) : (memref<2xi32>) -> ()
+ "test.same_operand_shape"(%alloc) : (memref<2xi32>) -> ()
return
^bb1:
%alloc = memref.alloc() : memref<2xi32>
@@ -581,7 +581,7 @@ func.func @blocks_not_preordered_by_dominance() {
// CHECK-NEXT: [[TRUE:%.+]] = arith.constant true
// CHECK-NEXT: cf.br [[BB1:\^.+]]
// CHECK-NEXT: [[BB2:\^[a-zA-Z0-9_]+]]:
-// CHECK-NEXT: "test.memref_user"([[ALLOC:%[a-zA-Z0-9_]+]])
+// CHECK-NEXT: "test.same_operand_shape"([[ALLOC:%[a-zA-Z0-9_]+]])
// CHECK-NEXT: bufferization.dealloc ([[ALLOC]] : {{.*}}) if ([[TRUE]])
// CHECK-NOT: retain
// CHECK-NEXT: return
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir
index 7014746e348d25..51cbfe6b764536 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-existing-deallocs.mlir
@@ -8,7 +8,7 @@ func.func @auto_dealloc() {
%c100 = arith.constant 100 : index
%alloc = memref.alloc(%c10) : memref<?xi32>
%realloc = memref.realloc %alloc(%c100) : memref<?xi32> to memref<?xi32>
- "test.memref_user"(%realloc) : (memref<?xi32>) -> ()
+ "test.same_operand_shape"(%realloc) : (memref<?xi32>) -> ()
return
}
@@ -17,7 +17,7 @@ func.func @auto_dealloc() {
// CHECK-NOT: bufferization.dealloc
// CHECK: [[V0:%.+]]:2 = scf.if
// CHECK-NOT: bufferization.dealloc
-// CHECK: test.memref_user
+// CHECK: test.same_operand_shape
// CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK-NEXT: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1)
// CHECK-NEXT: return
@@ -32,14 +32,14 @@ func.func @auto_dealloc_inside_nested_region(%arg0: memref<?xi32>, %arg1: i1) {
} else {
scf.yield %arg0 : memref<?xi32>
}
- "test.memref_user"(%0) : (memref<?xi32>) -> ()
+ "test.same_operand_shape"(%0) : (memref<?xi32>) -> ()
return
}
// CHECK-LABEL: func @auto_dealloc_inside_nested_region
// CHECK-SAME: (%arg0: memref<?xi32>, %arg1: i1)
// CHECK-NOT: dealloc
-// CHECK: "test.memref_user"([[V0:%.+]]#0)
+// CHECK: "test.same_operand_shape"([[V0:%.+]]#0)
// CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
// CHECK-NEXT: bufferization.dealloc ([[BASE]] : memref<i32>) if ([[V0]]#1)
// CHECK-NEXT: return
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
index 13c55d0289880e..d9c8dab77e7324 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
@@ -12,7 +12,7 @@
func.func private @emptyUsesValue(%arg0: memref<4xf32>) {
%0 = memref.alloc() : memref<4xf32>
- "test.memref_user"(%0) : (memref<4xf32>) -> ()
+ "test.same_operand_shape"(%0) : (memref<4xf32>) -> ()
return
}
@@ -37,7 +37,7 @@ func.func private @emptyUsesValue(%arg0: memref<4xf32>) {
func.func @emptyUsesValue(%arg0: memref<4xf32>) {
%0 = memref.alloc() : memref<4xf32>
- "test.memref_user"(%0) : (memref<4xf32>) -> ()
+ "test.same_operand_shape"(%0) : (memref<4xf32>) -> ()
return
}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir
index be894dbe4c5fc8..50093e4d472c94 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-other.mlir
@@ -21,11 +21,20 @@ func.func private @no_interface_no_operands(%t : tensor<?x?x?xf16>) -> memref<?x
// CHECK-LABEL: func private @no_interface(
// CHECK: %[[true:.*]] = arith.constant true
// CHECK: %[[alloc:.*]] = memref.alloc
-// CHECK: %[[foo:.*]] = "test.foo"(%[[alloc]])
+// CHECK: %[[foo:.*]] = "test.same_operand_and_result_type"(%[[alloc]])
// CHECK: %[[r:.*]] = bufferization.dealloc (%[[alloc]] : {{.*}}) if (%[[true]]) retain (%[[foo]] : {{.*}})
// CHECK: return %[[foo]]
func.func private @no_interface() -> memref<5xf32> {
%0 = memref.alloc() : memref<5xf32>
- %1 = "test.foo"(%0) : (memref<5xf32>) -> (memref<5xf32>)
+ %1 = "test.same_operand_and_result_type"(%0) : (memref<5xf32>) -> (memref<5xf32>)
return %1 : memref<5xf32>
}
+
+// -----
+
+func.func @unregistered_op() {
+ %0 = memref.alloc() : memref<5xf32>
+ // expected-error @below{{Unregistered ops are not supported}}
+ "test.unregistered_op_foo"(%0) : (memref<5xf32>) -> ()
+ return
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
index 1a8a930bc9002b..02726436376c27 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
@@ -71,7 +71,7 @@ func.func @nested_region_control_flow(
scf.yield %1 : memref<?x?xf32>
} else {
%3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
- "test.memref_user"(%3) : (memref<?x?xf32>) -> ()
+ "test.same_operand_shape"(%3) : (memref<?x?xf32>) -> ()
scf.yield %1 : memref<?x?xf32>
}
return %2 : memref<?x?xf32>
@@ -253,7 +253,7 @@ func.func @loop_alloc(
%buf: memref<2xf32>,
%res: memref<2xf32>) {
%0 = memref.alloc() : memref<2xf32>
- "test.memref_user"(%0) : (memref<2xf32>) -> ()
+ "test.same_operand_shape"(%0) : (memref<2xf32>) -> ()
%1 = scf.for %i = %lb to %ub step %step
iter_args(%iterBuf = %buf) -> memref<2xf32> {
%2 = arith.cmpi eq, %i, %ub : index
@@ -385,7 +385,7 @@ func.func @loop_nested_alloc(
%buf: memref<2xf32>,
%res: memref<2xf32>) {
%0 = memref.alloc() : memref<2xf32>
- "test.memref_user"(%0) : (memref<2xf32>) -> ()
+ "test.same_operand_shape"(%0) : (memref<2xf32>) -> ()
%1 = scf.for %i = %lb to %ub step %step
iter_args(%iterBuf = %buf) -> memref<2xf32> {
%2 = scf.for %i2 = %lb to %ub step %step
@@ -393,7 +393,7 @@ func.func @loop_nested_alloc(
%3 = scf.for %i3 = %lb to %ub step %step
iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
%4 = memref.alloc() : memref<2xf32>
- "test.memref_user"(%4) : (memref<2xf32>) -> ()
+ "test.same_operand_shape"(%4) : (memref<2xf32>) -> ()
%5 = arith.cmpi eq, %i, %ub : index
%6 = scf.if %5 -> (memref<2xf32>) {
%7 = memref.alloc() : memref<2xf32>
@@ -476,7 +476,7 @@ func.func @assumingOp(
// Confirm the alloc will be dealloc'ed in the block.
%1 = shape.assuming %arg0 -> memref<2xf32> {
%0 = memref.alloc() : memref<2xf32>
- "test.memref_user"(%0) : (memref<2xf32>) -> ()
+ "test.same_operand_shape"(%0) : (memref<2xf32>) -> ()
shape.assuming_yield %arg2 : memref<2xf32>
}
// Confirm the alloc will be returned and dealloc'ed after its use.
@@ -533,9 +533,9 @@ func.func @noRegionBranchOpInterface() {
func.func @noRegionBranchOpInterface() {
%0 = "test.bar"() ({
// expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
- %1 = "test.bar"() ({
- %2 = "test.get_memref"() : () -> memref<2xi32>
- "test.yield"(%2) : (memref<2xi32>) -> ()
+ %1 = "test.merge_blocks"() ({
+ %2 = memref.alloc() : memref<2xi32>
+ "test.return"(%2) : (memref<2xi32>) -> ()
}) : () -> (memref<2xi32>)
"test.yield"() : () -> ()
}) : () -> (i32)
@@ -545,13 +545,13 @@ func.func @noRegionBranchOpInterface() {
// -----
// Test Case: The op "test.bar" does not implement the RegionBranchOpInterface.
-// This is not allowed in buffer deallocation.
+// But it also does not operate on buffers, so we don't care.
func.func @noRegionBranchOpInterface() {
- // expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
%0 = "test.bar"() ({
- %2 = "test.get_memref"() : () -> memref<2xi32>
- %3 = "test.foo"(%2) : (memref<2xi32>) -> (i32)
+ %1 = memref.alloc() : memref<2xi32>
+ "test.same_operand_shape"(%1) : (memref<2xi32>) -> ()
+ %3 = "test.foo"() : () -> (i32)
"test.yield"(%3) : (i32) -> ()
}) : () -> (i32)
"test.terminator"() : () -> ()
diff --git a/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir b/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir
index 25349967e61d3e..5003e98fc87c6f 100644
--- a/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir
+++ b/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir
@@ -5,7 +5,7 @@ func.func @gpu_launch() {
gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) {
%alloc = memref.alloc() : memref<2xf32>
- "test.memref_user"(%alloc) : (memref<2xf32>) -> ()
+ "test.same_operand_shape"(%alloc) : (memref<2xf32>) -> ()
gpu.terminator
}
return
|
It's a bit concerning because a registered op that does not implement the memory effect interface does not provide more information either. |
|
That is a problem indeed. I just took a look at We currently assume throughout the buffer deallocation infrastructure that an op does not have an allocation side effect if the op does not declare If an op does not implement I think we can utilize a bunch of MLIR traits here. Instead of checking whether an op is unregistered or not, something along the lines of:
Unregistered ops will be rejected automatically because they do not implement any interface or trait. I'm going to give it a try and update this PR. |
|
Another tricky case would be ops that sometimes allocate and sometimes don't. This is something that we do not support properly at the moment. // With a probability of 40%: %0 = %m
// With a probability of 60%: %0 is a new allocation
%0 = test.probabilistic_alloc {p = 0.4 : f32} %m : memref<?xf32>Such an op would have to declare an (In the example above, without further knowledge about the op, @maerhart Does this make sense? |
This was referenced Dec 12, 2023
matthias-springer
force-pushed
the
buffer_deallocation_no_unregistered_ops
branch
from
7080e6c to
2b5a769
Compare
matthias-springer
changed the title
I am going to prepare a change that adds "maybe side effects", which should address the "may allocate" issue: https://discourse.llvm.org/t/how-to-model-maybe-side-effects/75674. (This may take a while...) |
matthias-springer
force-pushed
the
buffer_deallocation_no_unregistered_ops
branch
from
2b5a769 to
89fd111
Compare
aartbik
approved these changes
Jan 19, 2024
mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
Outdated
Show resolved
Hide resolved
mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
Outdated
Show resolved
Hide resolved
mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
Outdated
Show resolved
Hide resolved
...erization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
Show resolved
Hide resolved
...erization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
Outdated
Show resolved
Hide resolved
Memory side effects of unregistered ops are unknown. In particular, we do not know whether an unregistered op allocates memory or not. Therefore, unregistered ops cannot be handled safely in the buffer deallocation pass.
matthias-springer
force-pushed
the
buffer_deallocation_no_unregistered_ops
branch
from
89fd111 to
622e834
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The buffer deallocation pass checks the IR ("operation preconditions") to make sure that there is no IR that is unsupported. In such a case, the pass signals a failure.
The pass now rejects all ops with unknown memory effects. We do not know whether such an op allocates memory or not. Therefore, the buffer deallocation pass does not know whether a deallocation op should be inserted or not.
Memory effects are queried from the
MemoryEffectOpInterfaceinterface. Ops that do not implement this interface but have theRecursiveMemoryEffectstrait do not have any side effects (apart from the ones that their nested ops may have).Unregistered ops are now rejected by the pass because they do not implement the
MemoryEffectOpInterfaceand neither do we know if they haveRecursiveMemoryEffectsor not. All test cases that currently have unregistered ops are updated to use registered ops.