We Accidentally Built an LLVM Compiler for JAX While Shaving a Yak
While building Catalyst, a quantum compiler for PennyLane, we aimed to optimize hybrid workflows using MLIR. In the process, we discovered that our pipeline could compile pure JAX code directly to LLVM, bypassing the XLA runtime entirely. This accidental breakthrough enables standalone binaries, supports dynamically shaped arrays, and integrates native Python control flow without heavy dependencies.
Quantum algorithms are never just quantum; there is a tonne of classical software infrastructure and development needed to support quantum computing.
- ndesaulniers
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.)