[mlir] Allow users of `PromoteBuffersToStackPass` to customize `small buffer` func.

[rriddle]

mlir/lib/Transforms/BufferOptimizations.cpp

302

Use function_ref here instead.

377

This looks like something that should be computed during initialization.

388

The = nullptr is unnecessary.

[pifon2a]

Add explicit

[pifon2a]

Address River's comments.

[pifon2a]

mlir/lib/Transforms/BufferOptimizations.cpp

377

I tried it before, but then the command-line options won't work.

[rriddle]

mlir/lib/Transforms/BufferOptimizations.cpp

377

You can either override initialize or override initializeOptions and do it there.

[pifon2a]

mlir/lib/Transforms/BufferOptimizations.cpp

377

addressed in https://github.com/llvm/llvm-project/commit/891e769ab68f2cdbd105cb42f5c2f40051c70837

mlir/include/mlir/Transforms/Passes.h

 /// 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);
+
 /// Creates a pass that finalizes a partial bufferization by removing remaining
 /// tensor_load and tensor_to_memref operations.
 std::unique_ptr<FunctionPass> createFinalizingBufferizePass();
 
 /// Creates a pass that converts memref function results to out-params.
 std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
 
 /// Creates an instance of the Canonicalizer pass.

mlir/lib/Transforms/BufferOptimizations.cpp

 /// based control-flow interface.
 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 isSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
+static bool defaultIsSmallAlloc(Value alloc, unsigned maximumSizeInBytes,
                                 unsigned bitwidthOfIndexType,
                                 unsigned maxRankOfAllocatedMemRef) {
   auto type = alloc.getType().dyn_cast<ShapedType>();
   if (!type || !alloc.getDefiningOp<AllocOp>())
     return false;
   if (!type.hasStaticShape()) {
     // Check if the dynamic shape dimension of the alloc is produced by RankOp.
     // 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.

 
 /// Promotes heap-based allocations to stack-based allocations (if possible).
 class BufferPlacementPromotion : BufferPlacementTransformationBase {
 public:
   BufferPlacementPromotion(Operation *op)
       : BufferPlacementTransformationBase(op) {}
 
   /// Promote buffers to stack-based allocations.
-  void promote(unsigned maximumSize, unsigned bitwidthOfIndexType,
-               unsigned maxRankOfAllocatedMemRef) {
+  void promote(const std::function<bool(Value)> &isSmallAlloc) {
     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, bitwidthOfIndexType,
-                        maxRankOfAllocatedMemRef) ||
-          dealloc || !hasAllocationScope(alloc, aliases))
+      if (!isSmallAlloc(alloc) || dealloc ||
+          !hasAllocationScope(alloc, aliases))
         continue;
 
       Operation *startOperation = BufferPlacementAllocs::getStartOperation(
           alloc, alloc.getParentBlock(), liveness);
       // Build a new alloca that is associated with its parent
       // `AutomaticAllocationScope` determined during the initialization phase.
       OpBuilder builder(startOperation);
       Operation *allocOp = alloc.getDefiningOp();
void runOnFunction() override {
     BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(
         getFunction());
     optimizer.hoist();
   }
 };
 
 /// The promote buffer to stack pass that tries to convert alloc nodes into
 /// alloca nodes.
-struct PromoteBuffersToStackPass
-    : PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
-
+class PromoteBuffersToStackPass
+    : public PromoteBuffersToStackBase<PromoteBuffersToStackPass> {
+public:
   PromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
                             unsigned bitwidthOfIndexType,
                             unsigned maxRankOfAllocatedMemRef) {
     this->maxAllocSizeInBytes = maxAllocSizeInBytes;
     this->bitwidthOfIndexType = bitwidthOfIndexType;
     this->maxRankOfAllocatedMemRef = maxRankOfAllocatedMemRef;
   }
 
+  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());
-    optimizer.promote(this->maxAllocSizeInBytes, this->bitwidthOfIndexType,
-                      this->maxRankOfAllocatedMemRef);
+    if (isSmallAlloc == nullptr) {
+      isSmallAlloc = [=](Value alloc) {
+        return defaultIsSmallAlloc(alloc, maxAllocSizeInBytes,
+                                   bitwidthOfIndexType,
+                                   maxRankOfAllocatedMemRef);
+      };
+    }
+    optimizer.promote(isSmallAlloc);
   }
+
+private:
+  std::function<bool(Value)> isSmallAlloc = nullptr;
 };
 
 } // end anonymous namespace
 
 std::unique_ptr<Pass> mlir::createBufferHoistingPass() {
   return std::make_unique<BufferHoistingPass>();
 }
 
 std::unique_ptr<Pass> mlir::createBufferLoopHoistingPass() {
   return std::make_unique<BufferLoopHoistingPass>();
 }
 
 std::unique_ptr<Pass>
 mlir::createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes,
                                       unsigned bitwidthOfIndexType,
                                       unsigned maxRankOfAllocatedMemRef) {
   return std::make_unique<PromoteBuffersToStackPass>(
       maxAllocSizeInBytes, bitwidthOfIndexType, maxRankOfAllocatedMemRef);
 }
+
+std::unique_ptr<Pass>
+mlir::createPromoteBuffersToStackPass(std::function<bool(Value)> isSmallAlloc) {
+  return std::make_unique<PromoteBuffersToStackPass>(std::move(isSmallAlloc));
+}