Clarification: Does the standalone XLA CPU backend really lower HLO ops (e.g. dot/GEMM) into Eigen’s C++ routines?

[a80908]

Hi everyone,

I came across the following statement in some notes about JAX’s XLA CPU backend:

“The standalone XLA CPU backend will lower most HLO ops (including dot/generalized GEMM) into calls to Eigen’s C++ routines.”

However, I haven’t been able to find any official documentation or source in the JAX repository (or in the XLA docs) that confirms this behavior.
• Is this statement accurate?
• If so, where can I find the relevant docs or source code references?
• If not, what actually happens under the hood when HLO ops like dot or generalized GEMM are lowered on CPU?

Thanks in advance for any pointers!

[dfm]

It's hard to provide too much useful feedback without more details about why you want to know this information :D, but it looks like that statement is currently accurate for dot at least. Here's the relevant XLA source code AFAICT: https://github.com/openxla/xla/blob/9c8f9ea2bd606918b25d881cfcf78a68ca1bc642/xla/backends/cpu/runtime/dot_thunk.h

But, even as a full time JAX developer, I wouldn't say that I know all the details of exactly how every HLO op gets implemented in the CPU backend. If you have specific cases where it would be useful to understand the full lowering path for a particular operation, you might get more useful feedback on the XLA repo!

[a80908]

Thanks so much for digging into the XLA code and pointing me at dot_thunk.h! That definitely confirms that at least the dot HLO really is lowered into an Eigen contraction under the hood.

A couple of follow-ups if you don’t mind:
1. Generalized GEMM (“dot_general”)
Do you know if it follows the exact same pathway—i.e. its own thunk calling into Eigen’s C++ routines—or does it go through a slightly different helper? I had a quick look around xla/backends/cpu/runtime/ but I wasn’t 100% sure which files to pick apart.
2. Other HLOs and documentation
Is there a central mapping or “registry” in the XLA CPU backend where it decides “HLO op X → Eigen call Y” that I could study? Even a high-level design doc or comment would be super-helpful to understand how things like convolutions, element-wise ops, reductions, etc., get dispatched.

Thanks again for your pointers—appreciate any pointers at all!