[Mlir-commits] [mlir] 98770ec - [mlir][bufferization] Add `buffer_loop_hoisting` transform op

[Matthias Springer]

Author: Matthias Springer
Date: 2023-07-14T17:09:38+02:00
New Revision: 98770ecd76bb9cb1067c78c502be3899a28e5f92

URL: https://github.com/llvm/llvm-project/commit/98770ecd76bb9cb1067c78c502be3899a28e5f92
DIFF: https://github.com/llvm/llvm-project/commit/98770ecd76bb9cb1067c78c502be3899a28e5f92.diff

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

This op hoists buffer allocation from loops.

Differential Revision: https://reviews.llvm.org/D155289

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
    mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h
    mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
    mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
    mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
index 0a32afd0e19fe9..46a95ad8db2a6d 100644
--- a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
@@ -19,6 +19,35 @@ 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
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h
index d964be91f9626b..a0cfc811a0b50a 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h
@@ -41,6 +41,10 @@ LogicalResult eliminateEmptyTensors(RewriterBase &rewriter, Operation *op,
                                     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).

diff  --git a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
index f866484f485678..f1233739524d28 100644
--- a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
@@ -22,6 +22,23 @@ 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
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 278664abcf49d2..2a2baf39235250 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -14,6 +14,7 @@
 #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"
@@ -389,9 +390,7 @@ struct BufferLoopHoistingPass
 
   void runOnOperation() override {
     // Hoist all allocations out of loops.
-    BufferAllocationHoisting<BufferAllocationLoopHoistingState> optimizer(
-        getOperation());
-    optimizer.hoist();
+    hoistBuffersFromLoops(getOperation());
   }
 };
 
@@ -432,6 +431,11 @@ class PromoteBuffersToStackPass
 
 } // 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>();
 }

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 94550c8d4374a5..7c8bb422f533d7 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -180,3 +180,23 @@ func.func @empty_tensor_elimination(
       : 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
+}