[mlir][bufferization] Fix buffer-loop-hoisting with deallocs

[springerm]

Summary

buffer-loop-hoisting did not support IR with deallocs (and silently produced invalid IR; the alloc was hoisted but not the dealloc). This change adds support for simple cases where the alloc is in the same block as the dealloc (if any). buffer-loop-hoisting bails on all unsupported cases instead of producing incorrect IR.

Depends On D135054

mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp

for (auto allocValue : allocsAndAllocas) {
       // Find the actual placement block and determine the start operation using
       // an upper placement-block boundary. The idea is that placement block
       // cannot be moved any further upwards than the given upper bound.
       Block *placementBlock = findPlacementBlock(
           state, state.computeUpperBound(dominatorBlock, dependencyBlock));
       Operation *startOperation = BufferPlacementAllocs::getStartOperation(
           allocValue, placementBlock, liveness);
 
-      // Move the alloc in front of the start operation.
+      // Find the corresponding dealloc (if any).
       Operation *allocOperation = allocValue.getDefiningOp();
+      Operation *deallocOperation = nullptr;
+      bool isSupported = [&]() {
+        for (OpResult allocResult : allocOperation->getResults()) {
+          for (OpOperand &allocUse : allocResult.getUses()) {
+            Operation *deallocCandidate = allocUse.getOwner();
+            if (StateT::isMatchingDealloc(allocOperation, deallocCandidate)) {
+              // Multiple deallocs are not supported.
+              if (deallocOperation != nullptr &&
+                  deallocOperation != deallocCandidate)
+                return false;
+              // Deallocs in different blocks are not supported.
+              if (allocOperation->getBlock() != deallocCandidate->getBlock())
+                return false;
+              deallocOperation = deallocCandidate;
+            }
+          }
+        }
+        return true;
+      }();
+
+      // Skip unsupported allocs.
+      if (!isSupported)
+        continue;
+
+      // Move alloc and dealloc.
       allocOperation->moveBefore(startOperation);
+      if (deallocOperation)
+        deallocOperation->moveBefore(
+            allocOperation->getBlock()->getTerminator());
     }
   }
 
 private:
   /// Finds a valid placement block by walking upwards in the CFG until we
   /// either cannot continue our walk due to constraints (given by the StateT
   /// implementation) or we have reached the upper-most dominator block.
   Block *findPlacementBlock(StateT &state, Block *upperBound) {
bool isLegalPlacement(Operation *op) {
     return !BufferPlacementTransformationBase::isLoop(op);
   }
 
   /// Returns true if the given operation should be considered for hoisting.
   static bool shouldHoistOpType(Operation *op) {
     return llvm::isa<memref::AllocOp>(op);
   }
 
+  /// Returns true if the `deallocOp` is a deallocation of a result of
+  /// `allocOp`. One of the operands of `deallocOp` is guaranteed to be a result
+  /// of the `allocOp`.
+  static bool isMatchingDealloc(Operation *allocOp, Operation *deallocOp) {
+    return llvm::isa<memref::DeallocOp>(deallocOp);
+  }
+
   /// Sets the current placement block to the given block.
   void recordMoveToDominator(Block *block) { placementBlock = block; }
 
   /// Sets the current placement block to the given block.
   void recordMoveToParent(Block *block) { recordMoveToDominator(block); }
 };
 
 /// A state implementation compatible with the `BufferAllocationHoisting` class
return BufferPlacementTransformationBase::isLoop(op) &&
            !dominators->dominates(aliasDominatorBlock, op->getBlock());
   }
 
   /// Returns true if the given operation should be considered for hoisting.
   static bool shouldHoistOpType(Operation *op) {
     return llvm::isa<memref::AllocOp, memref::AllocaOp>(op);
   }
 
+  /// Returns true if the `deallocOp` is a deallocation of a result of
+  /// `allocOp`. One of the operands of `deallocOp` is guaranteed to be a result
+  /// of the `allocOp`.
+  static bool isMatchingDealloc(Operation *allocOp, Operation *deallocOp) {
+    if (!llvm::isa<memref::DeallocOp>(deallocOp))
+      return false;
+    assert(llvm::isa<memref::AllocOp>(allocOp) &&
+           "no deallocs expected for alloca allocations");
+    return true;
+  }
+
   /// Does not change the internal placement block, as we want to move
   /// operations out of loops only.
   void recordMoveToDominator(Block *block) {}
 
   /// Sets the current placement block to the given block.
   void recordMoveToParent(Block *block) { placementBlock = block; }
 };
 

mlir/test/Dialect/Bufferization/Transforms/buffer-hoisting.mlir

 }
 
 // CHECK-NEXT: %{{.*}} = memref.alloc()
 // CHECK-NEXT: %{{.*}} = memref.alloc()
 // CHECK-NEXT: cf.cond_br
 
 // -----
 
-// Test Case: Invalid position of the DeallocOp. There is a user after
-// deallocation.
+// Test Case:
 //   bb0
 //  /   \
 // bb1  bb2 <- Initial position of AllocOp
 //  \   /
 //   bb3
 // BufferHoisting expected behavior: It should move the AllocOp to the entry
 // block.
 
 // CHECK-LABEL: func @moving_invalid_dealloc_op_complex
 func.func @moving_invalid_dealloc_op_complex(
   %cond: i1,
     %arg0: memref<2xf32>,
     %arg1: memref<2xf32>) {
   cf.cond_br %cond, ^bb1, ^bb2
 ^bb1:
   cf.br ^exit(%arg0 : memref<2xf32>)
 ^bb2:
   %1 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%1: memref<2xf32>)
-  memref.dealloc %1 : memref<2xf32>
   cf.br ^exit(%1 : memref<2xf32>)
 ^exit(%arg2: memref<2xf32>):
   test.copy(%arg2, %arg1) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 // CHECK-NEXT: %{{.*}} = memref.alloc()
 // CHECK-NEXT: cf.cond_br

mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir

-// RUN: mlir-opt -buffer-loop-hoisting -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -buffer-loop-hoisting -allow-unregistered-dialect -split-input-file %s | FileCheck %s
 
 // This file checks the behavior of BufferLoopHoisting pass for moving Alloc
 // operations in their correct positions.
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
func.func @hoist_alloca(
   }
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
 // CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
+
+// -----
+
+// CHECK-LABEL: func @hoist_dealloc
+func.func @hoist_dealloc(%lb: index, %ub: index, %step: index, %f: f32) {
+  scf.for %iv = %lb to %ub step %step {
+    %0 = memref.alloc() : memref<5xf32>
+    memref.store %f, %0[%iv] : memref<5xf32>
+    "print_memref"(%0) : (memref<5xf32>) -> ()
+    memref.dealloc %0 : memref<5xf32>
+  }
+  return
+}
+
+// CHECK: memref.alloc
+// CHECK: scf.for {{.*}} {
+// CHECK:   memref.store
+// CHECK: }
+// CHECK: memref.dealloc
+// CHECK: return
+
+// -----
+
+// CHECK-LABEL: func @unsupported_hoist_dealloc
+func.func @unsupported_hoist_dealloc(%lb: index, %ub: index, %step: index, %f: f32) {
+  scf.for %iv = %lb to %ub step %step {
+    %0 = memref.alloc() : memref<5xf32>
+    memref.store %f, %0[%iv] : memref<5xf32>
+    "print_memref"(%0) : (memref<5xf32>) -> ()
+    %c = "some_condition"() : () -> (i1)
+    scf.if %c {
+      // Hoisting not supported because the dealloc appears in a block that is
+      // different from the alloc block.
+      memref.dealloc %0 : memref<5xf32>
+    }
+  }
+  return
+}
+
+// CHECK: scf.for {{.*}} {
+// CHECK:   memref.alloc
+// CHECK:   memref.store
+// CHECK:   scf.if {{.*}} {
+// CHECK:     memref.dealloc
+// CHECK:   }
+// CHECK: }
+// CHECK-NOT: memref.dealloc
+// CHECK: return