[mlir][bufferization] Drop equivalent buffer results - modify public functions option#177589
Merged
Merged
Conversation
Adds an option to allow modification of public functions to the drop equivalent buffer results API and the respective pass. llvm#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.
Member
|
@llvm/pr-subscribers-mlir-bufferization Author: Adam Siemieniuk (adam-smnk) ChangesAdds an option to allow modification of public functions to the drop equivalent buffer results API and the respective pass. #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 Full diff: https://github.com/llvm/llvm-project/pull/177589.diff 4 Files Affected:
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
|
Member
|
@llvm/pr-subscribers-mlir Author: Adam Siemieniuk (adam-smnk) ChangesAdds an option to allow modification of public functions to the drop equivalent buffer results API and the respective pass. #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 Full diff: https://github.com/llvm/llvm-project/pull/177589.diff 4 Files Affected:
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
|
Member
Author
|
Ping |
matthias-springer
approved these changes
Feb 2, 2026
rishabhmadan19
pushed a commit
to rishabhmadan19/llvm-project
that referenced
this pull request
Feb 9, 2026
…functions option (llvm#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. llvm#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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
#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
promoteBufferResultsToOutParamswhich also allows opt-in public function rewriting.