[Mlir-commits] [mlir] 891e769 - [mlir] Initialize `isSmallAlloc` in `initialize` method.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2021-02-12T10:28:58+01:00
New Revision: 891e769ab68f2cdbd105cb42f5c2f40051c70837

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

LOG: [mlir] Initialize `isSmallAlloc` in `initialize` method.

Added: 
    

Modified: 
    mlir/lib/Transforms/BufferOptimizations.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/BufferOptimizations.cpp b/mlir/lib/Transforms/BufferOptimizations.cpp
index 916007198b3e..6c1e1fe0c9cf 100644
--- a/mlir/lib/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Transforms/BufferOptimizations.cpp
@@ -371,9 +371,7 @@ class PromoteBuffersToStackPass
   explicit PromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc)
       : isSmallAlloc(std::move(isSmallAlloc)) {}
 
-  void runOnFunction() override {
-    // Move all allocation nodes and convert candidates into allocas.
-    BufferPlacementPromotion optimizer(getFunction());
+  LogicalResult initialize(MLIRContext* context) override {
     if (isSmallAlloc == nullptr) {
       isSmallAlloc = [=](Value alloc) {
         return defaultIsSmallAlloc(alloc, maxAllocSizeInBytes,
@@ -381,6 +379,12 @@ class PromoteBuffersToStackPass
                                    maxRankOfAllocatedMemRef);
       };
     }
+    return success();
+  }
+
+  void runOnFunction() override {
+    // Move all allocation nodes and convert candidates into allocas.
+    BufferPlacementPromotion optimizer(getFunction());
     optimizer.promote(isSmallAlloc);
   }