Add NIF for loading custom plugins

[seanmor5]

WIP

[josevalim]

Should we use a process instead such that, if someone does Application.stop(:exla) the process is shutdown as well as all plugins? 🤔

[polvalente]

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.

[polvalente]

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.

[josevalim]

prep_stop also works but you would need to iterate all persistent term to find the keys relevant to us. ETS would be better.

[polvalente]

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

[josevalim]

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?

[seanmor5]

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

[josevalim]

Sounds good! Happy to discuss sketches of the API any time!

[polvalente]

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

[josevalim]

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.

[josevalim]

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.