-
Notifications
You must be signed in to change notification settings - Fork 199
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
feat: add EXLA.to_mlir_module/2 #1497
Conversation
polvalente
commented
May 23, 2024
- feat: add EXLA.to_mlir_module/3
- doc: add doctest
@@ -71,6 +71,10 @@ defmodule EXLA.MLIR.Module do | |||
* `:use_spmd` - enables Single-Program Multiple-Data partioning. | |||
This is set to true if `:num_partitions` is more than one, otherwise is `false`. | |||
|
|||
* `:module_compilation` - either `:to_mlir` or `:to_pjrt`. The default is `:to_pjrt`. |
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.
@josevalim not sure about the option naming here
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.
It is private, so it is fine!
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.
If we return %EXLA.MLIR.Module{}
in the other function, shouldn't the whole module be public?
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.
Good catch. We should return the string only, not the module struct.
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.
The public-facing function is now returning the binary directly
{:ok, {_xla_time, executable, _extra, _outfeed}} = callback.() | ||
throw({:mlir_module, executable.ref}) |
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.
Gah, this is relying on internals of another module but I can't think of anything better for now, so ship it.