We accidentally built an LLVM compiler for Jax

Posted by infinitewalk 3 days ago

Comments

Comment by kingstnap 1 day ago

This 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.

Comment by kingstnap 1 day ago

Just for receipts I span up a basic test for XLA:

Can you figure out that 16 matrix vector multiplications into a concatenate is the same as concatenating first into a matrix matrix operation which can go on the GEMM. This is like the most basic thing you can imagine doing.

Turns out no it doesn't and there is a 23% performance difference by moving the concatenate up in the python code in my specific test.

Not to say that it didn't recognize it. I had an agent look at the XLA and the graph actually does a partial fusion into a sum and stack. But does not realize the whole thing is just a GEMM.

In more complicated examples the differences you can get can be much larger.

Comment by Gangway0829 1 day ago

Is this surprising? My experience with compilers from the old days of Fortran is that they care about correctness first, performance second. I used to spend plenty of time rewriting algorithms so they ended up in a form the compiler would like.

I think this is a great area for LLMs. You write the correct physics code, the LLM analyzes intent and goes back and forth with the compiler and your correct code to rewrite it in something that emits performant code

Comment by 1 day ago

Comment by Archit3ch 22 hours ago

1. Why not Julia? IIRC, some/all of this already exists in that ecosystem.

> Something to note is that quantum algorithms are never purely quantum. There is a tonne of classical processing needed to get inputs ready, post-process outputs, and even construct and represent the quantum circuits themselves.

2. Did you face any problems for the 'surrounding' areas by JAX imposing the limitation of pure functions?

Comment by infinitewalk 18 hours ago

Main reason is that the quantum computing/scientific research user base is still very Python entrenched

Comment by Archit3ch 18 hours ago

Sure, but if you're writing in the JAX DSL you've already abandoned python.

Comment by Gangway0829 14 hours ago

I mean, not really? Most Jax code is only marginally more restrictive than numpy, and most of the annoyances I have with Jax are shared with numpy.

And having spent too much of my life writing some extremely high performance physics code, 99% of the code complexity comes from the stupid stuff, like handling configuration and IO and restarting. Python handles that stuff great. Oh, and tests. Ever tried to write unit tests in Fortran? You'll hug python tight and promise it you'll never leave again.

Comment by steeve 1 day ago

> Plus, no need to build XLA using Bazel either! You’re welcome.

I don't get it, Bazel is incredible.

Comment by melodyogonna 1 day ago

Many people find it complicated.

Comment by fragmede 23 hours ago

It is. The question is if it's worth the complication. The cheeky answer would be to say it is, but that's not universally true.

Comment by ndesaulniers 1 day ago

Neat. 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.)

Comment by infinitewalk 1 day ago

> 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.

Comment by gregdaniels421 1 day ago

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.

Comment by mathisfun123 1 day ago

https://en.wikipedia.org/wiki/Shibboleth

Comment by gregdaniels421 1 day ago

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.

Comment by mathisfun123 1 day ago

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.

Comment by steeve 1 day ago

XLA uses MLIR very much, the interface to it is StableHLO, which is a MLIR dialect.

Comment by ndesaulniers 22 hours ago

Ah, I'm behind the times.

Comment by mathisfun123 1 day ago

> 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.

Comment by throwaway647272 1 day ago

This happened “despite”. MLIR is definitely a productivity boost and yet, the progress in adoption has been stalled many times by one aforementioned charismatic leader. It’s truly mind boggling that a personal spat would influence roadmap of an entire org for years.