How to build for TPU ?

[janpfeifer]

I'm an ML Framework (GoMLX, for Go language) developer using XLA (XlaBuilder for now, but at some point changing to PJRT), and I wanted to run some of my trainers in Cloud TPUs.

How do I build XLA for TPUs ? The configuire.py doesn't seem to have such a --build=TPU option. Or is there a pre-built distribution of XLA libraries for TPUs I could use ?

Thanks!

[cheshire]

Hey,

the TPU backend is non-OSS, so you can't build it from the GH repo.

[janpfeifer]

Oh, I'm sad to hear that ... I used to see TPU related code in the repo.

But if non-OSS, how can I use Cloud TPUs with my project ? Is there a pre-compiled binary/.so that I can link to ?

[cheshire]

Yeah there's a libtpu.so you get which gives you access: https://cloud.google.com/tpu/docs/runtimes

[janpfeifer]

Yeah there's a libtpu.so you get which gives you access: https://cloud.google.com/tpu/docs/runtimes

Thanks. But what is inside this libtpu.so ? Would I have to extract the exported symbols, and reverse-engineer some .h files, and somehow connect the OSS XLA to it ? The page https://cloud.google.com/tpu/docs/runtimes doesn't provide any details.

[cheshire]

I think people don't access it at that level - they would use XLA (or JAX/PT) interfaces, and the implementation details would be hidden behind PjRT/StableHLO APIs.

[janpfeifer]

(I updated the description of the bug to clarify that my question is from the point of view of someone using XLA to create an ML framework in another language that is not Python)

[cheshire]

Maybe the question should be then non-Python non-C++ StableHLO/PjRT/xla_builder bindings?

We've used to have quite a few other frontends: Julia, Elixir, etc, but those don't look very active these days.

[janpfeifer]

Well, I'm using xla_builder, but I'd happily change to use StableHLO/PjRT if it would enable using TPUs.

I modeled the very first version of the XLA bindings for Go on the Elixir version.