Migrating to One-Shot Bufferization #1027
Conversation
tzunghanjuang
force-pushed
the
bufferization-update
branch
2 times, most recently
from
2c43ec4
to
dce6cc2
Compare
tzunghanjuang
force-pushed
the
bufferization-update
branch
from
dce6cc2
to
6af3eef
Compare
|
This is awesome! Keep up the good work and let me know if you need some help :) |
|
Hello. You may have forgotten to update the changelog!
|
| struct CallbackOpInterface | ||
| : public bufferization::OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel< | ||
| CallbackOpInterface, CallbackOp> { | ||
| static bool supportsUnstructuredControlFlow() { return true; } |
There was a problem hiding this comment.
We never use unstructured control flow in callback op. I've set this to false and it works. Can you expand on why this is needed?
There was a problem hiding this comment.
I was following the FuncOpInterface. I will set to as false.
There was a problem hiding this comment.
Can you show also why exactly you needed to implement this interface? I am curious. Thanks!
There was a problem hiding this comment.
I tried to make the interface similar to FuncOp's. https://github.com/llvm/llvm-project/blob/f59b0c76030aff268b78d475e219708d06b982b5/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp#L325-L351
According to this LLVM commit llvm/llvm-project@6ecebb4,
we might only need it if cf.if appears in func.func. I remember last time I tested it and function like ops did not work without this interface.
I will try using the original interface BufferizableOpInterface::ExternalModel again to see what is going on.
There was a problem hiding this comment.
It seems we do not need this OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel for CallbackOp. I will remove it.
There was a problem hiding this comment.
I have tried removing OpWithUnstructuredControlFlowBufferizableOpInterfaceExternalModel from ForwardOp and ResverseOP but it would fail tests.
Sample error:
error: 'gradient.forward' op type of entry block argument #0('memref<f64, strided<[], offset: ?>>') must match the type of the corresponding argument in function signature('memref<f64>')There was a problem hiding this comment.
Yes, this is the same error you had in a different place about the layout right?
There was a problem hiding this comment.
Yes, this data layout issue should be removed by identity-layout-map flag. Not sure why this error comes back even though the flag is already set.
erick-xanadu
changed the title
This reverts commit 812bf7f.
Co-authored-by: erick-xanadu <[email protected]>
Co-authored-by: erick-xanadu <[email protected]>
There was a problem hiding this comment.
👍 @dime10 , I'll leave you some time to ask questions if you have any. If not let me know and I can merge.
|
Adding do-not merge label to wait after the release. |
Context:
MLIR removed old bufferization (
Conversion-Based bufferization) passes and introduced newOne-Shot-Bufferizationpass. [Link]Note that
Ownership-based Buffer Deallocationwill be handled in a different PR.Only update llvm and mhlo.Jax sticks with 0.4.28.Description of the Change:
Newly Added Bufferization Interface:
Comparing implemented interface methods with BufferizationOpInterface.td.
Operation List:
https://github.com/PennyLaneAI/catalyst/blob/8d0324db672d507d651b46fd65e9058455a9bb03/mlir/lib/Bufferization.md
Major Updates:
one-shot-bufferizeinterface forcatalyst-bufferize,gradient-bufferize, andquantum-bufferize.declarePromisedInterfaces.FunctionOPInterfacelikeForwardOPandReverseOpneed to accessOneShotModuleBufferize.Gradient.ReturnOPisReturnLikenow to support bufferization.function-boundary-type-conversiontoidentity-layout-mapto avoid invalid type conversion betweenMemref<T>andMemref<T, Stride<...>>allow-return-allocs-from-loopsflag to avoidscp.forandscp.whileto avoid arg and yield types mismatchcopy-before-writeflagRequired LLVM and MHLO Patches:
OneShotModuleBufferizesupportFunctionOpInterfaceinstead ofFuncOP. [MLIR] MakeOneShotModuleBufferizeuseOpInterfacellvm/llvm-project#110322FunctionOpInterfacepatch to mhlo.FunctionArgTypeConverterFnunderBufferizationOptionstakesFunctionOpInterfaceinstead offunc::FuncOpopenxla/xla#17935Related GitHub Issues:
[sc-71487]
[Discourse discussion about migration]
[Discord discussion about FunctionOpInterface and FuncOp]
[Discord discussion about memory layout]
[LLVM Upstream]
[LLVM PR1]
[LLVM PR2]
llvm/llvm-project@7f04a8a
youtube link with discussion about one-shot bufferizer
pdf with info