[Mlir-commits] [mlir] 6054bfa - [mlir] Support buffer hoisting on allocas

[Tres Popp]

Author: Tres Popp
Date: 2021-05-25T14:50:01+02:00
New Revision: 6054bfa8132ac06a140e1896d8cdd8dc79a01bd8

URL: https://github.com/llvm/llvm-project/commit/6054bfa8132ac06a140e1896d8cdd8dc79a01bd8
DIFF: https://github.com/llvm/llvm-project/commit/6054bfa8132ac06a140e1896d8cdd8dc79a01bd8.diff

LOG: [mlir] Support buffer hoisting on allocas

This adds support for hoisting allocas in both BufferHoisting and
BufferLoopHoisting.

Differential Revision: https://reviews.llvm.org/D102681

Added: 
    

Modified: 
    mlir/lib/Transforms/BufferOptimizations.cpp
    mlir/test/Transforms/buffer-loop-hoisting.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp
index 7d5b84db668b..b5dad2b5d393 100644
--- a/mlir/lib/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Transforms/BufferOptimizations.cpp
@@ -125,12 +125,19 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
 public:
   BufferAllocationHoisting(Operation *op)
       : BufferPlacementTransformationBase(op), dominators(op),
-        postDominators(op) {}
+        postDominators(op), scopeOp(op) {}
 
   /// Moves allocations upwards.
   void hoist() {
-    for (BufferPlacementAllocs::AllocEntry &entry : allocs) {
-      Value allocValue = std::get<0>(entry);
+    SmallVector<Value> allocsAndAllocas;
+    for (BufferPlacementAllocs::AllocEntry &entry : allocs)
+      allocsAndAllocas.push_back(std::get<0>(entry));
+    scopeOp->walk(
+        [&](memref::AllocaOp op) { allocsAndAllocas.push_back(op.memref()); });
+
+    for (auto allocValue : allocsAndAllocas) {
+      if (!StateT::shouldHoistOpType(allocValue.getDefiningOp()))
+        continue;
       Operation *definingOp = allocValue.getDefiningOp();
       assert(definingOp && "No defining op");
       auto operands = definingOp->getOperands();
@@ -222,6 +229,10 @@ class BufferAllocationHoisting : public BufferPlacementTransformationBase {
 
   /// The map storing the final placement blocks of a given alloc value.
   llvm::DenseMap<Value, Block *> placementBlocks;
+
+  /// The operation that this transformation is working on. It is used to also
+  /// gather allocas.
+  Operation *scopeOp;
 };
 
 /// A state implementation compatible with the `BufferAllocationHoisting` class
@@ -249,6 +260,11 @@ struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
     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);
+  }
+
   /// Sets the current placement block to the given block.
   void recordMoveToDominator(Block *block) { placementBlock = block; }
 
@@ -281,6 +297,11 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
            !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);
+  }
+
   /// Does not change the internal placement block, as we want to move
   /// operations out of loops only.
   void recordMoveToDominator(Block *block) {}

diff  --git a/mlir/test/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Transforms/buffer-loop-hoisting.mlir
index c5e727b89449..1b810b217d9f 100644
--- a/mlir/test/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Transforms/buffer-loop-hoisting.mlir
@@ -458,3 +458,32 @@ func @partial_hoist_multiple_loop_dependency(
 // CHECK-NEXT: {{.*}} = scf.for
 // CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc({{.*}})
 // CHECK-NEXT: {{.*}} = scf.for
+
+// -----
+
+// Test with allocas to ensure that op is also considered.
+
+// CHECK-LABEL: func @hoist_alloca
+func @hoist_alloca(
+  %lb: index,
+  %ub: index,
+  %step: index,
+  %buf: memref<2xf32>,
+  %res: memref<2xf32>) {
+  %0 = memref.alloca() : 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 = memref.alloca() : memref<2xf32>
+        scf.yield %0 : memref<2xf32>
+    }
+    scf.yield %0 : memref<2xf32>
+  }
+  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
+  return
+}
+
+//      CHECK: %[[ALLOCA0:.*]] = memref.alloca({{.*}})
+// CHECK-NEXT: %[[ALLOCA1:.*]] = memref.alloca({{.*}})
+// CHECK-NEXT: {{.*}} = scf.for