[mlir][bufferize] Fix buffer promotion to stack for index types

[herhut]

Summary

The index type does not have a bitsize and hence the size of corresponding allocations cannot be computed. Instead, the promotion pass now has an explicit option to specify the size of index.

[dfki-jugr]

mlir/lib/Transforms/BufferOptimizations.cpp

40

Just a question about the naming: bitWidth or bitwidth? It seems that there is no consistency in the naming in the whole project. Wouldn't it make sense to have only one name?

mlir/include/mlir/Transforms/Passes.h

 /// required copies that are inserted during the BufferDeallocation pass.
 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.
 std::unique_ptr<Pass>
-createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024);
+createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024,
+                                unsigned bitwidthOfIndexType = 64);
 
 /// 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 = [{
     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.">,
+    Option<"bitwidthOfIndexType", "bitwidth-of-index-type", "unsigned",
+           /*default=*/"64",
+           "Define the bitwidth of the index type. Used for size estimation.">,
+
   ];
 }
 
 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

 /// based control-flow interface.
 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) {
+static bool isSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
+                         unsigned bitwidthOfIndexType) {
   auto type = alloc.getType().dyn_cast<ShapedType>();
   if (!type || !type.hasStaticShape())
     return false;
-  return type.getSizeInBits() < maximumSizeInBytes * 8;
+  // 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.
 static bool
 leavesAllocationScope(Region *parentRegion,
                       const BufferAliasAnalysis::ValueSetT &aliases) {
   for (Value alias : aliases) {
     for (auto *use : alias.getUsers()) {

 
 /// 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) {
+  void promote(unsigned maximumSize, unsigned bitwidthOfIndexType) {
     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) || std::get<1>(entry) ||
+      if (!isSmallAlloc(alloc, maximumSize, bitwidthOfIndexType) || 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);
struct BufferLoopHoistingPass : BufferLoopHoistingBase<BufferLoopHoistingPass> {
   }
 };
 
 /// The promote buffer to stack pass that tries to convert alloc nodes into
 /// alloca nodes.
 struct PromoteBuffersToStackPass
     : PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
 
-  PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes)
-      : maximumSize(maxAllocSizeInBytes) {}
+  PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
+                            unsigned bitwidthOfIndexType) {
+    this->maxAllocSizeInBytes = maxAllocSizeInBytes;
+    this->bitwidthOfIndexType = bitwidthOfIndexType;
+  }
 
   void runOnFunction() override {
     // Move all allocation nodes and convert candidates into allocas.
     BufferPlacementPromotion optimizer(getFunction());
-    optimizer.promote(maximumSize);
+    optimizer.promote(this->maxAllocSizeInBytes, this->bitwidthOfIndexType);
   }
-
-private:
-  const unsigned maximumSize;
 };
 
 } // 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) {
-  return std::make_unique<PromoteBuffersToStackPass>(maxAllocSizeInBytes);
+mlir::createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
+                                      unsigned bitwidthOfIndexType) {
+  return std::make_unique<PromoteBuffersToStackPass>(maxAllocSizeInBytes,
+                                                     bitwidthOfIndexType);
 }

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

-// RUN: mlir-opt -promote-buffers-to-stack -split-input-file %s | FileCheck %s
+// 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
 
 // This file checks the behavior of PromoteBuffersToStack pass for converting
 // AllocOps into AllocaOps, if possible.
 
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp

 func @large_buffer_allocation(%arg0: memref<2048xf32>) {
   %0 = alloc() : memref<2048xf32>
   test.copy(%0, %arg0) : (memref<2048xf32>, memref<2048xf32>)
   return
 }
 
 // CHECK-NEXT: %[[ALLOC:.*]] = alloc()
 // CHECK-NEXT: test.copy
+
+// -----
+
+// Test Case: AllocOp with element type index.
+// PromoteBuffersToStack expected behavior: It should convert it to an
+// AllocaOp.
+
+// CHECK-LABEL: func @indexElementType
+func @indexElementType() {
+  %0 = alloc() : memref<4xindex>
+  return
+}
+// DEFINDEX-NEXT: alloca()
+// BIGINDEX-NEXT: alloca()
+// LOWLIMIT-NEXT: alloc()
+// CHECK-NEXT: return