[Mlir-commits] [mlir] [mlir][bufferization] Drop equivalent buffer results - modify public functions option (PR #177589)
Adam Siemieniuk
llvmlistbot at llvm.org
Fri Jan 23 05:48:59 PST 2026
https://github.com/adam-smnk created https://github.com/llvm/llvm-project/pull/177589
Adds an option to allow modification of public functions to the drop equivalent buffer results API and the respective pass.
A new standalone test for the pass is also added.
https://github.com/llvm/llvm-project/pull/163001 modified the default behavior by disallowing rewriting public function. This PR preserves the new behavior and only adds an opt-in flag to reenable changing public function.
The extra flag aligns dropping equivalent function results with another bufferization API `promoteBufferResultsToOutParams` which also allows opt-in public function rewriting.
>From 6aa515571145737df34989da7674bdb71f33ac83 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Fri, 23 Jan 2026 14:27:07 +0100
Subject: [PATCH] [mlir][bufferization] Drop buffer results - modify public
func option
Adds an option to allow modification of public functions to the drop
equivalent buffer results API and the respective pass.
https://github.com/llvm/llvm-project/pull/163001 modified the default
behavior by disallowing rewriting public function.
This PR preserves the new default behavior and only adds an opt-in
flag to reenable the pass on public function.
The extra flag aligns dropping equivalent function results with
another bufferization API `promoteBufferResultsToOutParams` which
also allows opt-in public function rewriting.
---
.../Dialect/Bufferization/Transforms/Passes.h | 9 ++-
.../Bufferization/Transforms/Passes.td | 5 ++
.../DropEquivalentBufferResults.cpp | 24 ++++--
.../drop-equivalent-buffer-results.mlir | 76 +++++++++++++++++++
4 files changed, 108 insertions(+), 6 deletions(-)
create mode 100644 mlir/test/Transforms/drop-equivalent-buffer-results.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
index ea158914e445b..90857358437a1 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h
@@ -182,8 +182,15 @@ LogicalResult
promoteBufferResultsToOutParams(ModuleOp module,
const BufferResultsToOutParamsOpts &options);
+/// Options for dropping equivalent memref buffer results.
+struct DropBufferResultsOpts {
+ /// If true, signatures of public functions are modified.
+ bool modifyPublicFunctions = false;
+};
+
/// Drop all memref function results that are equivalent to a function argument.
-LogicalResult dropEquivalentBufferResults(ModuleOp module);
+LogicalResult dropEquivalentBufferResults(
+ ModuleOp module, DropBufferResultsOpts options = DropBufferResultsOpts());
/// Creates a pass that promotes heap-based allocations to stack-based ones.
/// Only buffers smaller with `isSmallAlloc(alloc) == true` are promoted.
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 1eb692586bcfc..cd28bd6cf73a5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -276,6 +276,11 @@ def DropEquivalentBufferResultsPass
Note: If a bbArg buffer is not returned directly but casted to beforehand,
the buffer is still considered equivalent.
}];
+ let options = [
+ Option<"modifyPublicFunctions", "modify-public-functions", "bool",
+ /*default=*/"false", "Modify function signatures of public "
+ "functions.">,
+ ];
let dependentDialects = ["memref::MemRefDialect"];
}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp b/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
index bc1799099de31..e724af312652d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/DropEquivalentBufferResults.cpp
@@ -74,8 +74,8 @@ static bool operandsEqualFuncArgument(ArrayRef<Value> operands,
return true;
}
-LogicalResult
-mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) {
+LogicalResult mlir::bufferization::dropEquivalentBufferResults(
+ ModuleOp module, DropBufferResultsOpts options) {
IRRewriter rewriter(module.getContext());
DenseMap<func::FuncOp, DenseSet<func::CallOp>> callerMap;
@@ -83,13 +83,18 @@ mlir::bufferization::dropEquivalentBufferResults(ModuleOp module) {
module.walk([&](func::CallOp callOp) {
if (func::FuncOp calledFunc =
dyn_cast_or_null<func::FuncOp>(callOp.resolveCallable())) {
- if (!calledFunc.isPublic() && !calledFunc.isExternal())
+ if (calledFunc.isPublic() && !options.modifyPublicFunctions)
+ return WalkResult::advance();
+ if (!calledFunc.isExternal())
callerMap[calledFunc].insert(callOp);
}
+ return WalkResult::advance();
});
for (auto funcOp : module.getOps<func::FuncOp>()) {
- if (funcOp.isExternal() || funcOp.isPublic())
+ if (funcOp.isPublic() && !options.modifyPublicFunctions)
+ continue;
+ if (funcOp.isExternal())
continue;
SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
if (returnOps.empty())
@@ -166,9 +171,18 @@ namespace {
struct DropEquivalentBufferResultsPass
: bufferization::impl::DropEquivalentBufferResultsPassBase<
DropEquivalentBufferResultsPass> {
+ using Base::Base;
+
void runOnOperation() override {
- if (failed(bufferization::dropEquivalentBufferResults(getOperation())))
+ // Convert pass options.
+ options.modifyPublicFunctions = modifyPublicFunctions;
+
+ if (failed(bufferization::dropEquivalentBufferResults(getOperation(),
+ options)))
return signalPassFailure();
}
+
+private:
+ bufferization::DropBufferResultsOpts options;
};
} // namespace
diff --git a/mlir/test/Transforms/drop-equivalent-buffer-results.mlir b/mlir/test/Transforms/drop-equivalent-buffer-results.mlir
new file mode 100644
index 0000000000000..0375d3f0b16eb
--- /dev/null
+++ b/mlir/test/Transforms/drop-equivalent-buffer-results.mlir
@@ -0,0 +1,76 @@
+// RUN: mlir-opt -drop-equivalent-buffer-results -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -drop-equivalent-buffer-results=modify-public-functions=1 -split-input-file %s | \
+// RUN: FileCheck %s --check-prefix=MODIFY-PUBLIC
+
+
+// CHECK-LABEL: func private @single_buffer_return({{.*}}) {
+// CHECK: return
+
+!type = memref<?xf32, strided<[?], offset: ?>>
+func.func private @single_buffer_return(%buf: !type, %val: f32, %idx: index) -> !type {
+ memref.store %val, %buf[%idx] : !type
+ return %buf : !type
+}
+
+// -----
+
+// CHECK-LABEL: func private @multiple_buffer_returns({{.*}}) {
+// CHECK: return
+
+!type = memref<?xf32, strided<[?], offset: ?>>
+!type1 = memref<?x?xf32>
+func.func private @multiple_buffer_returns(
+ %buf: !type, %buf1: !type1, %val: f32, %idx: index) -> (!type1, !type) {
+ memref.store %val, %buf[%idx] : !type
+ memref.store %val, %buf1[%idx, %idx] : !type1
+ return %buf1, %buf : !type1, !type
+}
+
+// -----
+
+// CHECK-LABEL: func private @multiple_mixed_returns({{.*}}) -> i32 {
+// CHECK: %[[CST:.+]] = arith.constant 1 : i32
+// CHECK: return %[[CST]] : i32
+
+!type = memref<?xf32, strided<[?], offset: ?>>
+!type1 = memref<?x?xf32>
+func.func private @multiple_mixed_returns(
+ %buf: !type, %buf1: !type1, %val: f32, %idx: index) -> (!type1, i32, !type) {
+ memref.store %val, %buf[%idx] : !type
+ memref.store %val, %buf1[%idx, %idx] : !type1
+ %cst = arith.constant 1 : i32
+ return %buf1, %cst, %buf : !type1, i32, !type
+}
+
+// -----
+
+// Ensure public functions remain unchanged by default.
+// CHECK-LABEL: func @public_function(
+// CHECK-SAME: %[[BUF:.+]]: memref<?xf32, strided<[?], offset: ?>>,
+// CHECK-SAME: ) -> memref<?xf32, strided<[?], offset: ?>> {
+// CHECK: return %[[BUF]]
+
+// When explicitly requested, public functions can be modified.
+// MODIFY-PUBLIC-LABEL: func @public_function(
+// MODIFY-PUBLIC-SAME: %[[BUF:.+]]: memref<?xf32, strided<[?], offset: ?>>,
+// MODIFY-PUBLIC-SAME: ) {
+// MODIFY-PUBLIC: return
+
+!type = memref<?xf32, strided<[?], offset: ?>>
+func.func @public_function(
+ %buf: !type, %val: f32, %idx: index) -> !type {
+ memref.store %val, %buf[%idx] : !type
+ return %buf : !type
+}
+
+// -----
+
+// CHECK-LABEL: func private @negative_external_function(
+// CHECK-SAME: -> memref<?xf32, strided<[?], offset: ?>>
+
+// Ensure external function remains unchanged.
+// MODIFY-PUBLIC-LABEL: func private @negative_external_function(
+// MODIFY-PUBLIC-SAME: -> memref<?xf32, strided<[?], offset: ?>>
+
+!type = memref<?xf32, strided<[?], offset: ?>>
+func.func private @negative_external_function(%arg0: !type) -> !type
More information about the Mlir-commits
mailing list