[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.

Author: matthias-springer

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<BaseMemRefType>(); }
 
+/// Return "true" if the given op is guaranteed to have no "Allocate" or "Free"
+/// side effect.
+static bool hasNoAllocateOrFreeSideEffect(Operation *op) {
+  if (isa<MemoryEffectOpInterface>(op))
+    return hasEffect<MemoryEffects::Allocate>(op) ||
+           hasEffect<MemoryEffects::Free>(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. The ops inside the regions of this op will be
+  // processed separately.
+  return op->hasTrait<OpTrait::HasRecursiveMemoryEffects>();
+}
+
+/// 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 &region : 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 &region) {
-  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,29 @@ 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<DeallocOp>(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. Call ops
+  // typically do not implement the MemoryEffectOpInterface but usually do not
+  // have side effects (apart from the callee, which will be analyzed
+  // separately), so they are allowed.
+  if (!isa<MemoryEffectOpInterface>(op) &&
+      !op->hasTrait<OpTrait::HasRecursiveMemoryEffects>() &&
+      !isa<CallOpInterface>(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) && hasNoAllocateOrFreeSideEffect(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 +535,10 @@ 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.
-  // Alternatively, we could also remove all deallocations as a pre-pass.
-  if (isa<DeallocOp>(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.

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.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

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.read_buffer"(%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.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<?xi32>, %arg1: i1) {
   } else {
     scf.yield %arg0 : memref<?xi32>
   }
-  "test.memref_user"(%0) : (memref<?xi32>) -> ()
+  "test.read_buffer"(%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.read_buffer"([[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

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
 }
 

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.forward_buffer"(%[[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.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
+}

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.read_buffer"(%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.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,15 +385,15 @@ 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
       iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
       %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,58 +511,49 @@ 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"() : () -> ()
 }
 
 // -----
 
 func.func @while_two_arg(%arg0: index) {
   %a = memref.alloc(%arg0) : memref<?xf32>
   scf.while (%arg1 = %a, %arg2 = %a) : (memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>) {
-    %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<?xf32>, memref<?xf32>
   } do {
   ^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>):
@@ -591,7 +582,8 @@ func.func @while_two_arg(%arg0: index) {
 func.func @while_three_arg(%arg0: index) {
   %a = memref.alloc(%arg0) : memref<?xf32>
   scf.while (%arg1 = %a, %arg2 = %a, %arg3 = %a) : (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>, memref<?xf32>) {
-    %0 = "test.make_condition"() : () -> i1
+    // This op has a side effect but it's not an Allocate or Free side effect.
+    %0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1
     scf.condition(%0) %arg1, %arg2, %arg3 : memref<?xf32>, memref<?xf32>, memref<?xf32>
   } do {
   ^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: memref<?xf32>):

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

mlir/test/lib/Dialect/Test/TestDialect.cpp

@@ -1272,6 +1272,21 @@ static void printSumProperty(OpAsmPrinter &printer, Operation *op,
   printer << second << " = " << (second + first);
 }
 
+//===----------------------------------------------------------------------===//
+// Tensor/Buffer Ops
+//===----------------------------------------------------------------------===//
+
+void ReadBufferOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &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
 //===----------------------------------------------------------------------===//