-
Notifications
You must be signed in to change notification settings - Fork 212
Add NIF for loading custom plugins #1519
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
Conversation
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.
Should we use a process instead such that, if someone does Application.stop(:exla) the process is shutdown as well as all plugins? 🤔
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.
Regarding how to store things, I would rather do an ETS than a process for storing the custom call registry if persistent term is not what we want, to keep close to the same read concurrency.
Another possible alternative would be a GenServer that manages the persistent term on terminate, it would clean up the persitent term state.
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.
In this manager alternative, we'd have non-process functions that read first and write if needed to the persistent term, and the only purpose for the GenServer would be to ensure cleanup upon termination.
PS: Upon writing this I went reading and found https://hexdocs.pm/elixir/1.12/Application.html#c:prep_stop/1 which would serve this purpose nicely.
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.
prep_stop also works but you would need to iterate all persistent term to find the keys relevant to us. ETS would be better.
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 there's something like EXLA.CustomCalls.cleanup/0 we could call, then that function will already know about all of the keys that should be cleaned up.
I see no issue with using ETS however, as the speed difference here only matters at defn compile time and not runtime
The general idea looks neat to me! Although we still need to figure out how they would be specified via defn.
We probably want to use optional callbacks for this. Maybe we allow anyone to define optional callbacks and then we allow optional callbacks to be registered dynamically, when we register the plugin? Then we have the "built-in" optional callbacks and the third party ones. Thoughts?
Yeah I think it should be similar to optional. I talked to Paulo a bit yesterday and we are going to move the custom call registration to jit time so we can provide shape and type information to the implementer. That would require the user to register a library and then pass the name of the library plus the symbol to the custom call
Then we could have a fallback which is an Nx function
Sounds good! Happy to discuss sketches of the API any time!
The general idea looks neat to me! Although we still need to figure out how they would be specified via
defn.We probably want to use optional callbacks for this. Maybe we allow anyone to define optional callbacks and then we allow optional callbacks to be registered dynamically, when we register the plugin? Then we have the "built-in" optional callbacks and the third party ones. Thoughts?
I was thinking more along the ways of adding something like Nx.Defn.Expr.block(name, expr) which would mark a given subgraph with a specific name, and then EXLA would take advantage of the arity and name to delegate blocks to custom calls whenever relevant.
The mapping could even be block name -> custom_call name in an EXLA config key
I was thinking more along the ways of adding something like Nx.Defn.Expr.block(name, expr) which would mark a given subgraph with a specific name, and then EXLA would take advantage of the arity and name to delegate blocks to custom calls whenever relevant.
When I was thinking about this yesterday I came to the conclusion that we need the function name, the input and output types, and a default implementation for other backends, and that's exactly what optional provides. Which is why I thought about using instead of coming up with a new construct.
Yeah, and the main issue with ets is copying the data that you read. If the data is a reference, it should be plenty fast. Even a process should be good enough.
WIP