-
Notifications
You must be signed in to change notification settings - Fork 57
Issue #1684: Refactor Catalyst pass registering #1984
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Issue #1684: Refactor Catalyst pass registering #1984
Conversation
I'm not sure I fully understand your question. If you're asking why you can add passes created by these factory functions to the pipeline without registering them first, it's because pass registration and adding passes to the pipeline are independent in MLIR. Registering a pass is mainly intended to support the CLI, where passes can be referenced by name in a textual pipeline. So, you can add passes ( But, in the current design, the factory functions you mentioned (such as |
Many thanks for the explanation. That was exactly my doubt, yes.
Good spot. I will add those registrations to the PR then. The new |
…es, and 2) manually implementing a creation function for a new pass (e.g., std::unique_ptr<mlir::Pass> createQuantumBufferizationPass). I have understood that the first of those points was the main goal of issue PennyLaneAI#1684 (Refactor Catalyst pass registering). Add mlir/include/RegisterAllPasses with a catalyst::registerAllPasses convenience function. This function calls the pass registration functions for every group (e.g., registerQuantumPasses). And those pass registration functions 1) are auto generated, and 2) sit in the <module>/Transforms/Passes.h.inc file (e.g., catalyst/mlir/build/include/Quantum/Transforms/Passes.h.inc). Remove Catalyst/Transforms/RegisterAllPasses.cpp, and catalyst::registerAllCatalystPasses. This is the function where we were manually registering every pass. And that we needed to update for every new pass. Passes.h: - Remove manual pass creation and pass registration function declarations. The implementation of those functions is also removed from every file implementing a pass. - Define GEN_PASS_DECL and GEN_PASS_REGISTRATION. Include <module>/Transforms/Passes.h.inc. This adds the auto generated pass creation and pass registration function declarations. Passes.td: - Remove all constructor variables. These variables point to manual pass creation functions. If you don't use these constructor variables, constructors with the same names will be auto generated in the Passes.h.inc files. CMakeLists.txt: - Change mlir_tablegen group name to start with uppercase (e.g., catalyst to Catalyst). That simply translates into, e.g., a name such as registerCatalystPasses being generated instead of registercatalystPasses. - Change mlir_tablegen group name for mlir-hlo to Mhlo. That avoids registermlir-hloPasses being generated as a function name, which is an invalid name. Remove include of <module>/Transforms/Passes.h from cpp files where a Passes.h.inc is later included. Add GEN_PASS_DECL_<pass> to cpp files where GEN_PASS_DEF_<pass> was previously used. This is needed to add the declarations for the pass creation functions. All Quantum passes are implemented now under catalyst::quantum namespace. TODO: check if GEN_PASS_DECL_<pass> are really needed. TODO: CompilerDriver calls mlir::mhlo::registerAllMhloPasses (not catalyst::registerMhloPasses). However, it then goes on to add catalyst::createMhloLegalizeControlFlowPass, catalyst::createMhloLegalizeToStandardPass, and catalyst::createMhloLegalizeSortPass as nested passes, without previously registering them. Is this OK? TODO: I have moved Quantum passes to catalyst::quantum. Is this OK? May it cause problems from the Python connection?
…h DEF and DECL defines are needed, put the declarations before the definitions. TODO: CompilerDriver calls mlir::mhlo::registerAllMhloPasses (not catalyst::registerMhloPasses). However, it then goes on to add catalyst::createMhloLegalizeControlFlowPass, catalyst::createMhloLegalizeToStandardPass, and catalyst::createMhloLegalizeSortPass as nested passes, without previously registering them. Is this OK? TODO: update 'Writing and running your first Catalyst pass' section at doc/dev/transforms.rst.
…dev/transforms.rst. TODO: CompilerDriver calls mlir::mhlo::registerAllMhloPasses (not catalyst::registerMhloPasses). However, it then goes on to add catalyst::createMhloLegalizeControlFlowPass, catalyst::createMhloLegalizeToStandardPass, and catalyst::createMhloLegalizeSortPass as nested passes, without previously registering them. Is this OK?
- Add registerMhloPasses call to registerAllPasses function in catalyst/mlir/include/RegisterAllPasses.h. I've also taken the opportunity to do the following changes: - Move Passes.h, Passes.td, and mhlo_legalize_to_standard_patterns.td from catalyst/mlir/include/mlir-hlo to a Transforms subfolder. - Move mhlo_legalize_control_flow.cpp, mhlo_legalize_sort.cpp, and mhlo_legalize_to_std.cpp from catalyst/mlir/lib/mlir-hlo to a Transforms subfolder. - Catalyst mlir-hlo passes are now under catalyst::mhlo. - Rename mhlo_legalize_to_std.cpp to mhlo_legalize_to_standard.cpp. - Rename populateMhloToStdPatterns function to populateMhloToStandardPatterns in mhlo_legalize_to_standard.cpp. - Rename public tablegen targets MLIRHLOCatalystPassIncGen and MLIRMhloLegalizeToStandardIncGen to CatalystMhloPassIncGen and CatalystMhloLegalizeToStandardIncGen respectively. - Simplify header inclusion in Pipelines.cpp.
fc1c52c to
6cafe77
Compare
…for other groups of passes.
Possibly just historic when Quantum was the only dialect. I don't see an issue having these under the quantum namespace. In any case, nice work 💯 |
Thanks for working on this! 🥳 To answer this question: the function call Thus the three passes you mentioned should currently be registered in But do note that we are going through some migration for the hlo dependency. The registration process for these three (somewhat special) passes will change a bit. We can revisit this after #1921 ! For now, we can just work on the other passes! |
Appreciated!
That's it. We can see this in this PR's code. For example, in
They were there indeed, as @rniczh had also commented. And they still are in this PR's code, now in |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1984 +/- ##
=======================================
Coverage 97.42% 97.42%
=======================================
Files 92 92
Lines 10424 10424
Branches 1002 1002
=======================================
Hits 10156 10156
Misses 211 211
Partials 57 57 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
Hi! The stablehlo migration is now done, please update your local and PR with the following:
You can also find them in If you have any questions, you can look at the PR #1921 or ask me directly! |
Many thanks, I'm on it! |
- Move {include,lib}/hlo-extensions files to a Transforms subfolder.
- Move hlo-extensions passes under catalyst::hlo.
- Rename mlir_tablegen stablehlo group name to Stablehlo.
- Rename stablehlo_legalize_to_std.cpp to stablehlo_legalize_to_standard.cpp.
|
Any idea why docs may be failing? I can build them locally with some warnings: |
autoray released a new version a couple hours ago 🫠 we are handling this right now (Sorry for all the inconvenience, there has been many dependency issues this week...) |
Thanks for the explanation!
Nah, no problem at all. |
|
I deleted all |
Cool. Regarding this, as I commented above, we can add them back, for consistency and because they do no harm. It's maybe easier for future coders to just add both DEF and DECL by default. |
I'm fine with either adding them back or not adding them back, so let's go with whichever is easier for you : ) Just remember to add to the documentation saying |
joeycarter
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Follow the following order: - standard headers (or gmock, gtest) - llvm - mlir - stablehlo - internal headers (from this Catalyst project)
Also take the opportunity to cut some lines short.
Thanks for the comment!
Cool.
It says now, also as suggested by @joeycarter (more or less): But we're still on time to adjust that comment, of course. |
paul0403
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Awesome 🥳 🤩
joeycarter
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks great, thanks @rturrado! 💯
|
Thanks to you, guys! 🍕🎈 |
**Context:** Issue #1684 points to the maintenance problem of having to add a `mlir::registerPass(<module>::create<pass name>Pass);` at `registerAllCatalystPasses` function for every new pass. **Description of the Change:** An early look at the `<module>/Transforms/Passes.h.inc` files, e.g., `Catalyst/Transforms/Passes.h.inc`, shows that: - There is already an auto generated `register<module>Passes` function that registers each pass in a group of passes, e.g., Gradient or Quantum. This function is enabled by the GEN_PASS_REGISTRATION macro. - Function declarations and implementations for creating each pass are also auto generated. They have the name `create<pass name>Pass`. Declarations are enabled by the `GEN_PASS_DECL_<pass name in uppercase>PASS` macro. Implementations are enabled by the `GEN_PASS_DEF_<pass name in uppercase>PASS` macro. And, for a group of passes, all the declarations can be enabled at once with the `GEN_PASS_DECL` macro. This PR tries to use all that auto generated code, and so get rid of the manual pass registration and creation code. **Benefits:** - No need to maintain code when adding a new pass. Note though that if the pass is added to a new group of passes, some scaffolding will be needed. But this can be replicated from other groups of passes. - Less boilerplate code. **Possible Drawbacks:** I have a couple of questions that I would like to be considered during the code review: 1) CompilerDriver calls `mlir::mhlo::registerAllMhloPasses` (not `catalyst::registerMhloPasses`). However, it then goes on to add `catalyst::createMhloLegalizeControlFlowPass`, `catalyst::createMhloLegalizeToStandardPass`, and `catalyst::createMhloLegalizeSortPass` as nested passes, without previously registering them. It is a bit strange to me that you don't need to register nested passes. 2) All Quantum passes are implemented now under `catalyst::quantum` namespace. I did this to ease the refactoring. It also looks OK to my eyes because, for example, Gradient passes are implemented under `catalyst::gradient`. But a) there may be a reason I don't know why Quantum passes were just under `catalyst` namespace, and b) I don't know either if this change can affect the connection with Python code. **Related GitHub Issues:** (Paul: [sc-89760]) --------- Co-authored-by: ringo-but-quantum <> Co-authored-by: Joey Carter <joseph.carter@xanadu.ai>
Context:
Issue #1684 points to the maintenance problem of having to add a
mlir::registerPass(<module>::create<pass name>Pass);atregisterAllCatalystPassesfunction for every new pass.Description of the Change:
An early look at the
<module>/Transforms/Passes.h.incfiles, e.g.,Catalyst/Transforms/Passes.h.inc, shows that:register<module>Passesfunction that registers each pass in a group of passes, e.g., Gradient or Quantum. This function is enabled by the GEN_PASS_REGISTRATION macro.create<pass name>Pass. Declarations are enabled by theGEN_PASS_DECL_<pass name in uppercase>PASSmacro. Implementations are enabled by theGEN_PASS_DEF_<pass name in uppercase>PASSmacro. And, for a group of passes, all the declarations can be enabled at once with theGEN_PASS_DECLmacro.This PR tries to use all that auto generated code, and so get rid of the manual pass registration and creation code.
Benefits:
Possible Drawbacks:
I have a couple of questions that I would like to be considered during the code review:
mlir::mhlo::registerAllMhloPasses(notcatalyst::registerMhloPasses). However, it then goes on to addcatalyst::createMhloLegalizeControlFlowPass,catalyst::createMhloLegalizeToStandardPass, andcatalyst::createMhloLegalizeSortPassas nested passes, without previously registering them. It is a bit strange to me that you don't need to register nested passes.catalyst::quantumnamespace. I did this to ease the refactoring. It also looks OK to my eyes because, for example, Gradient passes are implemented undercatalyst::gradient. But a) there may be a reason I don't know why Quantum passes were just undercatalystnamespace, and b) I don't know either if this change can affect the connection with Python code.Related GitHub Issues:
#1684
(Paul: [sc-89760])