[MLIR] Added support for dynamic shaped allocas to promote-buffers-to-stack pass.

[dfki-jugr]

Summary

Extended promote buffers to stack pass to support dynamically shaped allocas.
The conversion is limited by the rank of the underlying tensor.
An option is added to the pass to adjust the given rank.

[pifon2a]

mlir/include/mlir/Transforms/Passes.td

216

I would remove "Define the" from all these 3 options. Just

Maximal size in bytes to promote allocations to stack.
Bitwidth of the index type. Used for size estimation.
Maximal tensor rank to promote dynamic buffers.

mlir/lib/Transforms/BufferOptimizations.cpp

39–40

nit: please, write this in one line:
return type.getRank() < maxRankOfTensors; you might need a static_cast though

322

would it work with Operation* instead of auto? then you wouldn't need to write getOperation() below.

[dfki-jugr]

Addressed comments.
Introduce transformation if the alloc operands are RankOps.

[herhut]

Thanks. Just some nits.

mlir/include/mlir/Transforms/Passes.td

222

Can you describe this in more detail? Not all buffers that have less than this rank will be promoted but only "small" ones.

mlir/lib/Transforms/BufferOptimizations.cpp

35

I think we we should limit this to alloc where the defining op is an AllocOp as otherwise we likely do not want to do this transformation.

44

Would return operand.getDefiningOp<RankOp>() also work?

mlir/test/Transforms/promote-buffers-to-stack.mlir

112

Can you add a check where this is more than 3?

[dfki-jugr]

Addressed comments.
Changed default max rank to 1 instead of 2 and changed condition to less or equal.
This way is more intuitive.

[dfki-jugr]

mlir/test/Transforms/promote-buffers-to-stack.mlir

112

I added another test pass that checks this behaviour. An extra test for this is redundant and skipped.

[herhut]

mlir/include/mlir/Transforms/Passes.td

220

Maybe max-rank-of-allocated-memref?

mlir/lib/Transforms/BufferOptimizations.cpp

36

Should this be an ||?

[silvas]

mlir/lib/Transforms/BufferOptimizations.cpp

39

nit: I feel like this comment should read "if the dynamic dimension of the alloc is produced by a RankOp, then we know it is likely to be small. Also limit this to maxRankOfTensors to prevent multiple small values produced by RankOp from multiplying to an excessively large value"

[silvas]

mlir/test/Transforms/promote-buffers-to-stack.mlir

104

nit: why does this test need control flow? this patch seems like a very localized change unrelated to control flow.

[dfki-jugr]

Addressed comments.

mlir/include/mlir/Transforms/Passes.h

 std::unique_ptr<Pass> createBufferHoistingPass();
 
 /// Creates a pass that moves allocations upwards out of loops. This avoids
 /// reallocations inside of loops.
 std::unique_ptr<Pass> createBufferLoopHoistingPass();
 
 /// Creates a pass that promotes heap-based allocations to stack-based ones.
 /// Only buffers smaller than the provided size are promoted.
+/// Dynamic shaped buffers are promoted up to the given rank.
 std::unique_ptr<Pass>
 createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024,
-                                unsigned bitwidthOfIndexType = 64);
+                                unsigned bitwidthOfIndexType = 64,
+                                unsigned maxRankOfTensor = 2);
 
 /// Creates a pass that converts memref function results to out-params.
 std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
 
 /// Creates an instance of the Canonicalizer pass.
 std::unique_ptr<Pass> createCanonicalizerPass();
 
 /// Create a pass that removes unnecessary Copy operations.

mlir/include/mlir/Transforms/Passes.td

let description = [{
     This pass implements a simple algorithm to convert heap-based memory
     allocations to stack-based ones. It uses a built-in heuristic to decide
     whether it makes sense to convert an allocation.
   }];
   let constructor = "mlir::createPromoteBuffersToStackPass()";
   let options = [
     Option<"maxAllocSizeInBytes", "max-alloc-size-in-bytes", "unsigned",
            /*default=*/"1024",
-           "Define the maximum size in bytes to promote allocations to stack.">,
+           "Maximal size in bytes to promote allocations to stack.">,
     Option<"bitwidthOfIndexType", "bitwidth-of-index-type", "unsigned",
            /*default=*/"64",
-           "Define the bitwidth of the index type. Used for size estimation.">,
-
+           "Bitwidth of the index type. Used for size estimation.">,
+    Option<"maxRankOfTensors", "max-rank-of-tensors", "unsigned",
+           /*default=*/"2",
+           "Maximal tensor rank to promote dynamic buffers.">,
   ];
 }
 
 def BufferResultsToOutParams : Pass<"buffer-results-to-out-params", "ModuleOp">  {
   let summary = "Converts memref-typed function results to out-params";
   let description = [{
     Some calling conventions prefer to pass output memrefs as "out params". The
     conversion to this calling convention must be done as an atomic

mlir/lib/Transforms/BufferOptimizations.cpp

 static bool isKnownControlFlowInterface(Operation *op) {
   return isa<LoopLikeOpInterface, RegionBranchOpInterface>(op);
 }
 
 /// Check if the size of the allocation is less than the given size. The
 /// transformation is only applied to small buffers since large buffers could
 /// exceed the stack space.
 static bool isSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
-                         unsigned bitwidthOfIndexType) {
+                         unsigned bitwidthOfIndexType,
+                         unsigned maxRankOfTensors) {
   auto type = alloc.getType().dyn_cast<ShapedType>();
-  if (!type || !type.hasStaticShape())
+  if (!type)
     return false;
+  if (!type.hasStaticShape()) {
+    // Check if the dynamic shape dimension is below the allowed rank.
+    if (type.getRank() < maxRankOfTensors) {
+      // Check if all operands of an alloc are RankOp.
+      return llvm::all_of(alloc.getDefiningOp()->getOperands(),
+                          [&](Value operand) {
+                            Operation *operandOp = operand.getDefiningOp();
+                            return operandOp && llvm::isa<RankOp>(operandOp);
+                          });
+    }
+    return false;
+  }
   // For index types, use the provided size, as the type does not know.
   unsigned int bitwidth = type.getElementType().isIndex()
                               ? bitwidthOfIndexType
                               : type.getElementTypeBitWidth();
   return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8;
 }
 
 /// Checks whether the given aliases leave the allocation scope.

 
 /// Promotes heap-based allocations to stack-based allocations (if possible).
 class BufferPlacementPromotion : BufferPlacementTransformationBase {
 public:
   BufferPlacementPromotion(Operation *op)
       : BufferPlacementTransformationBase(op) {}
 
   /// Promote buffers to stack-based allocations.
-  void promote(unsigned maximumSize, unsigned bitwidthOfIndexType) {
+  void promote(unsigned maximumSize, unsigned bitwidthOfIndexType,
+               unsigned maxRankOfTensors) {
     for (BufferPlacementAllocs::AllocEntry &entry : allocs) {
       Value alloc = std::get<0>(entry);
       Operation *dealloc = std::get<1>(entry);
       // Checking several requirements to transform an AllocOp into an AllocaOp.
       // The transformation is done if the allocation is limited to a given
       // size. Furthermore, a deallocation must not be defined for this
       // allocation entry and a parent allocation scope must exist.
-      if (!isSmallAlloc(alloc, maximumSize, bitwidthOfIndexType) || dealloc ||
-          !hasAllocationScope(alloc, aliases))
+      if (!isSmallAlloc(alloc, maximumSize, bitwidthOfIndexType,
+                        maxRankOfTensors) ||
+          dealloc || !hasAllocationScope(alloc, aliases))
         continue;
 
       Operation *startOperation = BufferPlacementAllocs::getStartOperation(
           alloc, alloc.getParentBlock(), liveness);
       // Build a new alloca that is associated with its parent
       // `AutomaticAllocationScope` determined during the initialization phase.
       OpBuilder builder(startOperation);
-      auto alloca = builder.create<AllocaOp>(
-          alloc.getLoc(), alloc.getType().cast<MemRefType>());
+      Operation *allocOp = alloc.getDefiningOp();
+      Operation *alloca = builder.create<AllocaOp>(
+          alloc.getLoc(), alloc.getType().cast<MemRefType>(),
+          allocOp->getOperands());
 
       // Replace the original alloc by a newly created alloca.
-      Operation *allocOp = alloc.getDefiningOp();
-      allocOp->replaceAllUsesWith(alloca.getOperation());
+      allocOp->replaceAllUsesWith(alloca);
       allocOp->erase();
     }
   }
 };
 
 //===----------------------------------------------------------------------===//
 // BufferOptimizationPasses
 //===----------------------------------------------------------------------===//

 };
 
 /// The promote buffer to stack pass that tries to convert alloc nodes into
 /// alloca nodes.
 struct PromoteBuffersToStackPass
     : PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
 
   PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
-                            unsigned bitwidthOfIndexType) {
+                            unsigned bitwidthOfIndexType,
+                            unsigned maxRankOfTensors) {
     this->maxAllocSizeInBytes = maxAllocSizeInBytes;
     this->bitwidthOfIndexType = bitwidthOfIndexType;
+    this->maxRankOfTensors = maxRankOfTensors;
   }
 
   void runOnFunction() override {
     // Move all allocation nodes and convert candidates into allocas.
     BufferPlacementPromotion optimizer(getFunction());
-    optimizer.promote(this->maxAllocSizeInBytes, this->bitwidthOfIndexType);
+    optimizer.promote(this->maxAllocSizeInBytes, this->bitwidthOfIndexType,
+                      this->maxRankOfTensors);
   }
 };
 
 } // end anonymous namespace
 
 std::unique_ptr<Pass> mlir::createBufferHoistingPass() {
   return std::make_unique<BufferHoistingPass>();
 }
 
 std::unique_ptr<Pass> mlir::createBufferLoopHoistingPass() {
   return std::make_unique<BufferLoopHoistingPass>();
 }
 
 std::unique_ptr<Pass>
 mlir::createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
-                                      unsigned bitwidthOfIndexType) {
-  return std::make_unique<PromoteBuffersToStackPass>(maxAllocSizeInBytes,
-                                                     bitwidthOfIndexType);
+                                      unsigned bitwidthOfIndexType,
+                                      unsigned maxRankOfTensors) {
+  return std::make_unique<PromoteBuffersToStackPass>(
+      maxAllocSizeInBytes, bitwidthOfIndexType, maxRankOfTensors);
 }

mlir/test/Transforms/promote-buffers-to-stack.mlir

 // RUN: mlir-opt -promote-buffers-to-stack -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix DEFINDEX
 // RUN: mlir-opt -promote-buffers-to-stack="bitwidth-of-index-type=256 max-alloc-size-in-bytes=128" -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix BIGINDEX
 // RUN: mlir-opt -promote-buffers-to-stack="bitwidth-of-index-type=256 max-alloc-size-in-bytes=64" -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix LOWLIMIT
+// RUN: mlir-opt -promote-buffers-to-stack="max-rank-of-tensors=3" -split-input-file %s | FileCheck %s --check-prefix=CHECK --check-prefix RANK
 
 // This file checks the behavior of PromoteBuffersToStack pass for converting
 // AllocOps into AllocaOps, if possible.
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // PromoteBuffersToStack expected behavior: It should convert %0 into an
 // AllocaOp.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @condBranch
 func @condBranch(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   br ^bb3(%arg1 : memref<2xf32>)
 ^bb2:
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)

 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // PromoteBuffersToStack expected behavior:
 // Since the alloc has dynamic type, it is not converted into an alloca.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @condBranchDynamicType
 func @condBranchDynamicType(
   %arg0: i1,
   %arg1: memref<?xf32>,
   %arg2: memref<?xf32>,
   %arg3: index) {
   cond_br %arg0, ^bb1, ^bb2(%arg3: index)
 ^bb1:

 // CHECK-NEXT: test.buffer_based
 //      CHECK: br ^bb3
 // CHECK-NEXT: ^bb3(%[[ALLOC0:.*]]:{{.*}})
 //      CHECK: test.copy(%[[ALLOC0]],
 // CHECK-NEXT: return
 
 // -----
 
+// CHECK-LABEL: func @condBranchDynamicRanked
+func @condBranchDynamicRanked(
+  %arg0: i1,
+  %tensor: tensor<*xf32>) {
+  cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  br ^bb3
+^bb2:
+  %0 = rank %tensor : tensor<*xf32>
+  %1 = alloc(%0) : memref<?xindex>
+  br ^bb3
+^bb3:
+  return
+}
+
+// CHECK-NEXT: cond_br
+//      CHECK: ^bb2
+//      CHECK: ^bb2
+// CHECK-NEXT: %[[RANK:.*]] = rank
+// CHECK-NEXT: %[[ALLOCA:.*]] = alloca(%[[RANK]])
+//      CHECK: br ^bb3
+
+// -----
+
+// CHECK-LABEL: func @condBranchDynamicRanked2D
+func @condBranchDynamicRanked2D(
+  %arg0: i1,
+  %tensor: tensor<*xf32>) {
+  cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  br ^bb3
+^bb2:
+  %0 = rank %tensor : tensor<*xf32>
+  %1 = alloc(%0, %0) : memref<?x?xindex>
+  br ^bb3
+^bb3:
+  return
+}
+
+// CHECK-NEXT: cond_br
+//      CHECK: ^bb2
+//      CHECK: ^bb2
+// CHECK-NEXT: %[[RANK:.*]] = rank
+//  RANK-NEXT: %[[ALLOC:.*]] = alloca(%[[RANK]], %[[RANK]])
+// DEFINDEX-NEXT: %[[ALLOC:.*]] = alloc(%[[RANK]], %[[RANK]])
+//      CHECK: br ^bb3
+
+// -----
+
+// CHECK-LABEL: func @condBranchDynamicNoRank
+func @condBranchDynamicNoRank(
+  %arg0: i1,
+  %arg1: index,
+  %tensor: tensor<*xf32>) {
+  cond_br %arg0, ^bb1, ^bb2
+^bb1:
+  br ^bb3
+^bb2:
+  %0 = alloc(%arg1) : memref<?xindex>
+  br ^bb3
+^bb3:
+  return
+}
+
+// CHECK-NEXT: cond_br
+//      CHECK: ^bb2
+//      CHECK: ^bb2
+// CHECK-NEXT: %[[ALLOC:.*]] = alloc
+//      CHECK: br ^bb3
+
+// -----
+
 // Test Case: Existing AllocOp with no users.
 // PromoteBuffersToStack expected behavior: It should convert it to an
 // AllocaOp.
 
 // CHECK-LABEL: func @emptyUsesValue
 func @emptyUsesValue(%arg0: memref<4xf32>) {
   %0 = alloc() : memref<4xf32>
   return
 }
 // CHECK-NEXT: %[[ALLOCA:.*]] = alloca()
 // CHECK-NEXT: return
 
 // -----
 
 // Test Case:
 //    bb0
 //   /   \
 //  |    bb1 <- Initial position of AllocOp
 //   \   /
 //    bb2
 // PromoteBuffersToStack expected behavior: It should convert it into an
 // AllocaOp.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @criticalEdge
 func @criticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
 ^bb1:
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   br ^bb2(%0 : memref<2xf32>)
 ^bb2(%1: memref<2xf32>):

 // Test Case:
 //    bb0 <- Initial position of AllocOp
 //   /   \
 //  |    bb1
 //   \   /
 //    bb2
 // PromoteBuffersToStack expected behavior: It converts the alloc in an alloca.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @invCriticalEdge
 func @invCriticalEdge(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cond_br %arg0, ^bb1, ^bb2(%arg1 : memref<2xf32>)
 ^bb1:
   br ^bb2(%0 : memref<2xf32>)
 ^bb2(%1: memref<2xf32>):

 // Test Case:
 //    bb0 <- Initial position of the first AllocOp
 //   /   \
 //  bb1  bb2
 //   \   /
 //    bb3 <- Initial position of the second AllocOp
 // PromoteBuffersToStack expected behavior: It converts the allocs into allocas.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @ifElse
 func @ifElse(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):

 // Test Case: No users for buffer in if-else CFG
 //    bb0 <- Initial position of AllocOp
 //   /   \
 //  bb1  bb2
 //   \   /
 //    bb3
 // PromoteBuffersToStack expected behavior: It converts the alloc into alloca.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @ifElseNoUsers
 func @ifElseNoUsers(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):

 //    |     /  \
 //    |   bb3  bb4
 //    \     \  /
 //     \     /
 //       bb5 <- Initial position of the second AllocOp
 // PromoteBuffersToStack expected behavior: The two allocs should be converted
 // into allocas.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @ifElseNested
 func @ifElseNested(%arg0: i1, %arg1: memref<2xf32>, %arg2: memref<2xf32>) {
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>)
   cond_br %arg0,
     ^bb1(%arg1, %0 : memref<2xf32>, memref<2xf32>),
     ^bb2(%0, %arg1 : memref<2xf32>, memref<2xf32>)
 ^bb1(%1: memref<2xf32>, %2: memref<2xf32>):

 // CHECK-NEXT: return
 
 // -----
 
 // Test Case: Dead operations in a single block.
 // PromoteBuffersToStack expected behavior: It converts the two AllocOps into
 // allocas.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @redundantOperations
 func @redundantOperations(%arg0: memref<2xf32>) {
   %0 = alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)
   %1 = alloc() : memref<2xf32>
   test.buffer_based in(%0: memref<2xf32>) out(%1: memref<2xf32>)
   return
 }

 //                                     bb0
 //                                    /   \
 // Initial pos of the 1st AllocOp -> bb1  bb2 <- Initial pos of the 2nd AllocOp
 //                                    \   /
 //                                     bb3
 // PromoteBuffersToStack expected behavior: Both AllocOps are converted into
 // allocas.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @moving_alloc_and_inserting_missing_dealloc
 func @moving_alloc_and_inserting_missing_dealloc(
   %cond: i1,
     %arg0: memref<2xf32>,
     %arg1: memref<2xf32>) {
   cond_br %cond, ^bb1, ^bb2
 ^bb1:
   %0 = alloc() : memref<2xf32>

 
 // -----
 
 // Test Case: Nested regions - This test defines a BufferBasedOp inside the
 // region of a RegionBufferBasedOp.
 // PromoteBuffersToStack expected behavior: The AllocOps are converted into
 // allocas.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @nested_regions_and_cond_branch
 func @nested_regions_and_cond_branch(
   %arg0: i1,
   %arg1: memref<2xf32>,
   %arg2: memref<2xf32>) {
   cond_br %arg0, ^bb1, ^bb2
 ^bb1:
   br ^bb3(%arg1 : memref<2xf32>)

 
 // -----
 
 // Test Case: buffer deallocation escaping
 // PromoteBuffersToStack expected behavior: The first alloc is returned, so
 // there is no conversion allowed. The second alloc is converted, since it
 // only remains in the scope of the function.
 
-#map0 = affine_map<(d0) -> (d0)>
-
 // CHECK-LABEL: func @memref_in_function_results
 func @memref_in_function_results(
   %arg0: memref<5xf32>,
   %arg1: memref<10xf32>,
   %arg2: memref<5xf32>) -> (memref<10xf32>, memref<15xf32>) {
   %x = alloc() : memref<15xf32>
   %y = alloc() : memref<5xf32>
   test.buffer_based in(%arg0: memref<5xf32>) out(%y: memref<5xf32>)

 // CHECK-LABEL: func @indexElementType
 func @indexElementType() {
   %0 = alloc() : memref<4xindex>
   return
 }
 // DEFINDEX-NEXT: alloca()
 // BIGINDEX-NEXT: alloca()
 // LOWLIMIT-NEXT: alloc()
+// RANK-NEXT: alloca()
 // CHECK-NEXT: return