[mlir][BufferOptimiization] Use datalayout instead of a flag to find index size

[bkramer]

Summary

This has the additional advantage of supporting more types.

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

 /// bufferization.to_tensor and bufferization.to_memref operations.
 std::unique_ptr<OperationPass<FuncOp>> createFinalizingBufferizePass();
 
 /// 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 maxRankOfAllocatedMemRef = 1);
 
 /// Creates a pass that promotes heap-based allocations to stack-based ones.
 /// Only buffers smaller with `isSmallAlloc(alloc) == true` are promoted.
 std::unique_ptr<Pass>
 createPromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc);
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

let description = [{
     shaped buffers that are limited by the rank of the tensor can be
     converted. They are only transformed if they are considered to be small.
   }];
   let constructor = "mlir::bufferization::createPromoteBuffersToStackPass()";
   let options = [
     Option<"maxAllocSizeInBytes", "max-alloc-size-in-bytes", "unsigned",
            /*default=*/"1024",
            "Maximal size in bytes to promote allocations to stack.">,
-    Option<"bitwidthOfIndexType", "bitwidth-of-index-type", "unsigned",
-           /*default=*/"64",
-           "Bitwidth of the index type. Used for size estimation.">,
     Option<"maxRankOfAllocatedMemRef", "max-rank-of-allocated-memref", "unsigned",
            /*default=*/"1",
            "Maximal memref rank to promote dynamic buffers.">,
   ];
 }
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES

mlir/lib/Dialect/Bufferization/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 defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
-                                unsigned bitwidthOfIndexType,
                                 unsigned maxRankOfAllocatedMemRef) {
   auto type = alloc.getType().dyn_cast<ShapedType>();
   if (!type || !alloc.getDefiningOp<memref::AllocOp>())
     return false;
   if (!type.hasStaticShape()) {
     // Check if the dynamic shape dimension of the alloc is produced by
     // `memref.rank`. If this is the case, it is likely to be small.
     // Furthermore, the dimension is limited to the maximum rank of the
     // allocated memref to avoid large values by multiplying several small
     // values.
     if (type.getRank() <= maxRankOfAllocatedMemRef) {
       return llvm::all_of(alloc.getDefiningOp()->getOperands(),
                           [&](Value operand) {
                             return operand.getDefiningOp<memref::RankOp>();
                           });
     }
     return false;
   }
-  // For index types, use the provided size, as the type does not know.
-  unsigned int bitwidth = type.getElementType().isIndex()
-                              ? bitwidthOfIndexType
-                              : type.getElementTypeBitWidth();
+  unsigned bitwidth = mlir::DataLayout::closest(alloc.getDefiningOp())
+                          .getTypeSizeInBits(type.getElementType());
   return type.getNumElements() * bitwidth <= maximumSizeInBytes * 8;
 }
 
 /// Checks whether the given aliases leave the allocation scope.
 static bool
 leavesAllocationScope(Region *parentRegion,
                       const BufferViewFlowAnalysis::ValueSetT &aliases) {
   for (Value alias : aliases) {

 };
 
 /// The promote buffer to stack pass that tries to convert alloc nodes into
 /// alloca nodes.
 class PromoteBuffersToStackPass
     : public PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
 public:
   PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
-                            unsigned bitwidthOfIndexType,
                             unsigned maxRankOfAllocatedMemRef) {
     this->maxAllocSizeInBytes = maxAllocSizeInBytes;
-    this->bitwidthOfIndexType = bitwidthOfIndexType;
     this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
   }
 
   explicit PromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc)
       : isSmallAlloc(std::move(isSmallAlloc)) {}
 
   LogicalResult initialize(MLIRContext *context) override {
     if (isSmallAlloc == nullptr) {
       isSmallAlloc = [=](Value alloc) {
         return defaultIsSmallAlloc(alloc, maxAllocSizeInBytes,
-                                   bitwidthOfIndexType,
                                    maxRankOfAllocatedMemRef);
       };
     }
     return success();
   }
 
   void runOnOperation() override {
     // Move all allocation nodes and convert candidates into allocas.
std::unique_ptr<Pass> mlir::bufferization::createBufferHoistingPass() {
   return std::make_unique<BufferHoistingPass>();
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createBufferLoopHoistingPass() {
   return std::make_unique<BufferLoopHoistingPass>();
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createPromoteBuffersToStackPass(
-    unsigned maxAllocSizeInBytes, unsigned bitwidthOfIndexType,
-    unsigned maxRankOfAllocatedMemRef) {
-  return std::make_unique<PromoteBuffersToStackPass>(
-      maxAllocSizeInBytes, bitwidthOfIndexType, maxRankOfAllocatedMemRef);
+    unsigned maxAllocSizeInBytes, unsigned maxRankOfAllocatedMemRef) {
+  return std::make_unique<PromoteBuffersToStackPass>(maxAllocSizeInBytes,
+                                                     maxRankOfAllocatedMemRef);
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createPromoteBuffersToStackPass(
     std::function<bool(Value)> isSmallAlloc) {
   return std::make_unique<PromoteBuffersToStackPass>(std::move(isSmallAlloc));
 }

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-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-allocated-memref=2" -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
 //   /   \

 // AllocaOp.
 
 // CHECK-LABEL: func @indexElementType
 func @indexElementType() {
   %0 = memref.alloc() : memref<4xindex>
   return
 }
 // DEFINDEX-NEXT: memref.alloca()
-// BIGINDEX-NEXT: memref.alloca()
+// LOWLIMIT-NEXT: memref.alloca()
+// RANK-NEXT: memref.alloca()
+// CHECK-NEXT: return
+
+// -----
+
+// CHECK-LABEL: func @bigIndexElementType
+module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 256>>} {
+  func @bigIndexElementType() {
+    %0 = memref.alloc() : memref<4xindex>
+    return
+  }
+}
+// DEFINDEX-NEXT: memref.alloca()
 // LOWLIMIT-NEXT: memref.alloc()
 // RANK-NEXT: memref.alloca()
 // CHECK-NEXT: return