-
Notifications
You must be signed in to change notification settings - Fork 2.8k
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
CPU profiling (not tracing) #24349
Comments
One difficulty with this feature request is that after a function is compiled into HLO, information about the original python function boundaries is lost. So we would not be able to automatically generate a profile that contains information about user-defined functions. You can look at the compiled code yourself by running One workaround for this could be to decorate all of your user functions using |
There are several ways of exporting HLO/StableHLO from JAX, and many (certainly the StableHLO MLIR bytecode portable artifacts) do export location information (which maps HLO/StableHLO ops to the Python code that created them). |
For example, this can be used (although it might be not the most compact representation):
|
Related discussion: #23251 |
I want to get some flamegraphs from some JIT'd JAX CPU code to understand where time is being spent (in terms of my user-defined functions).
My understanding (based on the docs) is that currently the recommended approach is to add some custom tracing events and run JAX's tracing feature.
This seems rather suboptimal. Is there a better way?
Related discussion: #19888
The text was updated successfully, but these errors were encountered: