[mlir][bufferization] Add `buffer_loop_hoisting` transform op

[springerm]

Summary

This op hoists buffer allocation from loops.

[springerm]

clang-format

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td

 include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
 def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
 
 //===----------------------------------------------------------------------===//
+// BufferLoopHoistingOp
+//===----------------------------------------------------------------------===//
+
+def BufferLoopHoistingOp
+    : Op<Transform_Dialect, "bufferization.buffer_loop_hoisting",
+        [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+         TransformEachOpTrait, TransformOpInterface]> {
+  let description = [{
+    Hoist buffer allocations ("memref.alloc" and "memref.alloca") from loops
+    within the targeted op. This transform assumes that there are no buffer
+    deallocation ops in the IR.
+
+    This transform reads the `target` handle and modifies the payload.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs);
+  let assemblyFormat = "$target attr-dict `:` type($target)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
+        ::mlir::Operation *target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // OneShotBufferizeOp
 //===----------------------------------------------------------------------===//
 
 def OneShotBufferizeOp
     : Op<Transform_Dialect, "bufferization.one_shot_bufferize",
         [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
          DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let description = [{

mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h

 ///   on the reverse SSA use-def chain, starting from the OpOperand and always
 ///   following the aliasing  OpOperand, that eventually ends at a single
 ///   tensor::EmptyOp.
 LogicalResult eliminateEmptyTensors(RewriterBase &rewriter, Operation *op,
                                     OneShotAnalysisState &state,
                                     AnchorMatchFn anchorMatchFunc,
                                     RewriteFn rewriteFunc);
 
+/// Within the given operation, hoist buffers from loops where possible. See
+/// "BufferLoopHoistingPass" for more information.
+void hoistBuffersFromLoops(Operation *op);
+
 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
 LogicalResult insertSliceAnchoredEmptyTensorEliminationStep(
     RewriterBase &rewriter, Operation *op, OneShotAnalysisState &state);
 
 /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
 /// After applying this transform, the IR can be bufferized without inserting

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

 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/IR/FunctionInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::transform;
 
 //===----------------------------------------------------------------------===//
+// BufferLoopHoistingOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::BufferLoopHoistingOp::applyToOne(
+    TransformRewriter &rewriter, Operation *target,
+    ApplyToEachResultList &results, TransformState &state) {
+  bufferization::hoistBuffersFromLoops(target);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::BufferLoopHoistingOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  onlyReadsHandle(getTarget(), effects);
+  modifiesPayload(effects);
+}
+
+//===----------------------------------------------------------------------===//
 // OneShotBufferizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::OneShotBufferizeOp::apply(transform::TransformRewriter &rewriter,
                                      TransformResults &transformResults,
                                      TransformState &state) {
   OneShotBufferizationOptions options;

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/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"
 
 namespace mlir {
 namespace bufferization {

 
 /// The buffer loop hoisting pass that hoists allocation nodes out of loops.
 struct BufferLoopHoistingPass
     : public bufferization::impl::BufferLoopHoistingBase<
           BufferLoopHoistingPass> {
 
   void runOnOperation() override {
     // Hoist all allocations out of loops.
-    BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(
-        getOperation());
-    optimizer.hoist();
+    hoistBuffersFromLoops(getOperation());
   }
 };
 
 /// The promote buffer to stack pass that tries to convert alloc nodes into
 /// alloca nodes.
 class PromoteBuffersToStackPass
     : public bufferization::impl::PromoteBuffersToStackBase<
           PromoteBuffersToStackPass> {
public:
   }
 
 private:
   std::function<bool(Value)> isSmallAlloc;
 };
 
 } // namespace
 
+void mlir::bufferization::hoistBuffersFromLoops(Operation *op) {
+    BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(
+        op);
+    optimizer.hoist();
+}
+
 std::unique_ptr<Pass> mlir::bufferization::createBufferHoistingPass() {
   return std::make_unique<BufferHoistingPass>();
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createBufferLoopHoistingPass() {
   return std::make_unique<BufferLoopHoistingPass>();
 }
 

mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir

 func.func @empty_tensor_elimination(
     %t: tensor<10xf32>, %f: f32) -> tensor<10xf32> {
   %0 = tensor.empty() : tensor<5xf32>
   %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
   %2 = tensor.insert_slice %1 into %t [1][5][1]
       : tensor<5xf32> into tensor<10xf32>
   return %2 : tensor<10xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.bufferization.buffer_loop_hoisting %0 : !transform.any_op
+}
+
+// CHECK-LABEL: func @buffer_loop_hoisting(
+//       CHECK:   memref.alloca
+//       CHECK:   scf.for
+//       CHECK:     memref.store
+func.func @buffer_loop_hoisting(%lb: index, %ub: index, %step: index, %f: f32, %pos: index) {
+  scf.for %iv = %lb to %ub step %step {
+    %0 = memref.alloca() : memref<5xf32>
+    memref.store %f, %0[%pos] : memref<5xf32>
+  }
+  return
+}