[Mlir-commits] [mlir] b70366c - [mlir][BufferOptimization] Use datalayout instead of a flag to find index size

[Benjamin Kramer]

Author: Benjamin Kramer
Date: 2022-01-27T13:50:29+01:00
New Revision: b70366c9c430e1eadd59d5a1dfbb9c4d84f83de5

URL: https://github.com/llvm/llvm-project/commit/b70366c9c430e1eadd59d5a1dfbb9c4d84f83de5
DIFF: https://github.com/llvm/llvm-project/commit/b70366c9c430e1eadd59d5a1dfbb9c4d84f83de5.diff

LOG: [mlir][BufferOptimization] Use datalayout instead of a flag to find index size

This has the additional advantage of supporting more types.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
    mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
    mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
    mlir/test/Transforms/promote-buffers-to-stack.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index 481eaa284d98b..29c0d7f741d93 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -34,7 +34,6 @@ std::unique_ptr<OperationPass<FuncOp>> createFinalizingBufferizePass();
 /// 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.

diff  --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index dcbdd52a5c81d..9890383931cbb 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -164,9 +164,6 @@ def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "FuncOp"> {
     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.">,

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 158c786739342..141c513c90cf7 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -32,7 +32,6 @@ static bool isKnownControlFlowInterface(Operation *op) {
 /// 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>())
@@ -51,10 +50,8 @@ static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
     }
     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;
 }
 
@@ -390,10 +387,8 @@ class PromoteBuffersToStackPass
     : public PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
 public:
   PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
-                            unsigned bitwidthOfIndexType,
                             unsigned maxRankOfAllocatedMemRef) {
     this->maxAllocSizeInBytes = maxAllocSizeInBytes;
-    this->bitwidthOfIndexType = bitwidthOfIndexType;
     this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
   }
 
@@ -404,7 +399,6 @@ class PromoteBuffersToStackPass
     if (isSmallAlloc == nullptr) {
       isSmallAlloc = [=](Value alloc) {
         return defaultIsSmallAlloc(alloc, maxAllocSizeInBytes,
-                                   bitwidthOfIndexType,
                                    maxRankOfAllocatedMemRef);
       };
     }
@@ -432,10 +426,9 @@ std::unique_ptr<Pass> mlir::bufferization::createBufferLoopHoistingPass() {
 }
 
 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(

diff  --git a/mlir/test/Transforms/promote-buffers-to-stack.mlir b/mlir/test/Transforms/promote-buffers-to-stack.mlir
index 2b6cd3185fa11..5dfe9ccf8fad7 100644
--- a/mlir/test/Transforms/promote-buffers-to-stack.mlir
+++ b/mlir/test/Transforms/promote-buffers-to-stack.mlir
@@ -1,6 +1,5 @@
 // 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
@@ -595,7 +594,20 @@ func @indexElementType() {
   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