Minor update to the new MLIR specification for custom ops #8
Conversation
|
@EiffL — Thank you!! I'm up to my eyeballs with work for the next little bit, but I'll take a look through and sanity check the GPU stuff ASAP. |
|
Also: @sharadmv and others might have thoughts on where this fits w.r.t. jax-ml/jax#12632 and its potential future implementation, but I still think it's probably sensible to update the doc here in the meantime! |
* Update Dockerfile * Update Dockerfile * Update Dockerfile * Update Dockerfile * Update Dockerfile
|
No problem and no rush @dfm ^^ I've checked it works on GPU, and updated the docker cuda container to work with current jax version, should be all good :-) |
|
@EiffL — I'm finally taking a look at this. Sorry for being so slow. I made some tiny superficial changes, but I'm otherwise totally happy. One thought is that it might be useful to have more sensible words about the relationship between MLIR and XLA, and "translation" vs. "lowering", since both this document and the JAX docs are a little muddled about this, I'd say. But, I don't think that should stop us merging this. As soon as you sign off, @EiffL, I'm happy to merge. Thanks again!! |
Following #7 , this small PR proposes to update this post to the new way that JAX uses to define custom ops by directly writing their specification using MLIR.
I'm pretty noob in MLIR but it seems very cool, and I guess at some point XLA will deprecate the old way of custom call building.
Here is the doc for the custom call op in MLIR: https://www.tensorflow.org/mlir/hlo_ops#mhlocustom_call_mlirmhlocustomcallop
And all the custom ops in jaxlib are now built this way.
Let me know if this looks good. I haven't run the checks on GPU