[Mlir-commits] [mlir] [mlir][bufferization] Buffer deallocation: Disallow unregistered ops (PR #75127)
Matthias Springer
llvmlistbot at llvm.org
Mon Dec 11 17:55:53 PST 2023
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/75127
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.
In test case, `test.memref_user` (unregistered op) is replaced with `test.same_operand_shape`, which is a registered op without any memory side effects that supports memref operand types.
Depends on #75126. Review only the top commit.
>From 9da3a82099cc53aa278ed40881de124e22b3e5a4 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 12 Dec 2023 10:48:28 +0900
Subject: [PATCH 1/2] [mlir][bufferization] Buffer deallocation: skip ops that
do not operate on buffers
Skip ops that do not operate on buffers during ownership-based buffer deallocation. Such ops can be ignored during buffer deallocation. (Except for terminators.)
---
.../OwnershipBasedBufferDeallocation.cpp | 48 ++++++++-----------
.../dealloc-region-branchop-interface.mlir | 3 +-
2 files changed, 21 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 38ffae68a43de2..f9ec6564995c83 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
@@ -512,9 +500,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 +635,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-region-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-region-branchop-interface.mlir
index 1a8a930bc9002b..f44b6d74051e2d 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
@@ -545,10 +545,9 @@ 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 at +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)
>From 7080e6c6b2e728fe04931a41500c12d6d3208f79 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 12 Dec 2023 10:53:15 +0900
Subject: [PATCH 2/2] [mlir][bufferization] Buffer deallocation: Disallow
unregistered ops
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.
In test case, `test.memref_user` (unregistered op) is replaced with `test.same_operand_shape`, which is a registered op without any memory side effects that supports memref operand types.
---
.../OwnershipBasedBufferDeallocation.cpp | 5 +++++
.../dealloc-branchop-interface.mlir | 4 ++--
.../dealloc-existing-deallocs.mlir | 8 +++----
.../dealloc-function-boundaries.mlir | 4 ++--
.../dealloc-other.mlir | 13 ++++++++++--
.../dealloc-region-branchop-interface.mlir | 21 ++++++++++---------
.../bufferization-buffer-deallocation.mlir | 2 +-
7 files changed, 36 insertions(+), 21 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index f9ec6564995c83..12b02d0f96313a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -489,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
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 f44b6d74051e2d..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 at +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)
@@ -549,8 +549,9 @@ func.func @noRegionBranchOpInterface() {
func.func @noRegionBranchOpInterface() {
%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
More information about the Mlir-commits
mailing list