Special node acceleration via metadata about nx HOT 6 CLOSED

Axon marks each layer with :op as :metadata, which means that you a compiler could replace operations with a better one depending on what it is. Plus one from me

from nx.

Should we also add a compiler callback for divulging which custom operation nodes are available?
Or should we just set as a best practice that we should document this in the docs for the specific compiler?

from nx.

I struggle to see how it would work in practice. If EXLA has to match on the metdata, then it requires EXLA itself to know about the operation, and in that case we should just expose the operation in Nx+EXLA themselves.

So for this to work as desired, we would need to make EXLA itself extensible. Custom calls are one mechanism to achieve this. Then someone could truly write extension nodes without ever touch Nx/EXLA itself. But it is yet unclear to me the relationship between custom calls and MLIR. Can we have custom MLIR code in our definitions? Or they must all go through custom calls?

from nx.

We could implement the operation node either as MLIR instructions or as a custom call. Both would be feasible, I believe.

Maybe EXLA could provide this functionality itself, then we don't have to worry about changing Nx and how to make it so other compilers know about it.

Nx.Defn.Expr is a public module, so there's no harm in making this part of EXLA only.

I'm almost sure we could also make it so people could provide their own symbols as a dynamically linked library for EXLA to link against and load the custom callbacks!

from nx.

Let's experiment a bit with the format I used with matching on :metadata for now, with implicit coupling between Axon/Bumblebee and EXLA, so we can start exploring this domain

from nx.

Closing this as for now we want to experiment rather than something actionable.

from nx.