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

feat: add EXLA.to_mlir_module/2 #1497

Merged
merged 7 commits into from
Jul 11, 2024
Merged

feat: add EXLA.to_mlir_module/2 #1497

merged 7 commits into from
Jul 11, 2024

Conversation

polvalente
Copy link
Contributor

  • feat: add EXLA.to_mlir_module/3
  • doc: add doctest

@polvalente polvalente self-assigned this May 23, 2024
@@ -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`.
Copy link
Contributor Author

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

Copy link
Collaborator

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!

Copy link
Member

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?

Copy link
Collaborator

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.

Copy link
Contributor Author

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

Comment on lines +381 to +382
{:ok, {_xla_time, executable, _extra, _outfeed}} = callback.()
throw({:mlir_module, executable.ref})
Copy link
Collaborator

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.

exla/lib/exla.ex Outdated Show resolved Hide resolved
@polvalente polvalente merged commit 92a0861 into main Jul 11, 2024
6 of 8 checks passed
@polvalente polvalente deleted the pv-feat/add-to-mlir-module branch July 11, 2024 05:04
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.

3 participants