Make buffer hoisting/promotion passes use AllocationOpInterface

[xiaoleis]

Summary

This update implements the usage of AllocationOpInterface in the buffer hoisting/promotion passes. Two interface methods, namely getHoistingKind and buildPromotedAlloc, have been added. The former indicates which kind of hoisting (loop, block) an allocation operation supports, while the latter builds a stack allocation operation for promotable allocations used by the promote-buffers-to-stack pass.

This update makes these passes be functional for user customized allocation operation.

[springerm]

mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h

25

I think we usually don't have this in the enum and instead write Loop | Block where HoistingKind is used.

[xiaoleis]

This comment was removed by xiaoleis.

[xiaoleis]

Remove LoopAndBlock from HoistingKind as suggested

[xiaoleis]

I do not have commit access, can anyone help me with this? Thanks

mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h

25

Sure, this make sense.

[MaskRay]

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

649

In other places, we just write uint8_t and omit std::

[xiaoleis]

Remove the std:: prefix for uint8_t

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

649

Thanks, std:: is omitted.

[xiaoleis]

Omit std:: for uint8_t in declaration of HoistingKind

[xiaoleis]

Further formatting on std::uint8_t

mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_
 #define MLIR_DIALECT_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_
 
 #include "mlir/IR/Builders.h"
 
+namespace mlir {
+// Enum class representing different hoisting kinds for the allocation
+// operation
+enum class HoistingKind : uint8_t {
+  None = 0,      // No hoisting kind selected
+  Loop = 1 << 0, // Indicates loop hoisting kind
+  Block = 1 << 1 // Indicates dominated block hoisting kind
+};
+} // namespace mlir
+
 #include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_

mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.td

StaticInterfaceMethod<[{
         allocation value (which refers to the current Op implementing this
         interface). The allocation value is a result of the current operation
         implementing this interface. If there is no compatible clone operation,
         this method can return ::std::nullopt.
       }],
       "::std::optional<::mlir::Value>", "buildClone",
       (ins "::mlir::OpBuilder&":$builder, "::mlir::Value":$alloc), [{}],
       /*defaultImplementation=*/[{ return std::nullopt; }]
+    >,
+    StaticInterfaceMethod<[{
+        Returns the kind of hoisting supported for the buffer allocated by this
+        operation.
+      }],
+      "::mlir::HoistingKind", "getHoistingKind",
+      (ins), [{}],
+      /*defaultImplementation=*/[{ return HoistingKind::None; }]
+    >,
+    StaticInterfaceMethod<[{
+        Builds a stack allocation operation using the provided builder and the
+        current allocation value (which refers to the current Op implementing this
+        interface). The allocation value is a result of the current
+        operation implementing this interface. If there is no compatible
+        stack allocation operation, this method can return ::std::nullopt.
+      }],
+      "::std::optional<::mlir::Operation*>", "buildPromotedAlloc",
+      (ins "::mlir::OpBuilder&":$builder, "::mlir::Value":$alloc), [{}],
+      /*defaultImplementation=*/[{ return std::nullopt; }]
     >
   ];
 }
 
 #endif  // ALLOCATION_OP_INTERFACE

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

 //===- BufferizationTransformOps.h - Bufferization transform ops ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/IR/FunctionInterfaces.h"
 
 using namespace mlir;

 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizationEnums.cpp.inc"
 
 void mlir::bufferization::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<BufferizationTransformDialectExtension>();
+  bufferization::registerAllocationOpInterfaceExternalModels(registry);
 }

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

static std::optional<Operation *> buildDealloc(OpBuilder &builder,
                                                  Value alloc) {
     return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
         .getOperation();
   }
   static std::optional<Value> buildClone(OpBuilder &builder, Value alloc) {
     return builder.create<bufferization::CloneOp>(alloc.getLoc(), alloc)
         .getResult();
   }
+  static ::mlir::HoistingKind getHoistingKind() {
+    return static_cast<HoistingKind>(static_cast<uint8_t>(HoistingKind::Loop) |
+                                     static_cast<uint8_t>(HoistingKind::Block));
+  }
+  static ::std::optional<::mlir::Operation *>
+  buildPromotedAlloc(OpBuilder &builder, Value alloc) {
+    Operation *definingOp = alloc.getDefiningOp();
+    return builder.create<memref::AllocaOp>(
+        definingOp->getLoc(), cast<MemRefType>(definingOp->getResultTypes()[0]),
+        definingOp->getOperands(), definingOp->getAttrs());
+  }
+};
+
+struct DefaultAutomaticAllocationHoistingInterface
+    : public bufferization::AllocationOpInterface::ExternalModel<
+          DefaultAutomaticAllocationHoistingInterface, memref::AllocaOp> {
+  static ::mlir::HoistingKind getHoistingKind() { return HoistingKind::Loop; }
 };
 
 struct DefaultReallocationInterface
     : public bufferization::AllocationOpInterface::ExternalModel<
           DefaultAllocationInterface, memref::ReallocOp> {
   static std::optional<Operation *> buildDealloc(OpBuilder &builder,
                                                  Value realloc) {
     return builder.create<memref::DeallocOp>(realloc.getLoc(), realloc)
LogicalResult bufferization::deallocateBuffers(Operation *op) {
 
   return success();
 }
 
 void bufferization::registerAllocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
     memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+    memref::AllocaOp::attachInterface<
+        DefaultAutomaticAllocationHoistingInterface>(*ctx);
     memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
   });
 }
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::bufferization::createBufferDeallocationPass() {
   return std::make_unique<BufferDeallocationPass>();
 }

mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp

 //===- BufferOptimizations.cpp - pre-pass optimizations for bufferization -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements logic for three optimization passes. The first two
 // passes try to move alloc nodes out of blocks to reduce the number of
 // allocations and copies during buffer deallocation. The third pass tries to
 // convert heap-based allocations to stack-based allocations, if possible.
 
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 
+#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"
 #include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Pass/Pass.h"
 

 using namespace mlir::bufferization;
 
 /// Returns true if the given operation implements a known high-level region-
 /// based control-flow interface.
 static bool isKnownControlFlowInterface(Operation *op) {
   return isa<LoopLikeOpInterface, RegionBranchOpInterface>(op);
 }
 
+/// Returns true if the given operation implements the AllocationOpInterface
+/// and it supports the dominate block hoisting.
+static bool allowAllocDominateBlockHoisting(Operation *op) {
+  auto allocOp = dyn_cast<AllocationOpInterface>(op);
+  return allocOp && (static_cast<uint8_t>(allocOp.getHoistingKind()) &
+                     static_cast<uint8_t>(HoistingKind::Block));
+}
+
+/// Returns true if the given operation implements the AllocationOpInterface
+/// and it supports the loop hoisting.
+static bool allowAllocLoopHoisting(Operation *op) {
+  auto allocOp = dyn_cast<AllocationOpInterface>(op);
+  return allocOp && (static_cast<uint8_t>(allocOp.getHoistingKind()) &
+                     static_cast<uint8_t>(HoistingKind::Loop));
+}
+
 /// 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 maxRankOfAllocatedMemRef) {
   auto type = dyn_cast<ShapedType>(alloc.getType());
   if (!type || !alloc.getDefiningOp<memref::AllocOp>())
     return false;
struct BufferAllocationHoistingState : BufferAllocationHoistingStateBase {
 
   /// Returns true if the given operation does not represent a loop.
   bool isLegalPlacement(Operation *op) {
     return !BufferPlacementTransformationBase::isLoop(op);
   }
 
   /// Returns true if the given operation should be considered for hoisting.
   static bool shouldHoistOpType(Operation *op) {
-    return llvm::isa<memref::AllocOp>(op);
+    return allowAllocDominateBlockHoisting(op);
   }
 
   /// Sets the current placement block to the given block.
   void recordMoveToDominator(Block *block) { placementBlock = block; }
 
   /// Sets the current placement block to the given block.
   void recordMoveToParent(Block *block) { recordMoveToDominator(block); }
 };
struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
   /// allocation is passed via a back edge.
   bool isLegalPlacement(Operation *op) {
     return BufferPlacementTransformationBase::isLoop(op) &&
            !dominators->dominates(aliasDominatorBlock, op->getBlock());
   }
 
   /// Returns true if the given operation should be considered for hoisting.
   static bool shouldHoistOpType(Operation *op) {
-    return llvm::isa<memref::AllocOp, memref::AllocaOp>(op);
+    return allowAllocLoopHoisting(op);
   }
 
   /// Does not change the internal placement block, as we want to move
   /// operations out of loops only.
   void recordMoveToDominator(Block *block) {}
 
   /// Sets the current placement block to the given block.
   void recordMoveToParent(Block *block) { placementBlock = block; }
for (BufferPlacementAllocs::AllocEntry &entry : allocs) {
         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();
-      Operation *alloca = builder.create<memref::AllocaOp>(
-          alloc.getLoc(), cast<MemRefType>(alloc.getType()),
-          allocOp->getOperands(), allocOp->getAttrs());
-
+      if (auto allocInterface = dyn_cast<AllocationOpInterface>(allocOp)) {
+        Operation *alloca =
+            allocInterface.buildPromotedAlloc(builder, alloc).value();
+        if (!alloca)
+          continue;
         // Replace the original alloc by a newly created alloca.
         allocOp->replaceAllUsesWith(alloca);
         allocOp->erase();
       }
     }
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // BufferOptimizationPasses
 //===----------------------------------------------------------------------===//
 
 /// The buffer hoisting pass that hoists allocation nodes into dominating
 /// blocks.