diff --git a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h @@ -0,0 +1,50 @@ +//===- 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 { + PassOptions::Option 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 diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h @@ -26,7 +26,8 @@ /// Creates an instance of the BufferDeallocation pass to free all allocated /// buffers. -std::unique_ptr createBufferDeallocationPass(); +std::unique_ptr +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 diff --git a/mlir/include/mlir/InitAllPasses.h b/mlir/include/mlir/InitAllPasses.h --- a/mlir/include/mlir/InitAllPasses.h +++ b/mlir/include/mlir/InitAllPasses.h @@ -20,6 +20,7 @@ #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" @@ -81,6 +82,7 @@ arm_sme::registerArmSMEPasses(); // Dialect pipelines + bufferization::registerBufferizationPipelines(); sparse_tensor::registerSparseTensorPipelines(); } diff --git a/mlir/lib/Dialect/Bufferization/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/CMakeLists.txt --- a/mlir/lib/Dialect/Bufferization/CMakeLists.txt +++ b/mlir/lib/Dialect/Bufferization/CMakeLists.txt @@ -1,3 +1,4 @@ add_subdirectory(IR) +add_subdirectory(Pipelines) add_subdirectory(TransformOps) add_subdirectory(Transforms) diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Bufferization/Pipelines/BufferizationPipelines.cpp @@ -0,0 +1,46 @@ +//===- 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( + memref::createExpandReallocPass(/*emitDeallocs=*/false)); + pm.addNestedPass(createCanonicalizerPass()); + pm.addNestedPass(createBufferDeallocationPass( + options.privateFunctionDynamicOwnership.getValue())); + pm.addNestedPass(createCanonicalizerPass()); + pm.addNestedPass(createBufferDeallocationSimplificationPass()); + pm.addPass(createLowerDeallocationsPass()); + pm.addNestedPass(createCSEPass()); + pm.addNestedPass(createCanonicalizerPass()); +} + +//===----------------------------------------------------------------------===// +// Pipeline registration. +//===----------------------------------------------------------------------===// + +void mlir::bufferization::registerBufferizationPipelines() { + PassPipelineRegistration( + "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); +} diff --git a/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/Bufferization/Pipelines/CMakeLists.txt @@ -0,0 +1,13 @@ +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 +) diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp @@ -994,6 +994,10 @@ struct BufferDeallocationPass : public bufferization::impl::BufferDeallocationBase< BufferDeallocationPass> { + BufferDeallocationPass(bool privateFuncDynamicOwnership) + : bufferization::impl::BufferDeallocationBase() { + this->privateFuncDynamicOwnership.setValue(privateFuncDynamicOwnership); + } void runOnOperation() override { func::FuncOp func = getOperation(); if (func.isExternal()) @@ -1024,6 +1028,7 @@ // BufferDeallocationPass construction //===----------------------------------------------------------------------===// -std::unique_ptr mlir::bufferization::createBufferDeallocationPass() { - return std::make_unique(); +std::unique_ptr mlir::bufferization::createBufferDeallocationPass( + bool privateFuncDynamicOwnership) { + return std::make_unique(privateFuncDynamicOwnership); } diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-branchop-interface.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-branchop-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-branchop-interface.mlir @@ -2,6 +2,8 @@ // 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 // / \ diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-callop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-callop-interface.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-callop-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-callop-interface.mlir @@ -3,6 +3,9 @@ // 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) -> memref { return %arg0 : memref } diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-existing-deallocs.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-existing-deallocs.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-existing-deallocs.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-existing-deallocs.mlir @@ -1,6 +1,8 @@ // 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 diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-function-boundaries.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-function-boundaries.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-function-boundaries.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-function-boundaries.mlir @@ -3,6 +3,9 @@ // 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. diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-memoryeffect-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-memoryeffect-interface.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-memoryeffect-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-memoryeffect-interface.mlir @@ -2,6 +2,8 @@ // 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. diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-region-branchop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-region-branchop-interface.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-region-branchop-interface.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-region-branchop-interface.mlir @@ -2,6 +2,8 @@ // 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 diff --git a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-subviews.mlir b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-subviews.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-subviews.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/BufferDeallocation/dealloc-subviews.mlir @@ -2,6 +2,8 @@ // 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) { %0 = memref.alloc() : memref<64x4xf32, strided<[4, 1], offset: 0>> diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -8253,6 +8253,7 @@ ":AsyncToLLVM", ":AsyncTransforms", ":BufferizationDialect", + ":BufferizationPipelines", ":BufferizationTransformOps", ":BufferizationTransforms", ":CastInterfaces", @@ -12184,6 +12185,21 @@ ], ) +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 = [