[Mlir-commits] [mlir] 4a77110 - [mlir][bufferize] Fix buffer promotion to stack for index types
Stephan Herhut
llvmlistbot at llvm.org
Fri Nov 13 00:30:16 PST 2020
Author: Stephan Herhut
Date: 2020-11-13T09:23:36+01:00
New Revision: 4a771108ac5ceeae6119185ae52ab6f5e6af2da1
URL: https://github.com/llvm/llvm-project/commit/4a771108ac5ceeae6119185ae52ab6f5e6af2da1
DIFF: https://github.com/llvm/llvm-project/commit/4a771108ac5ceeae6119185ae52ab6f5e6af2da1.diff
LOG: [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
Added:
Modified:
mlir/include/mlir/Transforms/Passes.h
mlir/include/mlir/Transforms/Passes.td
mlir/lib/Transforms/BufferOptimizations.cpp
mlir/test/Transforms/promote-buffers-to-stack.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index c9b90d15e79d..2e3437a46611 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/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();
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 3fcbb9a15087..da4ca24db499 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/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.">,
+
];
}
diff --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp
index 0a8802e439ec..8629b3395629 100644
--- a/mlir/lib/Transforms/BufferOptimizations.cpp
+++ b/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 @@ class BufferPlacementPromotion : BufferPlacementTransformationBase {
: 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);
}
diff --git a/mlir/test/Transforms/promote-buffers-to-stack.mlir b/mlir/test/Transforms/promote-buffers-to-stack.mlir
index 33f7c1a1a302..903b1f8cb748 100644
--- a/mlir/test/Transforms/promote-buffers-to-stack.mlir
+++ b/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
More information about the Mlir-commits
mailing list