[mlir] Support buffer hoisting on allocas

[tpopp]

Summary

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

[herhut]

mlir/test/Transforms/buffer-hoisting.mlir

555 ↗

(On Diff #346091)

For the non-loop case, why would the allocation be hoisted? The reason this is done is to avoid copies but alloca allocated memrefs are never copied anyway, as they have implicit lifetimes.

[tpopp]

mlir/test/Transforms/buffer-hoisting.mlir

555 ↗

(On Diff #346091)

This happened practically because the hoisting code is shared and I was assuming it might also hoist during unstructured control flow loops. As that is not the case, I'll add details to the subclassed portions to only work on Allocs for this case.

[tpopp]

Remove hoisting of allocas in non-loop cases.

mlir/lib/Transforms/BufferOptimizations.cpp

 };
 
 /// Implements the actual hoisting logic for allocation nodes.
 template <typename StateT>
 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();
       auto resultAliases = aliases.resolve(allocValue);
       // Determine the common dominator block of all aliases.
       Block *dominatorBlock =
           findCommonDominator(allocValue, resultAliases, dominators);
       // Init the initial hoisting state.
private:
   DominanceInfo dominators;
 
   /// The post dominator info to move the dependent allocs in the right
   /// position.
   PostDominanceInfo postDominators;
 
   /// 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
 /// that hoists allocations into dominator blocks while keeping them inside of
 /// loops.
 struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
   using BufferAllocationHoistingStateBase::BufferAllocationHoistingStateBase;
 
return dominators->properlyDominates(dominatorBlock, dependencyBlock)
                : dominatorBlock;
   }
 
   /// Returns true if the given operation does not represent a loop.
   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);
+  }
+
   /// 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
struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
   /// aliases caused the `aliasDominatorBlock` to be "above" the block of the
   /// given loop operation. If this is the case, it indicates that the
   /// allocation is passed via a back edge.
   bool isLegalPlacement(Operation *op) {
     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);
+  }
+
   /// 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/Transforms/buffer-loop-hoisting.mlir

func @partial_hoist_multiple_loop_dependency(
   test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
   return
 }
 
 //      CHECK: %[[ALLOC0:.*]] = memref.alloc({{.*}})
 // 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