Will mojo support JIT compilation?
From what I understand mojo does have already some JIT capabalities (like notebooks) but i'm talking about things like being able to pass run time variables to parameters (not arguments) in functions.
18 Replies
+1
I’m pretty sure it already does, just through the REPL
Could you share an example?
![No description](https://cdn.answeroverflow.com/1240739455778689074/IMG_5748.jpg)
JIT compilation could be really useful to handle computations on floats. During compilation, Mojo is using infinite precision, so specifying some runtime code to be jitted could be awesome 🙂
I imagine it would significantly increase the binary size though because there has to be said compiler to JIT it so it would come with a trade off
yeah but i'm talking about things like being able to pass run time variables to parameters (not arguments) in functions. Like right now there isn't a way to do this type of things, from what i know
Got it
+1
Mojo supports just-in-time compiled in the sense that the only the compiled code runs even in a REPL. PyPy etc. style "run time specialisation" is fundamentally incompatible with Mojo's compilation strategy, I think.
you think. I'm just asking about this because i feel like if mojo wants to be an AI language I think things like being able to pass run time values to parameters should be possible, because the reason jax, pytorch and others can do many things while the code is easy and simple to write, is thanks to this no?, like being able to use the python as the runtime and from there doing compiling to create efficient code. But maybe i'm wrong because i don't understand how the max engine works so because i don't know they do things in the max engine maybe that's why I think it is necessary.
@NKspartan Ah, I see. If what you have in mind is JAX, then it's a bit different. It's a domain specific compiler, MAX is like that. It's kinda a separate compiler. JAX uses XLA as the backend compiler, which is written in C++. But it's clear that the JIT related feature is not built into C++ the language, right? Similarly, the DS compiler can be built with Mojo, but I think it's the best to not see it as part of "the mojo compiler".
okay so max is also a separate compiler you are saying (that's what i was somewhat thinking), in the same way jax uses xla, okay. So now my question is how could i write that in mojo hahah. But thank you for you help sora as always.
IIUC, MAX has an additional advantage: it can compile/fuse your ML model code on the lower level as well.
Consider this code:
the front end might lower it to something like this
A ML compiler might simplify it further into something like
Now, the runtime provided by XLA will lower this further and call the corresponding prewritten kernels (in C++)
One can imagine the C++ implementation of
sin
and cos
sharing some intermediate results. However, since the DSL compiler cannot see through their implementations, it cannot reason across the 'kernel' boundaries, nor can it further fuse them. MAX, on the other hand, does not have this limitation. It benefits from techniques such as Common Subexpression Elimination (CSE), even at the lowest level, because the implementation of the kernels themselves goes through MLIROkay so even max can have more optimization thanks to also lowering the kernels to mlir. One question the first mlir code did you write it manually or did you get from running the function in mojo, because I wanted to know how to see the compiled mlir code from mojo
As far as i know, it's not a feature we can use yet. Sad.