We accidentally built an LLVM compiler for Jax
1–10 of 26 posts
Re: We accidentally built an LLVM compiler for Jax
#2Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X
> So if XLA already uses LLVM, why is our approach different?
Uses MLIR, XLA does not.
> So… what is the point?
> Honestly? We aren’t entirely sure yet.
> Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations for linear algebra on GPUs and TPUs. If you are training a massive transformer, stick to standard JAX.
> But what we do think is cool is what happens when you connect JAX directly to the broader LLVM ecosystem and drop the heavy XLA runtime. (Plus, no need to build XLA using Bazel either! You’re welcome.)
Re: We accidentally built an LLVM compiler for Jax
#3Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X > So if XLA already uses LLVM, why is our approach different? Uses MLIR, XLA does not. > So… what is the point? > Honestly? We aren’t entirely sure yet. > Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations…
> Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X
I would love to hear those stories! Sadly I'm based in Toronto, so dropping by the Bay Area meetups isn't in the cards anytime soon.
If any of that history ever makes it into a blog post, I'd be first in line to read it.
Re: We accidentally built an LLVM compiler for Jax
#4Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X > So if XLA already uses LLVM, why is our approach different? Uses MLIR, XLA does not. > So… what is the point? > Honestly? We aren’t entirely sure yet. > Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations…
> Uses MLIR, XLA does not.
https://github.com/search?q=repo%3Aopenxla%2Fxla+mlir&type=c...
> 1.5k files
you're behind the times. XLA moved over probably ~2 years ago - "Captain Awesome" eventually relented.
Re: We accidentally built an LLVM compiler for Jax
#5Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X > So if XLA already uses LLVM, why is our approach different? Uses MLIR, XLA does not. > So… what is the point? > Honestly? We aren’t entirely sure yet. > Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations…
I can't tell what you mean by this comment. Is XLA good? is XLA bad, cuz it doesn't use MLIR? Speak loud, for we are hard of hearing.
Re: We accidentally built an LLVM compiler for Jax
#6Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X > So if XLA already uses LLVM, why is our approach different? Uses MLIR, XLA does not. > So… what is the point? > Honestly? We aren’t entirely sure yet. > Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations…
I can't tell what you mean by this comment. Is XLA good? is XLA bad, cuz it doesn't use MLIR? Speak loud, for we are hard of hearing.
Re: We accidentally built an LLVM compiler for Jax
#7Neat. XLA predates MLIR. Interesting stories there. You'd have to stop by the Bay Area LLVM monthly meetup to hear them. :-X > So if XLA already uses LLVM, why is our approach different? Uses MLIR, XLA does not. > So… what is the point? > Honestly? We aren’t entirely sure yet. > Let me be perfectly clear: this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations…
> Uses MLIR, XLA does not. https://github.com/search?q=repo%3Aopenxla%2Fxla+mlir&type=c... > 1.5k files you're behind the times. XLA moved over probably ~2 years ago - "Captain Awesome" eventually relented.
lol, good luck to amazon
Re: We accidentally built an LLVM compiler for Jax
#8Earlier quoted context omitted.
I can't tell what you mean by this comment. Is XLA good? is XLA bad, cuz it doesn't use MLIR? Speak loud, for we are hard of hearing.
https://en.wikipedia.org/wiki/Shibboleth
Indeed, I am an outgroup member. look at me the dumbfuck. Now how does one become an ingroup member, by learning.
I am an active contributer to MLIR, if you knew my name and you weren't some weirdo, you could know me. This is a pseudonym account of which I have several, since I grew up in an age where you hid your identity. Now how XLA interacts with that I have no clue dude. Maybe help with that and don't be a low effort dick.
Re: We accidentally built an LLVM compiler for Jax
#9Earlier quoted context omitted.
https://en.wikipedia.org/wiki/Shibboleth
Indeed, I am an outgroup member. look at me the dumbfuck. Now how does one become an ingroup member, by learning. I am an active contributer to MLIR, if you knew my name and you weren't some weirdo, you could know me. This is a pseudonym account of which I have several, since I grew up in an age where you hid your identity. Now how XLA interacts with that I have no clue dude. Maybe help with that and don't be a low e…
brother i'm trying to tell you that the guy you're asking to be clearer is being vague because he's employing a shibboleth - a "wink wink i'm an og and you can only access me via the secret handshake at our club meeting". not you. i pasted the link because you explicitly say "I can't tell what you mean by this comment" so i'm proposing what they mean.
relax you'll live longer.
Re: We accidentally built an LLVM compiler for Jax
#10This is actually really neat, I had an agent create some experimental probes and check out whats going on here and it seems like a pretty cool project. I'm gonna need to digest how this might be useful but I think this could be cool for some sort of FPGA uses perhaps.
> this is not going to beat XLA for standard deep learning workloads. XLA has years of hyper-specific optimizations for linear algebra on GPUs and TPUs.
I'm not so convinced about this.
It's actually really easy to write two mathematically equivalent formulations in Python of something basic that have over a 2x performance difference in them after jax jit.
XLA is not that smart. And I'm not talking some niche nonsense I mean simple matrix multiplication graphs and residual connections on the CUDA backend.