[mlir][bufferization] Define a pipeline for buffer deallocation

[maerhart]

Summary

Since buffer deallocation requires a few passes to be run in a somewhat fixed
sequence, it makes sense to have a pipeline for convenience (and to reduce the
number of transform ops to represent default deallocation).

Depends on D158979

[maerhart]

rebase

[maerhart]

rebase

[maerhart]

add pipelines subdirectory to CMakeLists

[maerhart]

fix flag passed to expand reallocs

[springerm]

Note: This has landed as 08b7a71bcc951b9449a7c6462ec5f72c0522d67d.

mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h

+//===- Passes.h - Bufferization pipeline entry points -----------*- C++ -*-===//
+//
+// 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 header file defines prototypes of all bufferization pipelines.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
+#define MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H
+
+#include "mlir/Pass/PassOptions.h"
+
+namespace mlir {
+namespace bufferization {
+
+/// Options for the buffer deallocation pipeline.
+struct BufferDeallocationPipelineOptions
+    : public PassPipelineOptions<BufferDeallocationPipelineOptions> {
+  PassOptions::Option<bool> privateFunctionDynamicOwnership{
+      *this, "private-function-dynamic-ownership",
+      llvm::cl::desc(
+          "Allows to add additional arguments to private functions to "
+          "dynamically pass ownership of memrefs to callees. This can enable "
+          "earlier deallocations."),
+      llvm::cl::init(false)};
+};
+
+//===----------------------------------------------------------------------===//
+// Building and Registering.
+//===----------------------------------------------------------------------===//
+
+/// Adds the buffer deallocation pipeline to the `OpPassManager`. This
+/// is the standard pipeline for deallocating the MemRefs introduced by the
+/// One-Shot bufferization pass.
+void buildBufferDeallocationPipeline(
+    OpPassManager &pm, const BufferDeallocationPipelineOptions &options);
+
+/// Registers all pipelines for the `bufferization` dialect. Currently,
+/// this includes only the "buffer-deallocation-pipeline".
+void registerBufferizationPipelines();
+
+} // namespace bufferization
+} // namespace mlir
+
+#endif // MLIR_DIALECT_BUFFERIZATION_PIPELINES_PASSES_H

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

 // Passes
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
 
 /// Creates an instance of the BufferDeallocation pass to free all allocated
 /// buffers.
-std::unique_ptr<Pass> createBufferDeallocationPass();
+std::unique_ptr<Pass>
+createBufferDeallocationPass(bool privateFuncDynamicOwnership = false);
 
 /// Creates a pass that optimizes `bufferization.dealloc` operations. For
 /// example, it reduces the number of alias checks needed at runtime using
 /// static alias analysis.
 std::unique_ptr<Pass> createBufferDeallocationSimplificationPass();
 
 /// Creates an instance of the LowerDeallocations pass to lower
 /// `bufferization.dealloc` operations to the `memref` dialect.

mlir/include/mlir/InitAllPasses.h

 #define MLIR_INITALLPASSES_H_
 
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/AMDGPU/Transforms/Passes.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h"
 #include "mlir/Dialect/Async/Passes.h"
+#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/NVGPU/Transforms/Passes.h"
inline void registerAllPasses() {
   spirv::registerSPIRVPasses();
   tensor::registerTensorPasses();
   tosa::registerTosaOptPasses();
   transform::registerTransformPasses();
   vector::registerVectorPasses();
   arm_sme::registerArmSMEPasses();
 
   // Dialect pipelines
+  bufferization::registerBufferizationPipelines();
   sparse_tensor::registerSparseTensorPipelines();
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLPASSES_H_

mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp

+//===- BufferizationPipelines.cpp - Pipelines 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
+
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+
+//===----------------------------------------------------------------------===//
+// Pipeline implementation.
+//===----------------------------------------------------------------------===//
+
+void mlir::bufferization::buildBufferDeallocationPipeline(
+    OpPassManager &pm, const BufferDeallocationPipelineOptions &options) {
+  pm.addNestedPass<func::FuncOp>(
+      memref::createExpandReallocPass(/*doNotDealloc=*/true));
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  pm.addNestedPass<func::FuncOp>(createBufferDeallocationPass(
+      options.privateFunctionDynamicOwnership.getValue()));
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  pm.addNestedPass<func::FuncOp>(createBufferDeallocationSimplificationPass());
+  pm.addPass(createLowerDeallocationsPass());
+  pm.addNestedPass<func::FuncOp>(createCSEPass());
+  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+}
+
+//===----------------------------------------------------------------------===//
+// Pipeline registration.
+//===----------------------------------------------------------------------===//
+
+void mlir::bufferization::registerBufferizationPipelines() {
+  PassPipelineRegistration<BufferDeallocationPipelineOptions>(
+      "buffer-deallocation-pipeline",
+      "The default pipeline for automatically inserting deallocation "
+      "operations after one-shot bufferization. Deallocation operations "
+      "(except `memref.realloc`) may not be present already.",
+      buildBufferDeallocationPipeline);
+}

mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt

+add_mlir_dialect_library(MLIRBufferizationPipelines
+  BufferizationPipelines.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
+
+  LINK_LIBS PUBLIC
+  MLIRBufferizationTransforms
+  MLIRMemRefTransforms
+  MLIRFuncDialect
+  MLIRPass
+  MLIRTransforms
+)

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

 namespace {
 
 /// The actual buffer deallocation pass that inserts and moves dealloc nodes
 /// into the right positions. Furthermore, it inserts additional clones if
 /// necessary. It uses the algorithm described at the top of the file.
 struct BufferDeallocationPass
     : public bufferization::impl::BufferDeallocationBase<
           BufferDeallocationPass> {
+  BufferDeallocationPass(bool privateFuncDynamicOwnership)
+      : bufferization::impl::BufferDeallocationBase<BufferDeallocationPass>() {
+    this->privateFuncDynamicOwnership.setValue(privateFuncDynamicOwnership);
+  }
   void runOnOperation() override {
     func::FuncOp func = getOperation();
     if (func.isExternal())
       return;
 
     if (failed(deallocateBuffers(func, privateFuncDynamicOwnership)))
       signalPassFailure();
   }
bufferization::deallocateBuffers(FunctionOpInterface op,
   // Place all required temporary clone and dealloc nodes.
   return deallocation.deallocate(op);
 }
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
-std::unique_ptr<Pass> mlir::bufferization::createBufferDeallocationPass() {
-  return std::make_unique<BufferDeallocationPass>();
+std::unique_ptr<Pass> mlir::bufferization::createBufferDeallocationPass(
+    bool privateFuncDynamicOwnership) {
+  return std::make_unique<BufferDeallocationPass>(privateFuncDynamicOwnership);
 }

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-branchop-interface.mlir

 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation \
 // RUN:  -buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
 // Test Case:
 //    bb0
 //   /   \
 //  bb1  bb2 <- Initial position of AllocOp
 //   \   /
 //    bb3
 // BufferDeallocation expected behavior: bb2 contains an AllocOp which is
 // passed to bb3. In the latter block, there should be a deallocation.

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-callop-interface.mlir

 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=false \
 // RUN:   -buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=true \
 // RUN:   --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
+
 func.func private @f(%arg0: memref<f64>) -> memref<f64> {
   return %arg0 : memref<f64>
 }
 
 func.func @function_call() {
   %alloc = memref.alloc() : memref<f64>
   %alloc2 = memref.alloc() : memref<f64>
   %ret = call @f(%alloc) : (memref<f64>) -> memref<f64>

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-existing-deallocs.mlir

 // RUN: mlir-opt -verify-diagnostics -expand-realloc=emit-deallocs=false -buffer-deallocation \
 // RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
 func.func @auto_dealloc() {
   %c10 = arith.constant 10 : index
   %c100 = arith.constant 100 : index
   %alloc = memref.alloc(%c10) : memref<?xi32>
   %realloc = memref.realloc %alloc(%c100) : memref<?xi32> to memref<?xi32>
   "test.memref_user"(%realloc) : (memref<?xi32>) -> ()
   return
 }

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-function-boundaries.mlir

 // RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=false \
 // RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 // RUN: mlir-opt --allow-unregistered-dialect -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=true \
 // RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s --check-prefix=CHECK-DYNAMIC
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+// RUN: mlir-opt %s -buffer-deallocation-pipeline=private-function-dynamic-ownership --split-input-file > /dev/null
+
 // Test Case: Existing AllocOp with no users.
 // BufferDeallocation expected behavior: It should insert a DeallocOp right
 // before ReturnOp.
 
 func.func private @emptyUsesValue(%arg0: memref<4xf32>) {
   %0 = memref.alloc() : memref<4xf32>
   "test.memref_user"(%0) : (memref<4xf32>) -> ()
   return

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-memoryeffect-interface.mlir

 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation \
 // RUN:   --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
 // Test Case: Dead operations in a single block.
 // BufferDeallocation expected behavior: It only inserts the two missing
 // DeallocOps after the last BufferBasedOp.
 
 // CHECK-LABEL: func @redundantOperations
 func.func @redundantOperations(%arg0: memref<2xf32>) {
   %0 = memref.alloc() : memref<2xf32>
   test.buffer_based in(%arg0: memref<2xf32>) out(%0: memref<2xf32>)

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-region-branchop-interface.mlir

 // RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -buffer-deallocation \
 // RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file --verify-diagnostics > /dev/null
+
 // Test Case: Nested regions - This test defines a BufferBasedOp inside the
 // region of a RegionBufferBasedOp.
 // BufferDeallocation expected behavior: The AllocOp for the BufferBasedOp
 // should remain inside the region of the RegionBufferBasedOp and it should insert
 // the missing DeallocOp in the same region. The missing DeallocOp should be
 // inserted after CopyOp.
 
 func.func @nested_regions_and_cond_branch(

mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-subviews.mlir

 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation \
 // RUN:   --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -verify-diagnostics -buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null
 
+// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file > /dev/null
+
 // CHECK-LABEL: func @subview
 func.func @subview(%arg0 : index, %arg1 : index, %arg2 : memref<?x?xf32>) {
   %0 = memref.alloc() : memref<64x4xf32, strided<[4, 1], offset: 0>>
   %1 = memref.subview %0[%arg0, %arg1][%arg0, %arg1][%arg0, %arg1] :
     memref<64x4xf32, strided<[4, 1], offset: 0>>
   to memref<?x?xf32, strided<[?, ?], offset: ?>>
   test.copy(%1, %arg2) :
     (memref<?x?xf32, strided<[?, ?], offset: ?>>, memref<?x?xf32>)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":ArmSMETransforms",
         ":ArmSVEDialect",
         ":ArmSVETransforms",
         ":AsyncDialect",
         ":AsyncPassIncGen",
         ":AsyncToLLVM",
         ":AsyncTransforms",
         ":BufferizationDialect",
+        ":BufferizationPipelines",
         ":BufferizationTransformOps",
         ":BufferizationTransforms",
         ":CastInterfaces",
         ":ComplexDialect",
         ":ComplexToLLVM",
         ":ComplexToLibm",
         ":ComplexToSPIRV",
         ":ControlFlowDialect",
deps = [
         ":MemRefDialect",
         ":Pass",
         ":SCFDialect",
         ":Support",
         ":Transforms",
     ],
 )
 
+cc_library(
+    name = "BufferizationPipelines",
+    srcs = glob(["lib/Dialect/Bufferization/Pipelines/*.cpp"]),
+    hdrs = ["include/mlir/Dialect/Bufferization/Pipelines/Passes.h"],
+    includes = ["include"],
+    deps = [
+        ":BufferizationToMemRef",
+        ":BufferizationTransforms",
+        ":FuncDialect",
+        ":MemRefTransforms",
+        ":Pass",
+        ":Transforms",
+    ],
+)
+
 td_library(
     name = "DLTIDialectTdFiles",
     srcs = [
         "include/mlir/Dialect/DLTI/DLTI.td",
         "include/mlir/Dialect/DLTI/DLTIBase.td",
     ],
     includes = ["include"],
     deps = [":OpBaseTdFiles"],