Index: mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h
===================================================================
--- mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h
+++ mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.h
@@ -15,6 +15,17 @@
 
 #include "mlir/IR/Builders.h"
 
+namespace mlir {
+// Enum class representing different hoisting kinds for the allocation
+// operation
+enum class HoistingKind : std::uint8_t {
+  None = 0,       // No hoisting kind selected
+  Loop = 1 << 0,  // Indicates loop hoisting kind
+  Block = 1 << 1, // Indicates dominated block hoisting kind
+  LoopAndBlock = Loop | Block
+};
+} // namespace mlir
+
 #include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_ALLOCATIONOPINTERFACE_H_
Index: mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.td
===================================================================
--- mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.td
+++ mlir/include/mlir/Dialect/Bufferization/IR/AllocationOpInterface.td
@@ -51,6 +51,25 @@
       "::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; }]
     >
   ];
 }
Index: mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
===================================================================
--- mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+++ mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
@@ -11,6 +11,7 @@
 #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"
@@ -175,4 +176,5 @@
 void mlir::bufferization::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<BufferizationTransformDialectExtension>();
+  bufferization::registerAllocationOpInterfaceExternalModels(registry);
 }
Index: mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
===================================================================
--- mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
+++ mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp
@@ -644,6 +644,22 @@
     return builder.create<bufferization::CloneOp>(alloc.getLoc(), alloc)
         .getResult();
   }
+  static ::mlir::HoistingKind getHoistingKind() {
+    return HoistingKind::LoopAndBlock;
+  }
+  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
@@ -720,6 +736,8 @@
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
     memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+    memref::AllocaOp::attachInterface<
+        DefaultAutomaticAllocationHoistingInterface>(*ctx);
     memref::ReallocOp::attachInterface<DefaultReallocationInterface>(*ctx);
   });
 }
Index: mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
===================================================================
--- mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -13,6 +13,7 @@
 
 #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"
@@ -39,6 +40,22 @@
   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<std::uint8_t>(allocOp.getHoistingKind()) &
+                     static_cast<std::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<std::uint8_t>(allocOp.getHoistingKind()) &
+                     static_cast<std::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.
@@ -279,7 +296,7 @@
 
   /// 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.
@@ -316,7 +333,7 @@
 
   /// 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
@@ -356,13 +373,15 @@
       // `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());
-
-      // Replace the original alloc by a newly created alloca.
-      allocOp->replaceAllUsesWith(alloca);
-      allocOp->erase();
+      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();
+      }
     }
   }
 };