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.)
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.
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?
5 comments
[ 2816 ms ] story [ 4500 ms ] thread> 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.)
> 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.
I don't get it, Bazel is incredible.
> 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?