Skip to content
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

[MLIR] Mark debug.callbacks as inactive functions. #706

Merged
merged 17 commits into from
May 24, 2024

Conversation

erick-xanadu
Copy link
Contributor

@erick-xanadu erick-xanadu commented May 2, 2024

Context: debug.callbacks will always be inactive functions when taking the gradient of them.

Description of the Change: Change pyregistry to inactive_callback and mark all inactive_callbacks as inactive.

Benefits: Inactive callbacks are inactive.

Notes:

  • Future PRs will make active_callbacks through specialization and will call inactive_callback as a primitive.
  • Future PR will re-enable callbacks in gradients.

[sc-60496]

`PythonCallOp`s are lowered into LLVM functions. This commit
introduces the change that `PythonCallOp`s are lowered into LLVM
functions which are annotated with the `catalyst.debugCallback`
attribute.

The reason for this is that without this attribute, the only
way to determine if an LLVM call corresponds to a function call to
`@pyregistry` with a third parameter set to zero.

To avoid having to find the definition of the third parameter
(which may have been defined in a different function), we annotate
during the lowering.

This case is likely not going to happen in the current state of
the compiler, since it would depend on outlining. But, it is still
a good change to make since it reduces compilation time having to
traverse the use-def chain.
[skip ci]
@erick-xanadu erick-xanadu force-pushed the eochoa/2024-05-01/specialize-callbacks branch from 2fb0255 to bb78de4 Compare May 9, 2024 15:31
@erick-xanadu erick-xanadu changed the title 🚧 [MLIR] Mark debug.callbacks as inactive functions. [MLIR] Mark debug.callbacks as inactive functions. May 9, 2024
@erick-xanadu erick-xanadu marked this pull request as ready for review May 9, 2024 15:48
Copy link
Collaborator

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

doc/changelog.md Show resolved Hide resolved
mlir/test/Catalyst/AnnotateDebugAsEnzymeConst.mlir Outdated Show resolved Hide resolved
@erick-xanadu erick-xanadu requested a review from dime10 May 23, 2024 15:48
Copy link
Collaborator

@dime10 dime10 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💯

@erick-xanadu erick-xanadu merged commit 2636257 into main May 24, 2024
41 checks passed
@erick-xanadu erick-xanadu deleted the eochoa/2024-05-01/specialize-callbacks branch May 24, 2024 18:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants