Enable custom alloc-like ops in promoteBufferResultsToOutParams
#120288
Conversation
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
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-bufferization Author: None (srcarroll) ChangesIn Full diff: https://github.com/llvm/llvm-project/pull/120288.diff 1 Files Affected:
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index b7755b2be8483b..b4d2d6b0c5da8f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -6,6 +6,7 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -121,7 +122,8 @@ static LogicalResult updateReturnOps(func::FuncOp func,
OpBuilder builder(op);
for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
if (hoistStaticAllocs &&
- isa_and_nonnull<memref::AllocOp>(orig.getDefiningOp()) &&
+ isa_and_nonnull<bufferization::AllocationOpInterface>(
+ orig.getDefiningOp()) &&
mlir::cast<MemRefType>(orig.getType()).hasStaticShape()) {
orig.replaceAllUsesWith(arg);
orig.getDefiningOp()->erase();
|
|
i think in addition to this, we should also enable custom alloc op in |
srcarroll
changed the title
AllocationOpInterface with hoistStaticAllocs optionpromoteBufferResultsToOutParams
@matthias-springer i ended up doing this, but if you disagree i can always revert. |
| @@ -167,6 +174,10 @@ struct BufferResultsToOutParamsOpts { | |||
| return true; | |||
| }; | |||
|
|
|||
| /// Allocation function; used to allocate a memref. | |||
| /// If this is empty, memref.alloc is used | |||
| std::optional<AllocationFn> allocationFn; | |||
There was a problem hiding this comment.
Could this be avoided by adding a buildAlloc interface method to AllocationOpInterface? Maybe the transformation could also supported mixed alloc ops then>
There was a problem hiding this comment.
I'm not sure that would work actually. wouldn't you need an instance of an AllocationOpInterface to call the method? In this context there is no such instance, since we are creating a new one. Maybe i'm missing something. Also I do think I prefer this since it mirrors alloc and copy customization in bufferization.
There was a problem hiding this comment.
Oh, you only have a memref SSA value, right?
There was a problem hiding this comment.
yah, this is used in updateCalls for allocating at the call site and uses the arg types to allocate
matthias-springer
approved these changes
Dec 20, 2024
| @@ -224,7 +232,13 @@ LogicalResult mlir::bufferization::promoteBufferResultsToOutParams( | |||
| return failure(); | |||
| } | |||
| } | |||
| if (failed(updateCalls(module, options))) | |||
| auto defaultAllocationFn = [](OpBuilder &builder, Location loc, | |||
There was a problem hiding this comment.
Maybe we could give keep this as a default value in the struct, then the field doesn't have to be optional and you can pass the options to updateCalls. Not sure if it's possible without including MemRefOps.h in Passes.h though. (Maybe we can split the header?)
There was a problem hiding this comment.
sure that sounds good to me. i think i would want to do that with the copy op too
There was a problem hiding this comment.
i was just kinda following what had already been done with copy. i think they should be treated the same in this context either way
There was a problem hiding this comment.
Not sure if it's possible without including MemRefOps.h in Passes.h though. (Maybe we can split the header?)
If you are opposed to including MemRefOps.h I could just forward declare the appropriate ops if you prefer. I have no particular preference and honestly don't know what's best practice in this case.
There was a problem hiding this comment.
actually, no. i don't think forward declaration would work here
There was a problem hiding this comment.
i went ahead and just included MemRef.h for now. Passes.h was already included. Naively, I think this is ok since the default use of this would have to link the memref dialect anyway if I'm not mistaken. But let me know if this is no good
srcarroll
force-pushed
the
buffer-results-to-out-params-generalize
branch
from
f868c2a to
cc788e0
Compare
|
@matthias-springer i barely noticed you approved, but I think it was before my latest commit so just wanted to make sure it's still good to go. thanks! |
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/94/builds/3216 Here is the relevant piece of the build log for the reference |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
In
buffer-results-to-out-params, whenhoist-static-allocsoption is enabled the pass was looking formemref.allocs in order to attempt to avoid copies when it can. Which makes it not extensible to external ops that have allocation like properties. This patch simply changesmemref::AllocOptoAllocationOpInterfacein the check to enable for any allocation op.Moreover, for function call updates, we enable setting an allocation function callback in
BufferResultsToOutParamsOptsto allow users to emit their own alloc-like op.