[mlir][bufferization] Buffer deallocation: Make op preconditions stricter

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 neither "Allocate" nor
+/// "Free" side effects.
+static bool hasNeitherAllocateNorFreeSideEffect(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. All 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,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<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.
+  // - 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<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) && 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<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/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

@@ -1321,6 +1321,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
 //===----------------------------------------------------------------------===//