What is the recommended way to use XLA directly from C++?

TensorFlow General Discussion
,
1

I’m playing around with XLA, and the website (tensorflow.org/performance/xla) is not clear on how to use it, from the perspective of someone trying to use XLA directly (not through TensorFlow) in C++ (or some other language).

The site mentions HLO (I couldn’t find more documentation about it), and the C++ XlaBuilder API – which I’ve been successfully using.

How do these two (HLO and XlaBuilder) relate ?

Is there a recommendation for someone wanting to use XLA directly ? And how do Tensorflow/PyTorch/Jax/Nx/Julia use XLA ?

thanks!!

2

Non TF or TF bridges XLA related topics are going to be discussed in the new project/discussion space at:

(/cc @Mehdi_AMINI)

4

First, you need to turn your ML model into an HloModule. Then you can invoke tensorflow/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc at f8145fac4ffda49b6b58f8c93b4924f30c6c03cb · tensorflow/tensorflow · GitHub to compile the model. You must import the XLA modules into your C++ code as you might any other C++ library.

6

Thanks for the reply Thomas. It didn’t exactly answer my question: for instance, just randomly trying things, I got XlaBuilder API to work – I’m compiling and using computations even in GPU.

My question is: why should I use XlaBuilder API or figure out the HLO language and use that API instead ? What have other XLA clients used ? What are the trade-offs, any plans of future support ? I need guidance in choosing one or the other.

Btw, thanks for the pointer to the code, I didn’t know where to start with HLO either. Is there more detailed documentation ? The corresponding .h file is not very helpful either … the XlaBuilder is slightly better documented (but not by much, still lots of “code archeology” that one needs to do to use the library).

cheers