From 622e834ab081d9b2ba99d6e138f03c8eec82313a Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Thu, 4 Jan 2024 23:20:36 +0900 Subject: [PATCH] [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. --- .../OwnershipBasedBufferDeallocation.cpp | 79 +++++++++++++------ .../dealloc-branchop-interface.mlir | 4 +- .../dealloc-existing-deallocs.mlir | 8 +- .../dealloc-function-boundaries.mlir | 4 +- .../dealloc-other.mlir | 13 ++- .../dealloc-region-branchop-interface.mlir | 68 +++++++--------- .../bufferization-buffer-deallocation.mlir | 2 +- mlir/test/lib/Dialect/Test/TestDialect.cpp | 15 ++++ mlir/test/lib/Dialect/Test/TestOps.td | 37 ++++++++- 9 files changed, 153 insertions(+), 77 deletions(-) diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp index 529d5a808c012..c535d6c130edb 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp @@ -48,6 +48,32 @@ static Value buildBoolValue(OpBuilder &builder, Location loc, bool value) { static bool isMemref(Value v) { return v.getType().isa(); } +/// Return "true" if the given op is guaranteed to have neither "Allocate" nor +/// "Free" side effects. +static bool hasNeitherAllocateNorFreeSideEffect(Operation *op) { + if (isa(op)) + return hasEffect(op) || + hasEffect(op); + // If the op does not implement the MemoryEffectOpInterface but has has + // recursive memory effects, then this op in isolation (without its body) does + // not have any side effects. All the ops inside the regions of this op will + // be processed separately. + return op->hasTrait(); +} + +/// 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,21 +488,6 @@ BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref, return state.getMemrefWithUniqueOwnership(builder, memref, block); } -static bool regionOperatesOnMemrefValues(Region ®ion) { - WalkResult result = region.walk([](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(); - }); - return result.wasInterrupted(); -} - LogicalResult BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) { // (1) Ensure that there are supported loops only (no explicit control flow @@ -491,7 +502,32 @@ BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) { } LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) { - // (1) Check that the control flow structures are supported. + // (1) The pass does not work properly when deallocations are already present. + // Alternatively, we could also remove all deallocations as a pre-pass. + if (isa(op)) + return op->emitError( + "No deallocation operations must be present when running this pass!"); + + // (2) Memory side effects of unregistered ops are unknown. In particular, we + // do not know whether an unregistered op allocates memory or not. + // - Ops with recursive memory effects are allowed. All nested ops in the + // regions of `op` will be analyzed separately. + // - Call ops are allowed even though they typically do not implement the + // MemoryEffectOpInterface. They usually do not have side effects apart + // from the callee, which will be analyzed separately. (This is similar to + // "recursive memory effects".) + if (!isa(op) && + !op->hasTrait() && + !isa(op)) + return op->emitError( + "ops with unknown memory side effects are not supported"); + + // We do not care about ops that do not operate on buffers and have no + // Allocate/Free side effect. + if (!hasBufferSemantics(op) && hasNeitherAllocateNorFreeSideEffect(op)) + return success(); + + // (3) Check that the control flow structures are supported. auto regions = op->getRegions(); // Check that if the operation has at // least one region it implements the RegionBranchOpInterface. If there @@ -502,17 +538,10 @@ LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) { size_t size = regions.size(); if (((size == 1 && !op->getResults().empty()) || size > 1) && !dyn_cast(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. - // Alternatively, we could also remove all deallocations as a pre-pass. - if (isa(op)) - return op->emitError( - "No deallocation operations must be present when running this pass!"); - // (3) Check that terminators with more than one successor except `cf.cond_br` // are not present and that either BranchOpInterface or // RegionBranchTerminatorOpInterface is implemented. 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 3ae0529ab7d74..5e8104f83cc4d 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, %arg func.func @blocks_not_preordered_by_dominance() { cf.br ^bb1 ^bb2: - "test.memref_user"(%alloc) : (memref<2xi32>) -> () + "test.read_buffer"(%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.read_buffer"([[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 7014746e348d2..7a41ed2057031 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 %realloc = memref.realloc %alloc(%c100) : memref to memref - "test.memref_user"(%realloc) : (memref) -> () + "test.read_buffer"(%realloc) : (memref) -> () 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.read_buffer // 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, %arg1: i1) { } else { scf.yield %arg0 : memref } - "test.memref_user"(%0) : (memref) -> () + "test.read_buffer"(%0) : (memref) -> () return } // CHECK-LABEL: func @auto_dealloc_inside_nested_region // CHECK-SAME: (%arg0: memref, %arg1: i1) // CHECK-NOT: dealloc -// CHECK: "test.memref_user"([[V0:%.+]]#0) +// CHECK: "test.read_buffer"([[V0:%.+]]#0) // CHECK-NEXT: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0 // CHECK-NEXT: bufferization.dealloc ([[BASE]] : memref) 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 13c55d0289880..de71f23cfad6f 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.read_buffer"(%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.read_buffer"(%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 be894dbe4c5fc..9bfa91589482b 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) -> memref memref<5xf32> { %0 = memref.alloc() : memref<5xf32> - %1 = "test.foo"(%0) : (memref<5xf32>) -> (memref<5xf32>) + %1 = "test.forward_buffer"(%0) : (memref<5xf32>) -> (memref<5xf32>) return %1 : memref<5xf32> } + +// ----- + +func.func @no_side_effects() { + %0 = memref.alloc() : memref<5xf32> + // expected-error @below{{ops with unknown memory side effects 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 1a8a930bc9002..423fc4730b137 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 } else { %3 = memref.alloc(%arg0, %arg1) : memref - "test.memref_user"(%3) : (memref) -> () + "test.read_buffer"(%3) : (memref) -> () scf.yield %1 : memref } return %2 : memref @@ -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.read_buffer"(%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.read_buffer"(%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.read_buffer"(%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.read_buffer"(%0) : (memref<2xf32>) -> () shape.assuming_yield %arg2 : memref<2xf32> } // Confirm the alloc will be returned and dealloc'ed after its use. @@ -511,50 +511,40 @@ func.func @assumingOp( // ----- -// Test Case: The op "test.bar" does not implement the RegionBranchOpInterface. -// This is only allowed in buffer deallocation because the operation's region -// does not deal with any MemRef values. +// Test Case: The op "test.one_region_with_recursive_memory_effects" does not +// implement the RegionBranchOpInterface. This is allowed during buffer +// deallocation because the operation's region does not deal with any MemRef +// values. func.func @noRegionBranchOpInterface() { - %0 = "test.bar"() ({ - %1 = "test.bar"() ({ - "test.yield"() : () -> () + %0 = "test.one_region_with_recursive_memory_effects"() ({ + %1 = "test.one_region_with_recursive_memory_effects"() ({ + %2 = memref.alloc() : memref<2xi32> + "test.read_buffer"(%2) : (memref<2xi32>) -> () + "test.return"() : () -> () }) : () -> (i32) - "test.yield"() : () -> () + "test.return"() : () -> () }) : () -> (i32) - "test.terminator"() : () -> () + "test.return"() : () -> () } // ----- -// Test Case: The op "test.bar" does not implement the RegionBranchOpInterface. -// This is not allowed in buffer deallocation. +// Test Case: The second op "test.one_region_with_recursive_memory_effects" does +// not implement the RegionBranchOpInterface but has buffer semantics. This is +// not allowed during buffer deallocation. func.func @noRegionBranchOpInterface() { - %0 = "test.bar"() ({ + %0 = "test.one_region_with_recursive_memory_effects"() ({ // 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.one_region_with_recursive_memory_effects"() ({ + %2 = memref.alloc() : memref<2xi32> + "test.read_buffer"(%2) : (memref<2xi32>) -> () + "test.return"(%2) : (memref<2xi32>) -> () }) : () -> (memref<2xi32>) - "test.yield"() : () -> () + "test.return"() : () -> () }) : () -> (i32) - "test.terminator"() : () -> () -} - -// ----- - -// Test Case: The op "test.bar" does not implement the RegionBranchOpInterface. -// This is not allowed in buffer deallocation. - -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) - "test.yield"(%3) : (i32) -> () - }) : () -> (i32) - "test.terminator"() : () -> () + "test.return"() : () -> () } // ----- @@ -562,7 +552,8 @@ func.func @noRegionBranchOpInterface() { func.func @while_two_arg(%arg0: index) { %a = memref.alloc(%arg0) : memref scf.while (%arg1 = %a, %arg2 = %a) : (memref, memref) -> (memref, memref) { - %0 = "test.make_condition"() : () -> i1 + // This op has a side effect, but it's not an allocate/free side effect. + %0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1 scf.condition(%0) %arg1, %arg2 : memref, memref } do { ^bb0(%arg1: memref, %arg2: memref): @@ -591,7 +582,8 @@ func.func @while_two_arg(%arg0: index) { func.func @while_three_arg(%arg0: index) { %a = memref.alloc(%arg0) : memref scf.while (%arg1 = %a, %arg2 = %a, %arg3 = %a) : (memref, memref, memref) -> (memref, memref, memref) { - %0 = "test.make_condition"() : () -> i1 + // This op has a side effect, but it's not an allocate/free side effect. + %0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1 scf.condition(%0) %arg1, %arg2, %arg3 : memref, memref, memref } do { ^bb0(%arg1: memref, %arg2: memref, %arg3: memref): diff --git a/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir b/mlir/test/Dialect/GPU/bufferization-buffer-deallocation.mlir index 25349967e61d3..57260c86efca2 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.read_buffer"(%alloc) : (memref<2xf32>) -> () gpu.terminator } return diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp index cb5ee6014b611..1ee52fc08d77f 100644 --- a/mlir/test/lib/Dialect/Test/TestDialect.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -1321,6 +1321,21 @@ static void printSumProperty(OpAsmPrinter &printer, Operation *op, printer << second << " = " << (second + first); } +//===----------------------------------------------------------------------===// +// Tensor/Buffer Ops +//===----------------------------------------------------------------------===// + +void ReadBufferOp::getEffects( + SmallVectorImpl> + &effects) { + // The buffer operand is read. + effects.emplace_back(MemoryEffects::Read::get(), getBuffer(), + SideEffects::DefaultResource::get()); + // The buffer contents are dumped. + effects.emplace_back(MemoryEffects::Write::get(), + SideEffects::DefaultResource::get()); +} + //===----------------------------------------------------------------------===// // Test Dataflow //===----------------------------------------------------------------------===// diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td index 61e181999ff2f..11e409c6072f7 100644 --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -486,6 +486,17 @@ def AllocaScopeRegionOp : TEST_Op<"alloca_scope_region", let assemblyFormat = "attr-dict-with-keyword $region"; } +def OneRegionWithRecursiveMemoryEffectsOp + : TEST_Op<"one_region_with_recursive_memory_effects", [ + RecursiveMemoryEffects]> { + let description = [{ + Op that has one region and recursive side effects. The + RegionBranchOpInterface is not implemented on this op. + }]; + let results = (outs AnyType:$result); + let regions = (region SizedRegion<1>:$body); +} + //===----------------------------------------------------------------------===// // NoTerminator Operation //===----------------------------------------------------------------------===// @@ -1900,7 +1911,7 @@ def BlackHoleOp : TEST_Op<"blackhole">, //===----------------------------------------------------------------------===// def TestRegionBuilderOp : TEST_Op<"region_builder">; -def TestReturnOp : TEST_Op<"return", [ReturnLike, Terminator]> { +def TestReturnOp : TEST_Op<"return", [Pure, ReturnLike, Terminator]> { let arguments = (ins Variadic); let builders = [OpBuilder<(ins), [{ build($_builder, $_state, {}); }]> @@ -2106,8 +2117,10 @@ class BufferBasedOpBase traits> let description = [{ A buffer based operation, that uses memRefs as input and output. }]; - let arguments = (ins AnyRankedOrUnrankedMemRef:$input, - AnyRankedOrUnrankedMemRef:$output); + let arguments = (ins Arg:$input, + Arg:$output); } def BufferBasedOp : BufferBasedOpBase<"buffer_based", []>{ @@ -2138,6 +2151,24 @@ def TensorBasedOp : TEST_Op<"tensor_based", []> { }]; } +def ReadBufferOp : TEST_Op<"read_buffer", [DeclareOpInterfaceMethods]> { + let description = [{ + An operation that reads the buffer operand and dumps its contents. + }]; + let arguments = (ins AnyRankedOrUnrankedMemRef:$buffer); +} + +def ForwardBufferOp : TEST_Op<"forward_buffer", [Pure]> { + let description = [{ + A pure operation that takes a buffer and returns a buffer. This op does not + have any side effects, so it cannot allocate or read a buffer from memory. + It must return the input buffer (or a view thereof). This op purposely does + does not implement any interface. + }]; + let arguments = (ins AnyRankedOrUnrankedMemRef:$buffer); + let results = (outs AnyRankedOrUnrankedMemRef:$result); +} + //===----------------------------------------------------------------------===// // Test RegionBranchOpInterface //===----------------------------------------------------------------------===//