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

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.

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

mlir/include/mlir/Transforms/Passes.h

@@ -41,8 +41,10 @@ std::unique_ptr<Pass> createBufferHoistingPass();
 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();

mlir/include/mlir/Transforms/Passes.td

@@ -214,6 +214,10 @@ def PromoteBuffersToStack : FunctionPass<"promote-buffers-to-stack"> {
     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.">,
+
   ];
 }
 

mlir/lib/Transforms/BufferOptimizations.cpp

@@ -29,11 +29,16 @@ static bool isKnownControlFlowInterface(Operation *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.
@@ -281,14 +286,15 @@ public:
       : 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;
 
@@ -340,17 +346,17 @@ struct BufferLoopHoistingPass : BufferLoopHoistingBase<BufferLoopHoistingPass> {
 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
@@ -364,6 +370,8 @@ std::unique_ptr<Pass> mlir::createBufferLoopHoistingPass() {
 }
 
 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

@@ -1,4 +1,6 @@
-// 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.
@@ -566,3 +568,19 @@ func @large_buffer_allocation(%arg0: memref<2048xf32>) {
 
 // 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