xla Documentation

Technical Documentation: openxla/xla

ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference

Overview (README.md) Copy Section

XLA

XLA (Accelerated Linear Algebra) is an open-source machine learning (ML)
compiler for GPUs, CPUs, and ML accelerators.

The XLA compiler takes models from popular ML frameworks such as PyTorch,
TensorFlow, and JAX, and optimizes them for high-performance execution across
different hardware platforms including GPUs, CPUs, and ML accelerators.

openxla.org is the project's website.

Get started Copy Chapter

If you want to use XLA to compile your ML project, refer to the corresponding
documentation for your ML framework:

• PyTorch
• TensorFlow
• JAX

If you're not contributing code to the XLA compiler, you don't need to clone and
build this repo. Everything here is intended for XLA contributors who want to
develop the compiler and XLA integrators who want to debug or add support for ML
frontends and hardware backends.

Contribute Copy Chapter

If you'd like to contribute to XLA, review
How to Contribute and then see the
developer guide.

Contacts Copy Chapter

• For questions, contact the maintainers - maintainers at openxla.org

Resources Copy Chapter

• Community Resources

Code of Conduct Copy Chapter

While under TensorFlow governance, all community spaces for SIG OpenXLA are
subject to the
TensorFlow Code of Conduct.

Errors Overview (`docs/errors_overview.md`) Copy Section

XLA Errors Overview

XLA errors are categorized into different XLA error sources. Each source has a
list of an additional context other than the error message, which will be
attached to each error within the category.

🚧 Note that this standarization effort is a work in progress so not all error
messages will have an attached error code yet.

An example error log might look like:

XlaRuntimeError: RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G. Total hbm usage >= 49.34G: reserved 3.12M program unknown size arguments 49.34G

JaxRuntimeError: RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %ragged_latency_optimized_all_gather_lhs_contracting_gated_matmul_kernel.18 = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(%get-tuple-element.18273, %get-tuple-element.18274, %get-tuple-element.18275, %get-tuple-element.18276, %get-tuple-element.18277, /*index=5*/%bitcast.8695, %get-tuple-element.19201, %get-tuple-element.19202, %get-tuple-element.19203, %get-tuple-element.19204), custom_call_target=""

Statuses and CHECK failures Copy Chapter

In general, in XLA we can flag corrupted execution with two mechanisms: statuses
and CHECK macro failures.

Statuses are meant for non-fatal, recoverable errors. The assumption is that the
function returns, and execution continues down the path where the caller
explicitly checks the returned Status object. It's useful for handling invalid
user input or expected resource constraints.

On the other hand, CHECK failures cover programmer's errors or violations of
invariants that should never happen if the code is correct. In case of an
activated CHECK the program will log the error message and immediately
terminate. It could ensure internal consistency, such as checking that a pointer
is non-null before dereferencing it.

Error codes Copy Chapter

Here is an index list with all error codes.

Index (`docs/index.md`) Copy Section

XLA

XLA (Accelerated Linear Algebra) is an open-source compiler for machine
learning. The XLA compiler takes models from popular frameworks such as PyTorch,
TensorFlow, and JAX, and optimizes the models for high-performance execution
across different hardware platforms including GPUs, CPUs, and ML accelerators.

As a part of the OpenXLA project, XLA is built collaboratively by
industry-leading ML hardware and software companies, including
Alibaba, Amazon Web Services, AMD, Apple, Arm, Google, Intel, Meta, and NVIDIA.

Key benefits Copy Chapter

• Build anywhere: XLA is already integrated into leading ML frameworks
  such as TensorFlow, PyTorch, and JAX.

• Run anywhere: It supports various backends including GPUs, CPUs, and ML
  accelerators, and includes a pluggable infrastructure to add support for
  more.

• Maximize and scale performance: It optimizes a model's performance with
  production-tested optimization passes and automated partitioning for model
  parallelism.

• Eliminate complexity: It leverages the power of
  MLIR to bring the best capabilities into a single
  compiler toolchain, so you don't have to manage a range of domain-specific
  compilers.

• Future ready: As an open source project, built through a collaboration
  of leading ML hardware and software vendors, XLA is designed to operate at
  the cutting-edge of the ML industry.

Documentation Copy Chapter

To learn more about XLA, check out the links on the left. If you're a new XLA
developer, you might want to start with XLA architecture and
then read Contributing.

Megascale Overview (`docs/megascale/overview.md`) Copy Section

Megascale XLA

MegascaleXLA is a compiler + runtime system that powers large-scale TPU
training. It implements collective communication primitives that allow multiple
TPU slices to communicate, which allows running training jobs that span beyond
the limits of a single ICI domain.

The debugging guide discusses how to identify and
diagnose sources of performance issues such as slowness, hangs or errors in a
multi-slice job driven by Megascale.

The Megascale Performance Tuning Flags guide details
how to configure and tune Megascale XLA flags to optimize collective
communication performance and alleviate bottlenecks in multi-slice workloads.

Terminology Copy Chapter

• Slice

  • A slice is a collection of chips all located inside the same TPU Pod
    connected by high-speed inter chip interconnects (ICI). Slices are
    described in terms of chips or TensorCores, depending on the TPU
    version.
  • Multislice is a group of slices, extending TPU connectivity beyond the
    inter-chip interconnect (ICI) connections and leveraging the data-center
    network (DCN) for transmitting data beyond a slice. Data within each slice
    is still transmitted by ICI. Using this hybrid connectivity, Multislice
    enables parallelism across slices and lets you use a greater number of TPU
    cores for a single job than what a single slice can accommodate.
  • TPUs can be used to run a job either on a single slice or multiple
    slices.

• RapidEye

  • RapidEye is a system that aims to provide global ML debugging
    infrastructure to quickly identify and root cause issues caused by bad
    hardware or software bugs. It monitors Megascale jobs to automatically
    detect, analyze, and classify hang events. The process involves
    collecting data from all job workers, coordinating responses when a hang
    occurs, and generating a summary digest file for each event.
  • RapidEye is enabled by default for all multislice workloads. The
    diagnosis can be found under the Pathways resource manager or MXLA
    coordinator (slice0 task0 for multi-controller JAX workload). RapidEye
    is also used for auto-removing bad TPUs and NICs based on fleetwide
    rapideye data.

• Megascale Collective

  • XLA collectives are supported via the Megascale XLA (MXLA) primitives
    which are not directly usable by the end user. At the time of writing
    this the MXLA primitives include collectives include AllGather,
    AllReduce, AllToAll, ReduceScatter and OneToOne. The reduction
    operations that are currently supported include summation, max/min and
    product.

Pjrt Cpp Api Overview (`docs/pjrt/cpp_api_overview.md`) Copy Section

PJRT C++ Device API Overview

Background Copy Chapter

PJRT
is the uniform Device API that we want to add to the ML ecosystem. The long term
vision is that:

1. Frameworks (JAX, TF, etc.) will call PJRT, which has device-specific
   implementations that are opaque to the frameworks;
2. Each device focuses on implementing PJRT APIs, and can be opaque to the
   frameworks.

PJRT offers both a C API and C++ API. Plugging in at either layer is OK, the C++
API uses classes to abstract away some concepts, but also has stronger ties to
XLA datatypes. This page focuses on the C++ API.

PJRT Components Copy Chapter

Note: Most items in this diagram also have backpointers, memory spaces know
their device(s) and client, devices know their client, buffers know their memory
space.

PjRtClient Copy Section

Full reference at pjrt_client.h > PjRtClient.

Clients manage all communication between the device and framework, and
encapsulate all state used in the communication. They have a generic set of APIs
for interacting with a PJRT plugin, and they own the devices and memory spaces
for a given plugin.

PjRtDevice Copy Section

Full references at pjrt_client.h > PjRtDevice,
and pjrt_device_description.h

A device class is used to describe a single device. A device has a device
description to help identify its kind (unique hash to identify GPU/CPU/xPU), and
location within a grid of devices both locally and globally.

Devices also know their associated memory spaces and the client it is owned by.

A device does not necessarily know the buffers of actual data associated with
it, but it can figure that out by looking through its associated memory spaces.

PjRtMemorySpace Copy Section

Full reference at pjrt_client.h > PjRtMemorySpace.

Memory spaces can be used to describe a location of memory. These can either be
unpinned, and are free to live anywhere but be accessible from a device, or they
can be pinned and must live on a specific device.

Memory spaces know their associated buffers of data, and the devices (plural)
that a memory space is associated with, as well as the client it is a part of.

PjRtBuffer Copy Section

Full reference at pjrt_client.h > PjRtBuffer.

A buffer holds data on a device in some format that will be easy to work with
inside the plugin, such as an MLIR elements attr or a proprietary tensor format.
A framework may try to send data to a device in the form of an xla::Literal,
i.e. for an input argument to the module, which must be cloned (or borrowed), to
the device's memory. Once a buffer is no longer needed the Delete method is
invoked by the framework to clean up.

A buffer knows the memory space it is a part of, and transitively can figure out
which devices are able to access it, but buffers don't necessarily know their
devices.

For communicating with frameworks, buffers know how to convert to and from an
xla::Literal type:

// Literal to Buffer
absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(...) {...}

// Buffer to Literal
xla::Future<> ToLiteral(xla::MutableLiteralBase* literal) override {...}

APIs for creating a buffer have Buffer Semantics
which help dictate if literal data from the host buffer can be shared or copied
or mutated.

Lastly, a buffer may need last longer than the scope of its execution, if it is
assigned to a variable in the framework layer x = jit(foo)(10), in these cases
buffers allow building external references which provide a temporarily owned
pointer to the data held by the buffer, along with metadata (dtype / dim sizes)
for interpreting the underlying data.

PjRtCompiler Copy Section

Full reference at pjrt_compiler.h > PjRtCompiler.

The PjRtCompiler class provides useful implementation details for XLA
backends, but is not necessary for a plugin to implement. In theory, the
responsibility of a PjRtCompiler, or the PjRtClient::Compile method, is to
take an input module and return a PjRtLoadedExecutable.

PjRtExecutable / PjRtLoadedExecutable Copy Section

Full reference at pjrt_executable.h > PjRtExecutable,
and pjrt_client.h > PjRtLoadedExecutable.

A PjRtExecutable knows how to take a compiled artifact and execution options
and serialize/deserialize them so an executable can be stored and loaded as
needed.

The PjRtLoadedExecutable is the in-memory compiled executable which is ready
for input arguments to execute, it is a subclass of PjRtExecutable.

Executables are interfaced with via one of the client's Execute methods:

// Execute on addressable devices
absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
Execute(absl::Span<const std::vector<PjRtBuffer*>> argument_handles, ...) {...}

// Execute assigned replica/partition on the specified device
absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
ExecuteSharded(absl::Span<PjRtBuffer* const> argument_handles,
              PjRtDevice* device, ...) {...}

// Execute on specified device, single replica / partition
absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
ExecutePortable(absl::Span<PjRtBuffer* const> argument_handles,
                PjRtDevice* device, ...) {...}

Before calling Execute the framework will transfer all required data to
PjRtBuffers owned by the executing client, but returned for the framework to
reference. These buffers are then provided as arguments to the Execute method.

PJRT Concepts Copy Chapter

Futures & Async Computations Copy Section

If any part of a plugin is implemented asynchronously, it must properly
implement futures.

Consider the following program:

@jax.jit
def foo(x): return x + 1

x = foo(1)
# [...] other logic not using `x`
print(x + 1)

An async plugin would be able to enqueue the computation x, and immediately
return a buffer which isn't ready to be read yet, but execution will populate
it. Execution can continue to enqueue necessary computations after x, that
don't require x, including execution on other PJRT devices. Once the value of
x is needed, execution will block until the buffer declares itself ready via
the future returned by GetReadyFuture.

Futures can be useful to determine when an object becomes available, including
devices and buffers.

Advanced concepts Copy Section

Extending beyond implementing the base APIs will expand the features of JAX that
can be used by a plugin. These are all opt-in features in the sense that at
typical JIT and execute workflow will work without them, but for a production
quality pipeline some thought should likely be put into the degree of support
for any of these features supported by PJRT APIs:

• Memory spaces
• Custom layouts
• Communication ops like send/recv
• Host offloading
• Sharding

Typical PJRT framework-device communication Copy Chapter

Example Log Copy Section

The following is a log of the methods called to load the PJRT plugin and
execute y = jax.jit(lambda x: jnp.power(x, jnp.int32(2)))(1). In this case
we log JAX interacting with the StableHLO Reference PJRT plugin.

Example log

/* Detailed source-code truncated for AI context efficiency. */

Pjrt Index (`docs/pjrt/index.md`) Copy Section

Project: /_project.yaml
Book: /_book.yaml

PJRT - Uniform Device API

PJRT C API is the uniform Device API that we want to add to the ML ecosystem.
The long term vision is that: (1) frameworks (TF, JAX, etc.) will call PJRT,
which has device-specific implementations that are opaque to the frameworks; (2)
each device focuses on implementing PJRT APIs as PJRT plugins, which can be
opaque to the frameworks.

Communication channels Copy Chapter

• Issues and feature requests can be filed in the OpenXLA/xla repo.
• Questions regarding PJRT can be asked on the OpenXLA Discord.

Resources Copy Chapter

• PJRT C API header
• PJRT C API changelog
• PJRT integration guide
• PJRT design docs
• PJRT API ABI versioning and compatibility
• PJRT Plugin Mechanism design doc
• OpenXLA 2024 Fall DevLab PJRT plugin tutorial slides
  (recording)
• IREE PJRT plugin implementation

Tf2xla Index (`docs/tf2xla/index.md`) Copy Section

XLA: Optimizing Compiler for Machine Learning

OpenXLA is a domain-specific compiler for linear
algebra that can accelerate TensorFlow models with potentially no source code
changes.

Introduction Copy Chapter

When a TensorFlow program is run, all of the operations are executed
individually by the TensorFlow executor. Each TensorFlow operation has a
precompiled GPU kernel implementation that the executor dispatches to.

XLA provides an alternative mode of running models: it compiles the TensorFlow
graph into a sequence of computation kernels generated specifically for the
given model. Because these kernels are unique to the model, they can exploit
model-specific information for optimization. For example, let's look at an
optimization XLA does in the context of a simple TensorFlow computation:

def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)

Run without XLA, the graph launches three kernels: one for the multiplication,
one for the addition and one for the reduction. However, XLA can optimize the
graph so that it computes the result in a single kernel launch. It does this by
"fusing" the addition, multiplication and reduction into a single GPU kernel.
Moreover, this fused operation does not write out the intermediate values
produced by y*z and x+y*z to memory; instead it "streams" the results of
these intermediate computations directly to their users while keeping them
entirely in GPU registers. Fusion is XLA's single most important optimization.
Memory bandwidth is typically the scarcest resource on hardware accelerators, so
removing memory operations is one of the best ways to improve performance.

Enable XLA for TensorFlow models Copy Chapter

Explicit compilation with `tf.function(jit_compile=True)` Copy Section

Explicit compilation API offers a fine-grained control for choosing which
functions should be compiled. For example, the following TensorFlow function
which performs the MNIST training is compiled with XLA:

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

The jit_compile API has must-compile semantics: either the entire
function is compiled with XLA, or an errors.InvalidArgumentError exception is
thrown. XLA can not currently compile functions where dimensions are not
inferrable: that is, if it's not possible to infer the dimensions of all
tensors without running the entire computation. For example, the following
function will not compile:

@tf.function
def not_compilable(x):
  return tf.unique(x)

Shapes can vary across the runs though:

@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

Note: Nesting behavior: the function will be compiled if at least one function
in its call stack has jit_compile=True.

See the tutorial colab for a more detailed
usage example, and a
tutorial video on
jit_compile=True usage.

Usage with Keras Copy Section

For Keras models, jit_compile=True can be set as an argument to
model.compile:

model.compile(optimizer="adam", jit_compile=True)

Usage with distributed strategy Copy Section

XLA:GPU can be used with TF distributed strategy
(MirroredStrategy
or
MultiWorkerMirroredStrategy)
by annotating step function with jit_compile=True:

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

Auto-clustering Copy Section

A simple way to start using XLA in TensorFlow models without any changes is to
enable auto-clustering, which automatically finds clusters (connected
subgraphs) within the TensorFlow functions which can be compiled and executed
using XLA. Auto-clustering on GPU can be enabled by setting the TF_XLA_FLAGS
environment variable:

Note: In TF2, only the code inside tf.function will be clustered.

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

Auto-clustering is currently optimized for GPU workloads, but it can also be
enabled on CPU by additionally using the flag --tf_xla_cpu_global_jit:

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

Note: Auto-clustering support on CPU and on multi-GPU environments is
experimental.

For a detailed usage example see the
auto-clustering tutorial colab.

Inspect compiled programs Copy Chapter

XLA provides introspection facilities which let you inspect the generated
programs. To dump the generated programs, use the environment variable
XLA_FLAGS:

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

After the dumping is performed, you can find the following files in
/tmp/generated:

• module_XXXX.*_optimizations.txt Generated
  XLA programs, one per each compiled cluster.
  Attaching those when submitting XLA bug reports is extremely helpful!

• module_XXXX.ir-*.ll Generated files in
  LLVM intermediate representation, with
  NVPTX intrinsics.

• module_XXXX.ptx Generated
  PTX
  files.

You can also dump the graph visualizing the embedding of XLA clusters inside of
the TensorFlow graph with:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

Reproducible bug reports Copy Chapter

A bug report is much easier to reproduce if it includes dumps for the generated
XLA programs and the used auto-clustering embedding.
To generate them for a TensorFlow program running with auto-clustering, launch:

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"

When filing bugs, attach the contents of the /tmp/generated directory
(referenced above).

If possible, try to isolate
a bug to a single XLA program by using the
run_hlo_module
and iteratively running it on generated programs.

Further reading Copy Chapter

• OpenXLA Documentation OpenXLA Documentation
• Known Issues List of known issues with XLA+TF
• XLA - TensorFlow, Compiled:
  Read on Google Developers Blog
• Check out the
  XLA source
  on Github!

XLA Frontends Copy Chapter

Apart from TensorFlow, XLA programs can be generated by:

• JAX: Composable transformations of
  Python+NumPy programs
• Julia: The Julia language for
  scientific computing
• PyTorch: PyTorch framework
• Nx: Numerical computing library for the
  Elixir programming language

Talks Copy Chapter

Using XLA from TF using `jit_compile=True` Copy Section

XLA Overview Copy Section

Aliasing (`docs/aliasing.md`) Copy Section

Aliasing in XLA

This document describes the XLA aliasing API, which lets you specify the
aliasing between the input and output buffers when building an XLA program.

Defining aliasing at compile-time Copy Chapter

For example, consider a trivial HLO module which simply adds 1 to its input:

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

This module will allocate two 4-byte buffers: one for the input %p, and one
for the output %out.

However, it is often desirable to perform the update in-place (for example, if
in the frontend generating the expression the input variable is no longer alive
after the computation, as in the increment p++).

To perform such an update efficiently, you can specify the input aliasing:

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

The format specifies that the entire output (marked by {}) is aliased to the
input parameter 0.

To specify the aliasing programmatically, see the
XlaBuilder::SetUpAlias
API.

Defining aliasing at runtime Copy Chapter

The aliasing defined in the previous step is specified during compilation.
During execution, you can use the
LocalClient::RunAsync
API to choose whether to donate the buffer.

Input buffers to the program are wrapped in
ExecutionInputs,
which in turn contain a tree of MaybeOwningDeviceMemory. If memory is
specified as owning (ownership of the buffer is passed to the XLA runtime),
the buffer is actually donated, and the update is executed in place, as
requested by the compile-time aliasing API.

If, however, the buffer that is aliased at compile time is not donated at
runtime, copy-protection kicks in: an extra output buffer O is allocated,
and the contents of the input buffer P that was meant to be aliased are copied
into O (so effectively the program can execute as if the buffer O was
donated at runtime).

Architecture (`docs/architecture.md`) Copy Section

XLA architecture

XLA (Accelerated Linear Algebra) is a machine learning (ML) compiler that
optimizes linear algebra, providing improvements in execution speed and memory
usage. This page provides a brief overview of the objectives and architecture of
the XLA compiler.

Objectives Copy Chapter

Today, XLA supports several ML framework frontends (including PyTorch,
TensorFlow, and JAX) and is part of the OpenXLA project – an ecosystem of
open-source compiler technologies for ML that's developed collaboratively by
leading ML hardware and software organizations. Before the OpenXLA project was
created, XLA was developed inside the TensorFlow project, but the fundamental
objectives remain the same:

• Improve execution speed. Compile subgraphs to reduce the execution time
  of short-lived ops and eliminate overhead from the runtime, fuse pipelined
  operations to reduce memory overhead, and specialize known tensor shapes to
  allow for more aggressive constant propagation.

• Improve memory usage. Analyze and schedule memory usage, eliminating
  many intermediate storage buffers.

• Reduce reliance on custom ops. Remove the need for many custom ops by
  improving the performance of automatically fused low-level ops to match the
  performance of custom ops that were originally fused by hand.

• Improve portability. Make it relatively easy to write a new backend for
  novel hardware, so that a large fraction of ML models can run unmodified on
  that hardware. This is in contrast with the approach of specializing
  individual monolithic ops for new hardware, which requires models to be
  rewritten to make use of those ops.

How it works Copy Chapter

The XLA compiler takes model graphs from ML frameworks defined in
StableHLO and compiles them into machine
instructions for various architectures. StableHLO defines a versioned operation
set (HLO = high level operations) that provides a portability layer between ML
frameworks and the compiler.

In general, the compilation process that converts the model graph into a
target-optimized executable includes these steps:

1. XLA performs several built-in optimization and analysis passes on the
   StableHLO graph that are target-independent, such as
   CSE,
   target-independent operation fusion, and buffer analysis for allocating
   runtime memory for the computation. During this optimization stage, XLA also
   converts the StableHLO dialect into an internal HLO dialect.

2. XLA sends the HLO computation to a backend for further HLO-level
   optimizations, this time with target-specific information and needs in mind.
   For example, the GPU backend may perform operation fusions that are
   beneficial specifically for the GPU programming model and determine how to
   partition the computation into streams. At this stage, backends may also
   pattern-match certain operations or combinations thereof to optimized
   library calls.

3. The backend then performs target-specific code generation. The CPU and GPU
   backends included with XLA use LLVM for low-level IR,
   optimization, and code generation. These backends emit the LLVM IR necessary
   to represent the HLO computation in an efficient manner, and then invoke
   LLVM to emit native code from this LLVM IR.

Within this process, the XLA compiler is modular in the sense that it is easy to
slot in an alternative backend to
target some novel HW architecture. The GPU
backend currently supports NVIDIA GPUs via the LLVM NVPTX backend. The CPU
backend supports multiple CPU ISAs.

Async Ops (`docs/async_ops.md`) Copy Section

Async HLO Instructions

1. Adding async operations to HLO is cumbersome (i.e. all-reduce-start and
   all-reduce-done).
2. The start and done split may be inadequate for some of the asynchronous use
   cases.

To target the first shortcoming, we propose to introduce one last set of new
asynchronous opcodes: kAsyncStart, kAsyncUpdate, and kAsyncDone. The idea
is to create a generic asynchronous opcode that can wrap any HLO instruction.
The actual operation that will be performed asynchronously will be encoded using
a called computation that only has the instruction as its root and any
parameters for inputs. The in-flight input/output buffer handling and aliasing
can then be shared for any asynchronous operation. The async-start instruction’s
output shape will then be a tuple of the input operands, output values, and any
intermediate state that is needed for the async-update or async-done
instructions.

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(f32[64] %operand),
                                           calls=%async_op
%async-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-start)

In the representation above, only async-start has a called computation since
it is trivial to find what the async-done does by following its operand to
find the corresponding async-start to find the called computation.

Also note that the first element in the output tuple of async-start is a
tuple containing the operands. The elements of this operand tuple alias with
the respective operands, so their buffers stay alive until at least the
async-done instruction. Similarly, the second element aliases with the output
of async-done, and the third element is the context state that is used to
keep track of the asynchronous operation. This representation naturally
supports multiple tensors in the asynchronous operation input and/or output:

%async_op {
  %param0 = f32[64] parameter(0)
  %param1 = f32[64] parameter(1)
  ROOT %op = (f32[32], f32[32]) op(f32[64] %param0, f32[64] %param1),
                                op_specific_attr=”foo”
}

%async-start = ((f32[64], f32[64]), (f32[32], f32[32]), s32[])
               async-start(f32[64] %operand0, f32[64] %operand1),
               calls=%async_op
%async-done = (f32[32], f32[32]) async-done(%async-start)

In addition, the op can further be decomposed into zero or more async-update
steps that perform intermediate computations. The input/output aliasing works
the same way with the async-update instruction and each async-start and
async-update instructions must have one user that is either another
async-update or an async-done:

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=”foo”
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-update0 = ((f32[64]), f32[32], s32[]) async-update(
                           ((f32[64]), f32[32], s32[]) %async-start)
%async-update1 = ((f32[64]), f32[32], s32[]) async-update(
                           ((f32[64]), f32[32], s32[]) %async-update0)
%async-done = f32[32] async-done(((f32[64]), f32[32], s32[]) %async-update1)

Syntax sugar Copy Chapter

The HLO parser supports syntax sugar to automatically parse and print
asynchronous operations as if they are first-class opcodes. The parser treats
the -start, -update, and -done suffixes specially by automatically
creating the async computation and the wrapped instruction (without the suffix).

For example, an asynchronous custom-call can be written as:

%cc-start = ((f32[64]), f32[32], s32[]) custom-call-start(%operand),
                                        custom_call_target="foo"
%cc-update = ((f32[64]), f32[32], s32[]) custom-call-update(%cc-start)
%result = f32[32] custom-call-done(%cc-update)

The parser desugars this into the following equivalent HLO:

%async_computation {
  %p0 = f32[64] parameter(0)
  ROOT %custom-call = f32[32] custom-call(%p0), custom_call_target="foo"
}

%async-start = ((f32[64]), f32[32], s32[]) async-start(%operand),
                calls=%async_computation
%async-update = ((f32[64]), f32[32], s32[]) async-update(%async-start)
%result = f32[32] async-done(%async-update)

This desugaring is supported for most HLO opcodes (e.g., custom-call, dot,
all-reduce, etc.).

Exceptions Copy Section

In order not to create ambiguities, the parser will not desugar operations that
have explicit first-class opcodes defined with the -start and/or -done
suffixes (e.g., copy-start/copy-done,
collective-permute-start/collective-permute-done). These will continue to
use their respective first-class opcodes.

Late Binding Copy Chapter

In some cases, the operands (inputs) or outputs of an asynchronous
operation are not all available or allocated when the operation starts.
XLA supports late binding, which allows operands to be incrementally
bound during async-update steps, and outputs to be bound during either
async-update or async-done steps.

Representation in HLO Copy Section

For a called computation that expects $N$ parameters, we can start the
asynchronous execution with fewer than $N$ operands. The remaining
operands are passed in subsequent async-update instructions.

• async-start binds the first $K$ operands ($K < N$).
• async-update instructions bind the remaining $N - K$ operands.

Operand bindings must happen in left-to-right order. That is, if a
computation expects parameters $P_0, P_1, \dots, P_{N-1}$, they must be bound in
that order across the async chain.

The async-start and async-update shapes reflect the incrementally
bound parameters. Specifically, the first element of the tuple shape
(the operand shapes) grows as more operands are bound.

Output binding is independent of operand binding and can happen at any
step in the async chain (either in an async-update or at the final
async-done).

Example with `kCall` Copy Section

Consider a called computation %foo that takes two parameters:

%foo {
  %p0 = f32[] parameter(0)
  %p1 = f32[] parameter(1)
  ROOT %add = f32[] add(%p0, %p1)
}

We can call this computation asynchronously, binding %p0 at start and
%p1 at update:

%call-start = ((f32[]), (), s32[]) call-start(%operand0), to_apply=%foo
%call-update = ((f32[], f32[]), f32[], s32[]) call-update(%call-start, %operand1)
%result = f32[] call-done(%call-update)

The parser desugars this into the following HLO:

%async-start = ((f32[]), (), s32[]) async-start(%operand0), calls=%foo
%async-update = ((f32[], f32[]), f32[], s32[]) async-update(%async-start, %operand1)
%result = f32[] async-done(%async-update)

Late-Bound Outputs Copy Section

In addition to operands (inputs), the outputs of an asynchronous
operation can also be bound late. This is useful when the output
buffers are not known or allocated at the start of the operation.

To represent late-bound outputs:

1. The async-start (or call-start) instruction is defined with an
   empty tuple () at index 1 of its output shape (the result slot).
2. A subsequent async-update (or call-update) instruction
   specifies the actual output shape at index 1, replacing the empty
   tuple.
3. Alternatively, the output can be bound at the end of the chain by
   the async-done (or call-done) instruction, which returns the
   final output shape. This can be done regardless of whether there are
   intermediate async-update steps in the chain.

Example with `async-update`

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at update (index 1 becomes (f32[1024]))
%call-update = ((f32[1024]), (f32[1024]), s32[]) call-update(%call-start, %output_buffer)

%result = (f32[1024]) call-done(%call-update)

The parser desugars this into:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%async-update = ((f32[1024]), (f32[1024]), s32[]) async-update(%async-start, %output_buffer)
%result = (f32[1024]) async-done(%async-update)

Example with `async-done` (without `async-update`)

If there are no intermediate update steps, the output can be bound
directly at async-done:

// Output is not bound at start (index 1 is ())
%call-start = ((f32[1024]), (), s32[]) call-start(%input_buffer), to_apply=%foo

// Output is bound at done
%result = (f32[1024]) call-done(%call-start)

The parser desugars this into:

%async-start = ((f32[1024]), (), s32[]) async-start(%input_buffer), calls=%foo
%result = (f32[1024]) async-done(%async-start)

Example with intermediate `async-update` and output bound at `async-done`

If there are intermediate update steps to bind operands, but the output
is still bound at the very end:

// Output is not bound at start, no operands bound
%call-start = ((), (), s32[]) call-start(), to_apply=%foo

// Operands are bound at update, but output remains unbound (index 1 is ())
%call-update = ((f32[], f32[]), (), s32[]) call-update(%call-start, %operand0, %operand1)

// Output is bound at done
%result = f32[] call-done(%call-update)

The parser desugars this into:

%async-start = ((), (), s32[]) async-start(), calls=%foo
%async-update = ((f32[], f32[]), (), s32[]) async-update(%async-start, %operand0, %operand1)
%result = f32[] async-done(%async-update)

Avoiding Ifdefs (`docs/avoiding_ifdefs.md`) Copy Section

The Hitchhiker's Guide to fewer `#ifdef`s in XLA

TL;DR: #ifdefs are bad for maintainability. Please try to avoid them.
Scroll down for details on how.

[TOC]

Introduction Copy Chapter

Preprocessor conditionals (#if, #ifdef, #ifndef, etc.) are a convenient
way to make existing code work in a different environment than it was originally
intended for. Examples include:

• Making existing CUDA-specific code work in a ROCm or SYCL environment.
• Adding new library features (for example, from a newer cuDNN or hipDNN
  release) while maintaining support for older versions.

However, they come with a high maintenance cost (more details in the next
section), especially when refactoring existing code.

This guide explains the downsides of preprocessor conditionals and offers
alternatives for the most common use cases in a cookbook style. It is intended
to help contributors design maintainable, platform-portable changes from the
start and to serve as a shared reference during code reviews.

Motivation Copy Chapter

The C preprocessor is the first stage in a C++ compilation pipeline.
Preprocessor conditionals manipulate the stream of preprocessor tokens that is
read from the source code and header files and handed over to the C++ compiler.
The fact that it runs before the compiler and operates on the preprocessor token
level is the main reason why it results in a hard-to-maintain codebase:

• Code Complexity: #ifdefs are convenient to use because they can be
  inserted almost anywhere in the code. Since they operate on tokens, there
  are almost no syntactic or semantic limits. Developers don't need to think
  about a proper abstraction; they can just insert the conditionals right
  where they are needed. Not being forced to design a proper abstraction makes
  the code harder to read and harder to reason about. Having even a simple
  abstraction like a free function allows easier testing of this function or
  mocking its behavior—either now or later.

  It also renders a subset of compiler warnings meaningless because the
  compiler will only ever see one evaluated token stream. A common example is
  the need for the [[maybe_unused]] attribute on a function parameter that
  is only used in one branch of a preprocessor conditional.

• Testability (Build): Since preprocessor conditionals operate on token
  streams, all the untaken branches neither need to be syntactically nor
  semantically correct and are therefore not checked for correctness. This
  severely hinders larger refactorings where a developer applies a
  find-and-replace and relies on the compiler to tell them where to make
  manual fixups. Incorrect changes in uncompiled preprocessor branches either
  go in unnoticed or get detected late in the process (if CI happens to build
  that particular preprocessor branch). The latter is also relevant for
  day-to-day development (see the next point).

• Testability (Coverage): Having more preprocessor conditionals means
  having more build configurations to test, and the number of configurations
  increases exponentially with every new condition introduced. It is already
  impossible to test all of them. For example, XLA has quite a few
  conditionals based on cuDNN version numbers that are not all exercised in
  CI. One could argue that this is less relevant as long as we test the
  configurations we care about—which is probably true. But unnecessarily
  broken code leads to bug reports that someone needs to address, even if it
  is just saying that a configuration is not supported.

  More importantly, potentially broken build configurations increase the
  developer's iteration dead time. A change might work fine when running
  bazel test locally, but CI might build a slightly different configuration
  and fail. Fixes are often benign—like adding a [[maybe_unused]]
  attribute—but this additional round trip costs the developer extra time on
  every pull request.

  Therefore, decreasing the number of preprocessor conditionals decreases
  the number of build configurations, which in turn decreases the likelihood
  of additional CI round trips on a PR.

• Tooling: C++ parsing is such a complex task that most tools nowadays
  rely on a compiler frontend for parsing and semantic analysis and then
  operate directly on the AST. Notable examples include clang-tidy,
  include-cleaner, and language servers (clangd) that provide code
  completion, navigation, and syntax highlighting in your IDE. Another class
  of tools relies on compiler instrumentation, including sanitizers and code
  coverage tools.

  All of these tools only see a single build configuration, so extensive use
  of preprocessor conditionals hinders their usability. include-cleaner, for
  example, suggests removing #includes that are only used in an unevaluated
  preprocessor branch. Similarly, your IDE won't show syntax highlighting or
  code navigation for ROCm code when it is configured for a CUDA or CPU build,
  making it a pain to edit.

Mitigations Copy Chapter

Category I - Skipping test cases in unit tests Copy Section

It is very common that a certain test case is only supported:

• On a certain backend.
• With a certain GPU model.
• When library X is at least of version Y.

Previously, it was common to skip tests using preprocessor conditionals:

TEST(Foo, Bar) {
#ifdef TENSORFLOW_USE_ROCM
  GTEST_SKIP();
#endif
  // ...
}

Alternative: Ask StreamExecutor

StreamExecutor is XLA's hardware abstraction layer, and its
stream_executor::DeviceDescription
has the information needed to make the same decision at runtime:

TEST_F(FooTest, Bar) {
  // `device_description()` is provided by `HloPjRtGpuTestBase`. Other test
  // fixtures expose it via `executor->GetDeviceDescription()`.
  const se::DeviceDescription& device = device_description();

  // Skip based on the backend platform.
  if (device.gpu_compute_capability().IsRocm()) {
    GTEST_SKIP() << "Not supported on ROCm.";
  }

  // Skip based on the GPU model / compute capability.
  if (const auto* cc =
          device.gpu_compute_capability().cuda_compute_capability();
      cc != nullptr && !cc->IsAtLeastHopper()) {
    GTEST_SKIP() << "Requires Hopper or newer.";
  }

  // Skip based on the runtime or library version.
  if (device.runtime_version() < se::SemanticVersion{12, 2, 0}) {
    GTEST_SKIP() << "Requires CUDA runtime >= 12.2.";
  }
  if (device.dnn_version() < se::SemanticVersion{9, 0, 0}) {
    GTEST_SKIP() << "Requires cuDNN >= 9.0.";
  }
  // ...
}

DeviceDescription exposes both compute capabilities and structured version
numbers via
se::SemanticVersion:

• Backend & architecture: device.gpu_compute_capability().IsCuda(),
  device.gpu_compute_capability().IsRocm(),
  device.gpu_compute_capability().IsOneAPI(), and accessors for
  CudaComputeCapability, RocmComputeCapability, and
  OneAPIComputeCapability.
• Runtime & driver versions: device.runtime_version(),
  device.driver_version(), device.kernel_mode_driver_version(),
  device.compile_time_toolkit_version().
• Library versions: device.dnn_version(), device.cub_version().

If a specific version or hardware property is not yet available in
DeviceDescription, please add it (or ask the contributor to add it) to
DeviceDescription instead of falling back to #ifdef.

Long term (author's opinion)

In the long term, none of our higher-level tests should need to make decisions
based on backends, hardware variants, or runtime/driver versions. Instead, tests
should ask StreamExecutor whether a certain feature is available, and
StreamExecutor determines that based on all the necessary details. All this
logic should live in one place, though this does not yet have a concrete design.

Category II - Could be a runtime `if` statement Copy Section

Another class of preprocessor conditionals guards code that would also compile
fine without the conditional:

void Foo::Bar() {
#if TENSORFLOW_USE_ROCM
  // Do something that would also compile in CUDA/CPU mode
#endif
  // ...
}

Alternative: Use a runtime `if`

Replace the preprocessor conditional with a runtime conditional:

void Foo::Bar() {
  if (stream_executor_.GetDeviceDescription()
          .gpu_compute_capability()
          .IsRocm()) {
    // Do something that compiles fine everywhere
  }
  // ...
}

Whether this should all be factored out into a separate function is up to the
reviewer's discretion. At this point it is just "normal" code and normal code
review practices apply.

This class of preprocessor conditionals appears most often when checking a
condition that is not readily available in StreamExecutor's
DeviceDescription. Everyone is encouraged to ask contributors to add the
relevant pieces of information to DeviceDescription instead of accepting the
preprocessor conditional.

Category III - Requires access to the low-level runtime Copy Section

The most common (and also the best justified) class of preprocessor conditionals
are those that guard code that wouldn't compile otherwise. Most often this is
because it uses something from a low-level backend-specific header (such as the
CUDA or ROCm API):

#if TENSORFLOW_USE_ROCM
#include <something/something/rocm.h>
#endif

void Foo::Bar() {
#if TENSORFLOW_USE_ROCM
  // Do something that would *NOT* compile in CUDA/CPU mode
#endif
  // ...
}

The same applies to version guards on library headers (for example,
#if CUDNN_VERSION >= 90000) when the code references symbols that do not exist
in older versions of the header.

Alternative I: Let StreamExecutor deal with it

The first question to ask is whether the changed component should depend on
low-level runtime specifics at all. Ideally, only our hardware abstraction layer
StreamExecutor should. So if this code is not part of StreamExecutor, the
first option is to see if it can be pushed down into xla/stream_executor/.

That means creating an API (a function, a class, or whatever makes sense) in
StreamExecutor that exists for all backends. This API can have different
implementations based on the backend. It might also allow querying whether this
particular feature is available, which enables the runtime-if solution
described above.

Of course, that means StreamExecutor now has to deal with accessing the
low-level backend APIs. Alternative II describes how this can be modeled without
(too many) #ifdefs.

Alternative II: Code in separate targets

If the usage of backend-specific headers can't be pushed into StreamExecutor
or the change is inside StreamExecutor, the best option is to split the code
into separate compilation units. In our example, we have one build target with a
single C++ source file containing the ROCm code and another build target with a
single C++ source file containing the non-ROCm code—typically a stub
implementation that returns absl::UnimplementedError on all the defined
functions.

This has the following advantages:

• Splitting code across multiple files means there needs to be some layer of
  abstraction in between—at the very least a free function. This encourages
  contributors and reviewers alike to think about the right level of
  abstraction, leading to higher-quality code in the long term.
• Both files only contain a single build configuration—giving you full syntax
  highlighting and IDE tooling in both of them.
• Both build targets can be built separately in the same build graph. No need
  to change your Bazel build flags and evict your build cache.
• Both build targets can have their own tests. For example, if a certain API
  only exists for ROCm, we can have unit tests only testing the ROCm
  implementation instead of higher-level tests that get skipped in all other
  cases, which can speed up testing. (To be fair: in many cases there are good
  reasons for having the higher-level tests as well.)

A concrete design could look like this:

// feature.h
// Do NOT include any platform-specific headers (CUDA/ROCm) here.
#include "absl/status/status.h"

absl::Status Foo();

// feature_rocm.cc
#include "feature.h"
#include "something/something/rocm.h"

absl::Status Foo() {
  // Do something ROCm-specific.
}

// feature_stub.cc
#include "feature.h"
#include "absl/status/status.h"

absl::Status Foo() {
  return absl::UnimplementedError("This is a ROCm-only feature.");
}

# BUILD
load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured")

cc_library(
    name = "feature_rocm",
    srcs = [
        "feature.h",
        "feature_rocm.cc",
    ],
    tags = ["manual"],  # Exclude this from wildcard builds when ROCm is not enabled
    deps = [
        "@com_google_absl//absl/status",
        "@local_config_rocm//rocm:rocm_headers",
    ],
)

cc_library(
    name = "feature_stub",
    srcs = [
        "feature.h",
        "feature_stub.cc",
    ],
    deps = ["@com_google_absl//absl/status"],
)

cc_library(
    name = "feature",
    hdrs = ["feature.h"],
    deps = if_rocm_is_configured(
        [":feature_rocm"],
        [":feature_stub"],
    ) + [
        "@com_google_absl//absl/status",
    ],
)

Note that feature.h is only in the hdrs attribute of the :feature target
(and in srcs of the two implementation targets). This ensures that consumers
including feature.h (and automated dependency tools) depend on :feature
rather than directly on :feature_rocm or :feature_stub.

It is also important that feature.h does not include any platform-specific
headers (such as CUDA, ROCm, or cuDNN headers); only .cc files (and private
headers) should include those. The same build target pattern works for CUDA
(if_cuda_is_configured from //xla/tsl/platform/default:cuda_build_defs.bzl)
and SYCL.

Long term (author's opinion)

In the previous example, the targets :feature_rocm and :feature_stub define
the same symbol Foo, so they can't be linked into the same binary without
causing an ODR violation. Long term, XLA could have a plugin infrastructure
where multiple backends can be loaded at the same time. In that case, the design
above needs to be changed slightly so that :feature_rocm and :feature_stub
define symbols with different names and :feature has a dispatch function with
the original name—though it is currently unclear whether this will become a high
enough priority.

Broadcasting (`docs/broadcasting.md`) Copy Section

Broadcasting

This document describes the broadcasting semantics of XLA.

What is broadcasting? Copy Chapter

Broadcasting is the process of making arrays with different shapes have
compatible shapes for arithmetic operations. The terminology is borrowed from
NumPy broadcasting.

Broadcasting may be required for operations between multi-dimensional arrays of
different ranks, or between multi-dimensional arrays with different but
compatible shapes. Consider the addition X+v where X is a matrix (an array
with 2 dimensions) and v is a vector (an array with 1 dimension). To perform
element-wise addition, XLA needs to "broadcast" the vector v to the same
number of dimensions as the matrix X, by replicating v a certain number of
times. The vector's length has to match at least one of the dimensions of the
matrix.

For example:

|1 2 3| + |7 8 9|
|4 5 6|

The matrix's dimensions are (2,3), and the vector's dimension is (3). The vector
is broadcast by replicating it over rows to get:

|1 2 3| + |7 8 9| = |8  10 12|
|4 5 6|   |7 8 9|   |11 13 15|

In NumPy, this is called
broadcasting.

Principles Copy Chapter

The XLA language is as strict and explicit as possible, avoiding implicit
"magical" features. Such features might make some computations slightly easier
to define, but at the cost of more assumptions baked into user code that will be
difficult to change in the long term. If necessary, implicit magical features
can be added in client-level wrappers.

With regard to broadcasting, XLA requires explicit broadcasting specifications
on operations between arrays of different ranks. This is different from NumPy,
which infers the specification when possible.

Broadcasting a lower-dimensional array onto a higher-dimensional array Copy Chapter

Scalars can always be broadcast over arrays without an explicit specification
of broadcasting dimensions. An element-wise binary operation between a scalar
and an array means applying the operation with the scalar to each element in the
array. For example, adding a scalar to a matrix means producing a matrix in
which each element is a sum of the scalar and the corresponding element of the
input matrix.

|1 2 3| + 7 = |8  9  10|
|4 5 6|       |11 12 13|

Most broadcasting needs can be captured by using a tuple of dimensions on a
binary operation. When the inputs to the operation have different ranks, this
broadcasting tuple specifies which dimension(s) in the higher-dimensional
array to match with the lower-dimensional array.

Consider the previous example. Instead of adding a scalar to a (2,3) matrix, add
a vector of dimension (3) to a matrix of dimensions (2,3). Without specifying
broadcasting, this operation is invalid. To correctly request matrix-vector
addition, specify the broadcasting dimension to be (1), meaning the vector's
dimension is matched to dimension 1 of the matrix. In 2D, if dimension 0
represents rows and dimension 1 represents columns, this means that each element
of the vector becomes a column of a size matching the number of rows in the
matrix:

|7 8 9| ==> |7 8 9|
            |7 8 9|

As a more complex example, consider adding a 3-element vector (dimension (3)) to
a 3x3 matrix (dimensions (3,3)). There are two ways broadcasting can happen for
this example:

(1) A broadcasting dimension of 1 can be used. Each vector element becomes a
column and the vector is duplicated for each row in the matrix.

|7 8 9| ==> |7 8 9|
            |7 8 9|
            |7 8 9|

(2) A broadcasting dimension of 0 can be used. Each vector element becomes a row
and the vector is duplicated for each column in the matrix.

|7| ==> |7 7 7|
 |8|     |8 8 8|
 |9|     |9 9 9|

Note: When adding a 2x3 matrix to a 3-element vector, a broadcasting dimension
of 0 is invalid.

The broadcasting dimensions can be a tuple that describes how a
smaller-dimensional shape is broadcast into a larger-dimensional shape. For
example, given a 2x3x4 cuboid and a 3x4 matrix, a broadcasting tuple (1,2)
means matching the matrix to dimensions 1 and 2 of the cuboid.

This type of broadcast is used in the binary ops in XlaBuilder, if the
broadcast_dimensions argument is given. For example, see
XlaBuilder::Add.
In the XLA source code, this type of broadcasting is sometimes called "InDim"
broadcasting.

Formal definition Copy Section

The broadcasting attribute allows matching a lower-dimensional array to a
higher-dimensional array by specifying which dimensions of the
higher-dimensional array to match. For example, for an array with dimensions
MxNxPxQ, a vector with dimension T can be matched as follows:

MxNxPxQ

dim 3:          T
dim 2:        T
dim 1:      T
dim 0:    T

In each case, T has to be equal to the matching dimension of the
higher-dimensional array. The vector's values are then broadcast from the
matched dimension to all the other dimensions.

To match a TxV matrix onto the MxNxPxQ array, a pair of broadcasting dimensions
is used:

MxNxPxQ
dim 2,3:      T V
dim 1,2:    T V
dim 0,3:  T     V
etc...

The order of dimensions in the broadcasting tuple must be the order in which the
dimensions of the lower-dimensional array are expected to match the dimensions
of the higher-dimensional array. The first element in the tuple specifies which
dimension in the higher-dimensional array has to match dimension 0 in the
lower-dimensional array. The second element in the tuple specifies which
dimension in the higher-dimensional array has to match dimension 1 in the
lower-dimensional array, and so on. The order of broadcast dimensions must be
strictly increasing. For example, in the previous example it is illegal to
match V to N and T to P; it is also illegal to match V to both P and N.

Broadcasting similar-dimensional arrays with degenerate dimensions Copy Chapter

A related problem is broadcasting two arrays that have the same number of
dimensions but different dimension sizes. As with NumPy, this is only possible
when the arrays are compatible. Two arrays are compatible when all their
dimensions are compatible. Two dimensions are compatible if:

• They are equal, or
• One of them is 1 (a "degenerate" dimension)

When two compatible arrays are encountered, the result shape has the maximum of
the two inputs at every dimension index.

Examples:

1. (2,1) and (2,3) broadcast to (2,3).
2. (1,2,5) and (7,2,5) broadcast to (7,2,5).
3. (7,2,5) and (7,1,5) broadcast to (7,2,5).
4. (7,2,5) and (7,2,6) are incompatible and cannot be broadcast.

A special case arises, and is also supported, where each of the input arrays has
a degenerate dimension at a different index. In this case, the result is an
"outer operation": (2,1) and (1,3) broadcast to (2,3). For more examples,
consult the
NumPy documentation on broadcasting.

Broadcast composition Copy Chapter

Broadcasting of a lower-dimensional array to a higher-dimensional array and
broadcasting using degenerate dimensions can both be performed in the same
binary operation. For example, a vector of size 4 and a matrix of size 1x2 can
be added together using broadcast dimensions of value (0):

|1 2 3 4| + [5 6]    // [5 6] is a 1x2 matrix, not a vector.

First the vector is broadcast up to 2-dimensional (matrix) using the broadcast
dimensions. The single value (0) in the broadcast dimensions indicates that
dimension zero of the vector matches dimension zero of the matrix. This produces
a matrix of size 4xM where the value M is chosen to match the corresponding
dimension size in the 1x2 array. Therefore, a 4x2 matrix is produced:

|1 1| + [5 6]
|2 2|
|3 3|
|4 4|

Then "degenerate dimension broadcasting" broadcasts dimension zero of the 1x2
matrix to match the corresponding dimension size of the right hand side:

|1 1| + |5 6|     |6  7|
|2 2| + |5 6|  =  |7  8|
|3 3| + |5 6|     |8  9|
|4 4| + |5 6|     |9 10|

A more complicated example is a matrix of size 1x2 added to an array of size
4x3x1 using broadcast dimensions of (1, 2). First the 1x2 matrix is broadcast up
to 3 dimensions using the broadcast dimensions to produce an intermediate Mx1x2
array where the dimension size M is determined by the size of the larger operand
(the 4x3x1 array) producing a 4x1x2 intermediate array. The M is at dimension 0
(the left-most dimension) because the dimensions 1 and 2 are mapped to the
dimensions of the original 1x2 matrix as the broadcast dimensions are (1, 2).
This intermediate array can be added to the 4x3x1 matrix using broadcasting of
degenerate dimensions to produce a 4x3x2 array result.

Build From Source (`docs/build_from_source.md`) Copy Section

Build from source

This document describes how to build XLA components.

If you did not clone the XLA repository or install Bazel, check out the initial
sections of the XLA Developer Guide.

Linux Copy Chapter

Configure Copy Section

XLA builds are configured by the .bazelrc file in the repository's root
directory. The ./configure.py script can be used to adjust common settings.

If you need to change the configuration, run the ./configure.py script from
the repository's root directory. This script has flags for the location of XLA
dependencies and additional build configuration options (compiler flags, for
example). Refer to the Sample session section for details.

CPU support Copy Section

We recommend using a suitable Docker image - such as
ml-build
, which is also used in XLA's CI workflows on GitHub - for building and testing
XLA. The ml-build image comes with Clang 18 pre-installed.

docker run -itd --rm \
--name xla \
-w /xla \
-v $PWD:/xla \
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
bash

Using a Docker container, you can build XLA with CPU support by running the
following commands:

docker exec xla ./configure.py --backend=CPU

docker exec xla bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

If you want to build XLA targets with CPU support without using Docker,
you’ll need to install Clang. XLA is currently built with Clang 18 in CI,
but earlier versions should also work.

To configure and build the targets, run the following commands:

./configure.py --backend=CPU

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

GPU support Copy Section

We recommend using the same Docker container mentioned above to build XLA with
GPU support.

To start Docker container with access to all GPUs, run the following command:

docker run -itd --rm \
  --gpus all \
  --name xla_gpu \
  -w /xla \
  -v $PWD:/xla \
  us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
  bash

To build XLA with GPU support, run the following commands:

docker exec xla_gpu ./configure.py --backend=CUDA

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Note: You can build XLA on a machine without GPUs. In that case:

• Do not use --gpus all flag when starting the Docker container.
• Specify CUDA compute capabilities manually, For example:

docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

For more details regarding
TensorFlow's GPU docker images you can check out this document.

You can build XLA targets with GPU support without Docker as well. Configure and
build targets using the following commands:

./configure.py --backend=CUDA

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

For more details regarding
hermetic CUDA you can check out this document.

Build XLA with CUDA/cuDNN Support Using the JAX CI/Release Container Copy Section

XLA is a compiler used internally by JAX.
JAX is distributed via PyPI wheels.
The JAX Continuous Integration documentation
explains how to build JAX wheels using
the tensorflow/ml-build:latest Docker container.

We can extend these instructions to build XLA targets within the JAX container
as well. This ensures that the XLA targets' build configuration is consistent
with the JAX/XLA build configuration, which can be useful if we want to
reproduce workload results using XLA tools that were originally created in JAX.

Build XLA Targets in the JAX CI Container

1. Clone the JAX repository and navigate to the 'jax' directory

git clone https://github.com/jax-ml/jax.git

cd jax

2. Start JAX CI/Release Docker container by running:

./ci/utilities/run_docker_container.sh

This will start a Docker container named 'jax'.

3. Build the jax-cuda-plugin target inside the container using:

docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

This will create the .jax_configure.bazelrc file with the required build
configuration, including CUDA/cuDNN support

4. Access an interactive shell inside the container:

docker exec -ti jax /bin/bash

You should now be in the /jax directory within the container

5. Build the XLA target with the following command, e.g.:

/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

Optionally, you can overwrite HERMETIC envs, e.g.:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

6. Copy the resulting artifacts to /jax/dist to access them from the host OS
   if needed

cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/

7. Exit the interactive shell:

exit

Windows Copy Chapter

Building XLA on Windows natively is a CPU-only process, as CUDA is not
supported directly on Windows; you must use WSL2 if you need CUDA support. It
also requires specific environmental configurations, including a Bash shell, the
Clang compiler, and Visual Studio.

Prerequisites Copy Section

1. Visual Studio: You must install Visual Studio 2019 version 16.5 or newer
   to set up a C++ toolchain (which provides the necessary system headers and
   libraries).
2. Bash Environment: You must use a Bash shell (such as MSYS2 or Git Bash)
   to build XLA on Windows.
3. Clang Compiler: XLA Windows builds use clang-cl rather than the
   standard MSVC compiler. Ensure LLVM/Clang is installed.
4. Python: Python 3 must be installed and available in your system's
   PATH.
5. Developer Mode: You must enable Developer Mode in Windows or run your
   Bash shell as an Administrator. This is required because Bazel relies on
   creating symlinks for the runfiles tree during the build process.

Building from Source Copy Section

To compile the CPU backend of XLA on Windows, run the following command from
your Bash terminal. We use the xla_windows_x86_cpu_2022 config to
automatically set up the clang-cl toolchain.

Note that certain targets and tags must be explicitly excluded from the Windows
CPU build, such as GPU targets and features not yet supported on Windows:

# Note: Ensure your PATH includes Python and Bazel before running this
bazel build \
  --config=xla_windows_x86_cpu_2022 \
  --keep_going \
  --build_tag_filters=-no_oss,-oss_excluded,-gpu,-no_windows,-windows_excluded \
  -- //xla/... -//xla/hlo/experimental/... -//xla/python_api/... -//xla/python/...

Command Buffer Tutorial (`docs/command_buffer_tutorial.md`) Copy Section

Tuning GPU command buffers

Command buffers reduce the CPU work needed to launch GPU operations. XLA records
an eligible sequence of operations into a GPU graph, then replays that graph on
subsequent executions. On NVIDIA GPUs, this uses CUDA Graphs; supported AMD GPU
operations use HIP graphs.

This is especially useful when a training step or inference request repeatedly
launches many small kernels. A workload dominated by a few long-running kernels
may benefit less. Recording, updating, and retaining graphs also have costs, so
measure the complete workload before adopting a setting.

This guide describes upstream XLA main. Your framework's bundled XLA version
may expose different flags or defaults. The defaults below are XLA defaults; a
framework can override them.

Set flags before starting your application Copy Chapter

For applications that accept XLA environment flags, set XLA_FLAGS before
starting the process. In Python, set it before importing or initializing your
XLA-based framework. Restart the process for each configuration you compare.

The examples use Bash and python train.py as a placeholder for your
application:

# Keep the default command types and also allow library collectives.
XLA_FLAGS="--xla_gpu_enable_command_buffer=+COLLECTIVES" python train.py

Separate flags with spaces and list entries with commas. These examples replace
XLA_FLAGS for that invocation; include any other flags your application needs
in the same string. Use the uppercase enum names shown in this guide.

Start with the defaults Copy Chapter

Command buffers are already enabled for several operation types. You do not
need an extra flag to enable the default behavior.

Choose which operations to capture Copy Chapter

--xla_gpu_enable_command_buffer accepts a list of operation categories, not a
boolean. It has three useful forms:

# Add to the defaults.
XLA_FLAGS="--xla_gpu_enable_command_buffer=+COLLECTIVES,+WHILE" python train.py

# Remove one category from the defaults.
XLA_FLAGS="--xla_gpu_enable_command_buffer=-CUSTOM_CALL" python train.py

# Replace the defaults: allow only these categories.
XLA_FLAGS="--xla_gpu_enable_command_buffer=FUSION,CUBLASLT" python train.py

# Disable command buffer conversion for a baseline comparison.
XLA_FLAGS="--xla_gpu_enable_command_buffer=" python train.py

Use either a plain replacement list or a list of +/- modifiers. A plain list
does not preserve omitted defaults. Disabling capture does not disable the
operations themselves: they execute through the normal GPU runtime.

These categories follow the compiled execution path, not the name of your
framework operation. For example, a matrix multiplication implemented as a
generated kernel belongs to FUSION, while a cuBLASLt call belongs to
CUBLASLT.

Enabling a category makes it eligible for capture; it does not force every
operation into a graph. XLA still checks runtime support, dependencies, nested
operations, and the minimum graph size.

Select library collectives Copy Section

To capture only selected library collective types, enable COLLECTIVES and set
--xla_gpu_enable_collectives_command_buffer_filter:

XLA_FLAGS="--xla_gpu_enable_command_buffer=+COLLECTIVES \
--xla_gpu_enable_collectives_command_buffer_filter=ALLREDUCE,ALLGATHER" \
python train.py

The named operation types are ALLREDUCE, ALLGATHER, REDUCESCATTER,
COLLECTIVEBROADCAST, ALLTOALL, COLLECTIVEPERMUTE, and RAGGEDALLTOALL.
ALLCOLLECTIVES allows all supported types and is the default. An empty filter
also means all types, so it is not a way to disable collective capture.

The filter does not enable COLLECTIVES by itself. It applies to mapped
collective operation types within that category; it is not a universal filter
for groups, send/receive, or COLLECTIVES_KERNEL. Use a plain list such as the
example above to narrow the filter: subtracting a type while retaining
ALLCOLLECTIVES does not exclude it.

To exclude both library collectives and collective kernels from capture:

XLA_FLAGS="--xla_gpu_enable_command_buffer=-COLLECTIVES,-COLLECTIVES_KERNEL" \
python train.py

Adjust the minimum graph size Copy Chapter

--xla_gpu_graph_min_graph_size=5 controls how large an eligible region must be
before XLA captures it. The current implementation counts leaf runtime
operations, including those inside nested structures. This is not a count of
Python operators, tensor elements, or necessarily the final GPU graph nodes.

Lowering the threshold can capture smaller regions:

XLA_FLAGS="--xla_gpu_graph_min_graph_size=1" python train.py

Try this when launch overhead matters and eligible regions are small. Raising
the threshold avoids building graphs for small regions whose capture and update
costs outweigh replay savings. A threshold of 1 still does not make
unsupported operations capturable.

Choose a scheduling mode Copy Chapter

--xla_gpu_command_buffer_scheduling_mode controls concurrency inside command
buffers. Start with LHS, then compare one alternative at a time.

XLA_FLAGS="--xla_gpu_command_buffer_scheduling_mode=CONCURRENT_REGIONS" \
python train.py

More concurrency is not always faster. Kernels can compete for compute
resources, memory bandwidth, or instruction-cache capacity. Compare both elapsed
time and peak memory usage.

Reduce graph update overhead Copy Chapter

A graph records the device addresses used by its operations. When addresses
change, XLA may need to update graph nodes or capture library calls again.
--xla_gpu_command_buffer_update_mode controls the allocation strategy used to
reduce this work.

Virtual memory management (VMM) lets XLA present a fixed device virtual address
to the graph even when the underlying allocation changes. Remapping also has a
cost; it is most useful when it saves more graph-update work than it adds.

XLA_FLAGS="--xla_gpu_command_buffer_update_mode=SKIP_TEMP" python train.py

# Run separately to compare profile-guided selection.
XLA_FLAGS="--xla_gpu_command_buffer_update_mode=SKIP_PROFILED" python train.py

Both skip modes require a compatible VMM allocator. The upstream StreamExecutor
GPU PJRT client automatically selects its VMM allocator when either is
requested, overriding the requested allocator kind. Other integrations must
supply compatible allocator support; the flag alone does not guarantee that
remapping is active. Compare memory behavior as well as latency because the
allocator changes.

In the current implementation, SKIP_PROFILED observes the first three
successful executions per executable/device state and uses ordinary
operation-by-operation execution during that observation window. Subsequent
executions can use the selected remappings and command buffers. Warm up beyond
both the observation window and initial graph setup before timing steady-state
execution.

This internal address profiling is separate from an external performance
profiler. It does not require
xla_enable_command_buffers_during_profiling=true. A later address change is
handled through remapping rather than assuming the original address is still
valid. For distributed runs, use consistent settings and representative,
symmetric warmup behavior across ranks.

Neither skip mode promises to eliminate all graph updates: other dynamic buffers
and changing operation parameters can still require them.

Tune the trace cache Copy Chapter

--xla_cmd_buffer_trace_cache_size=16 sets the capacity of each traced-command
cache. Library calls can be captured by observing the GPU work they enqueue;
the cache reuses those captures for previously seen buffer-address combinations.
This is not the framework's compilation cache or a global limit on graph count.

XLA_FLAGS="--xla_cmd_buffer_trace_cache_size=32" python train.py

A larger cache may help when a workload repeatedly cycles through more address
combinations than the cache holds. It retains more graph resources and increases
cache lookup work. It may not help if addresses rarely repeat. Use a positive
integer; 0 is not a supported way to disable this cache.

Capture and unroll loops Copy Chapter

--xla_gpu_command_buffer_unroll_loops=true allows eligible loops with a known
trip count to be expanded into repeated graph commands. The default is false.
Enable WHILE capture as well when testing whole-loop capture:

XLA_FLAGS="--xla_gpu_enable_command_buffer=+WHILE \
--xla_gpu_command_buffer_unroll_loops=true" python train.py

Unrolling can reduce loop-control overhead and enables some cases with
loop-dependent dynamic-slice offsets. It can also produce a much larger graph,
increasing setup time and memory use. Start with short, fixed-trip-count loops;
unknown trip counts or unsupported body operations prevent this transformation.

Profile the execution path you intend to measure Copy Chapter

--xla_enable_command_buffers_during_profiling defaults to false. During an
XLA-detected active profiling session, XLA normally falls back to individual
operations. A profiled run can therefore behave differently from an unprofiled
run even when capture is enabled.

To keep graph execution enabled during such a session:

XLA_FLAGS="--xla_enable_command_buffers_during_profiling=true" python train.py

This option is marked experimental. Use it when investigating graph replay, and
compare against unprofiled wall-clock measurements. It does not enable
additional command categories or bypass capture eligibility checks. External
profilers that do not activate XLA's profiling-session detection may behave
differently.

Run a useful comparison Copy Chapter

1. Measure the defaults and a separate run with command buffers disabled.
2. Warm up each compiled workload before measuring. Exclude compilation, initial
   capture, and SKIP_PROFILED observation/setup from steady-state timings, but
   measure startup separately if it matters to your application.
3. Wait for GPU work to finish before stopping the timer. For example, with JAX,
   call jax.block_until_ready(result); timing asynchronous dispatch alone does
   not measure execution time.
4. Change one setting at a time. Keep shapes, inputs, device count, and profiler
   configuration comparable. Apply the same configuration on every worker.
5. Compare output correctness, steady-state latency or throughput, startup time,
   and peak memory. Keep a setting only when it improves your target workload.

Troubleshooting and version compatibility Copy Chapter

Backend support differs. CUDA command buffer capture is disabled if either the
known runtime or driver version is below 12.3. DynamicSliceFusionV2 capture
requires runtime, driver, and build-time CUDA toolkit versions of at least 12.9.
ROCm disables conditional/while graph capture, and oneAPI command buffer
conversion is disabled. Enabling a category does not override these checks.

For implementation details, see the
flag definitions,
option enums,
capture eligibility checks,
and VMM allocation policy.

Contributing (`docs/contributing.md`) Copy Section

Contributing to OpenXLA

Everyone can contribute to OpenXLA, and we value everyone’s contributions. There
are several ways to contribute, including:

• Answering questions on OpenXLA’s discussions forums (openxla-discuss)

• Improving or expanding OpenXLA’s documentation

• Contributing to OpenXLA’s code-base

• Contributing in any of the above ways to the broader ecosystem of libraries
  built on OpenXLA

The OpenXLA project follows
Google’s Open Source Community Guidelines.

Before you begin Copy Chapter

Sign the Contributor License Agreement Copy Section

Contributions to this project must be accompanied by a
Contributor License Agreement (CLA).
You (or your employer) retain the copyright to your contribution; this simply
gives us permission to use and redistribute your contributions as part of the
project.

If you or your current employer have already signed the Google CLA (even if it
was for a different project), you probably don't need to do it again.

Visit https://cla.developers.google.com/ to see your current agreements or to
sign a new one.

Review the Code of Conduct Copy Section

This project follows
Tensorflow's Code of Conduct.

Contribution process Copy Chapter

Developer Guide Copy Section

For a guide on how to setup a development environment for OpenXLA, including
getting code, building it, running tests and submitting changes, please refer to
the Developer guide.

Contribution guide Copy Section

The architecture of the compiler consists of the following components.

Optimization passes

Optimization passes execute transformations on the HLO to enhance computational
efficiency. These transformations span from architecture-agnostic, high-level
improvements to hardware-specific adjustments (e.g., for NVIDIA GPUs).

What we generally accept:

• Passes that generalize across multiple workloads and demonstrate a clear and
  significant positive impact on performance benchmarks.

What we generally reject:

• Passes that perform unique optimizations targeting specific models.

Fusion passes

Fusion is a critical optimization that combines multiple HLO operations into a
single kernel to reduce memory I/O and kernel launch overhead.

All fusion passes should be added to the fusion pipeline only, not before or
after. That also means that the pre-optimized HLO module should not contain
fusion instructions. If the fusion is formed early in the pipeline, it becomes a
barrier for the optimization passes. If the fusion is formed late, then we lose
an ability to select and tune the backend for the generated fusion.

Fusion into the custom calls, i.e. pattern-matching custom calls with the
producers/consumers and rewriting them into the new custom calls is not allowed.
In that case, it should be replaced with a proper fusion pass.

Backends & Autotuning

Backends for the unnested ops, e.g. custom calls and fusions, should implement
CodegenBackend
interface.

This interface is necessary to enable optimal backend selection, because it
provides the methods to include the parameters for the given HLO instructions
into the search space of the autotuner.

// Returns all supported configs for the given HLO instruction.
virtual absl::StatusOr<std::vector<std::unique_ptr<BackendConfig>>>
  GetSupportedConfigs(const HloInstruction& instr);

// Returns a default config for the given HLO instruction.
virtual absl::StatusOr<std::unique_ptr<BackendConfig>> GetDefaultConfig(
  const HloInstruction& instr);

Runtime

The end result of the XLA compilation pipeline is a thunk sequence that can be
serialized.

All of the new thunk types should be serializable, i.e. GpuCompiler or
CpuCompiler should be able to compile the program, serialize it, so that later
the XLA runner could load and execute the program. That means that there should
be no pointers to HloInstruction or to other parts of the compiler or the
StreamExecutor.

GPU Compatibility Window {#gpu-compatibility-window}

XLA:GPU guarantees 6 months of backward compatibility (runtimes must execute
programs compiled up to 6 months earlier) and 2 weeks of forward compatibility
(runtimes up to 2 weeks older must execute newly compiled programs).

For detailed rules, change triage, rollout procedures, and testing guidance,
see the GPU AOT compatibility guide.

Code standards Copy Section

• Coding style: We follow Google's code style guide.
  Specifically see the C/C++ and Python guides. All
  code submitted must strictly conform to these style guides.

• Compact changes: We follow
  Google's engineering practices.
  In particular, please observe the
  guide on writing compact changes.
  Doing so will greatly increase the speed at which you can get your code
  merged due to improve reviewability, and reducing the likelihood of
  unintentional side effects of change. Even if you have a large change, there
  are many strategies for breaking it up into more incremental changes. If
  your PR is too large, it will receive an automated comment asking you to
  break it down into smaller PRs.

• Test Coverage: All changes should include appropriate unit tests. Unit
  tests should not be dependent on specific hardware (CPU, GPU, etc.) timings,
  and should make liberal use of mocks and fakes in order to make
  deterministic and focused tests. Changes seeking to extend existing code
  that’s currently hard to test should make appropriate improvements to
  testability.

  All changes should include appropriate benchmark results as well in the
  change title to ensure the benefits are clearly understood.

• Feature Flags: All somewhat complicated new features should be guarded
  with a flag first (e.g., via DebugOptions). This allows for easy rollback
  of the flag flip if problems arise, and affected users can temporarily
  set the flag themselves before a rollback is performed.

• Preprocessor Conditionals (#ifdef): Avoid using #if / #ifdef for
  backend-, hardware-, or version-specific logic whenever possible. Prefer
  runtime checks via StreamExecutor's DeviceDescription or separate build
  targets. See The Hitchhiker's Guide to Fewer #ifdefs in XLA
  for details and examples.

• When in doubt as to conventions within the code, it is always a good idea to
  examine pre-existing code and to try to follow the patterns already in place
  in OpenXLA.

Review Process Copy Section

All submissions, including submissions by project members, require review. We
use GitHub pull requests for this purpose. Consult
GitHub Help for more
information on using pull requests.

• Code must follow all standards listed above prior to review. These are not
  optional and it is critical that the submitter ensure their code conforms
  before requesting review in order to assure timely acceptance of changes.

• All tests and additional checks on GitHub must pass. If you find that a
  test is broken and the issue is not related to your build environment or
  otherwise your changes, please contact the maintainers.

• Avoid scope creep during the review process. This is the
  responsibility of both the submitter and the reviewer. If a change starts to
  get too large, consider breaking it up into multiple changes.

• After a change is approved on GitHub but before it is merged, it will
  undergo internal testing that uses code
  internal to Google and other hardware vendors. This can potentially add extra
  steps to the review process if there are failures on internal tests that our
  public CI doesn't catch. The Googler reviewing your change will communicate
  any internal test failures and describe what needs to be fixed. The overall
  state of the internal checks is visible in the checks list on GitHub:

  • import/copybara — Change imported to the internal review system:
    Your PR has been imported in Google's internal system and checks are
    running.
  • import/copybara — An error happened while migrating the change: Your PR
    could not be imported into Google's internal system. This very rarely
    happens. If you see this state please ping your reviewer.
  • feedback/copybara — Google internal checks PASS for runs with create time...:
    All internal checks pass. Your PR should be merged soon.
  • feedback/copybara — Google internal checks FAILED for runs with create time ...:
    Some internal checks failed. A Google engineer will soon post a comment with
    more details. If you don't get any info about the failures within a day
    please ping your reviewer.

AI-generated contributions Copy Section

AI-generated PRs can be accepted in many cases. However, we experience a
large review load due to AI-generated code and we do not want to waste valuable
time reviewing PRs that are not beneficial to the broader community. To
streamline code reviews, each AI-generated PR should be:

• Labeled - The first line after the title in the PR description should
  be this tag: #ai-generated

• Relevant - Just because a PR can be generated does not mean it should be
  submitted. We do not want to merge PRs that only defend against theoretical
  edge cases without providing any real value. Such "fixes" only clutter the
  code-base. You can demonstrate the relevance of a PR by providing the
  following:

  • A convincing explanation of where you encountered the problem in
    practice.
  • A convincing HLO or especially JAX test.
  • Realistic benchmarks showing improvements.

• Understood by its author - PR authors must understand all the code in
  their PRs in detail regardless of whether it was generated by AI or not.
  Moreover, if we ask questions about the PR we expect a human-written answer.

• Concise - The PR description should be short and to the point. AI often
  generates long-winded text with superfluous detail and this needs to be
  edited for brevity. The same applies to code and comments (e.g. don't
  check for error conditions that can't occur).

WARNING: Expect longer review times for AI-generated PRs.

If we suspect that a PR is AI-generated and it does not have the qualities
above we may reject it or ask you to revise it with this canned
response:

text Copy

This looks like an AI-generated PR that does not fully follow the guidance at:
https://openxla.org/xla/contributing#ai-generated-contributions

Please revise the PR according to the guidance above or it will be rejected.

If we receive an AI-generated answer to one of our comments or
questions on the PR, the PR will be immediately rejected without further
discussion.

Frequently asked questions (FAQ) Copy Chapter

"This infrastructure change is not related to my PR. Why should I do it?" Copy Section

The XLA team doesn't have a dedicated infrastructure team, so it's up to us all
to build helper libraries and avoid technical debt. We consider it to be a
regular part of making changes to XLA, and everyone is expected to participate.
We generally build infrastructure as needed when writing code.

XLA reviewers may ask you to build some infrastructure (or otherwise make a
large change to a PR) along with a PR that you've written. This request may seem
unnecessary or orthogonal to the change you're trying to make. This is likely
because of a mismatch between your expectations about how much infra you need to
build and your reviewer's expectations for the same.

A mismatch in expectations is okay! That's expected when you're new to a project
(and it sometimes even happens to us old hats). It's likely that projects you've
worked on in the past have different expectations. That's also okay and
expected! It doesn't mean either one of these projects has the wrong approach;
they're just different. We invite you to take infra requests alongside all other
review comments as an opportunity to learn what we expect on this project.

"Can I address your comment in a future PR?" Copy Section

A frequent question with respect to infrastructure requests (or other large
requests) in PRs is whether or not the change must be made in the original PR,
or whether it can be done as a follow-up in a future PR.

In general, XLA does not allow PR authors to address review comments with a
follow-up PR. When a reviewer decides that something needs to be addressed in a
given PR, we generally expect authors to address it in that PR, even if what's
requested is a large change. This standard applies externally and also
internally within Google.

There are a few reasons that XLA takes this approach.

• Trust: Having earned the reviewer's trust is a key component. In an
  open-source project, contributors can appear or disappear at will. After we
  approve a PR, reviewers have no way to ensure that any promised follow-ups
  actually get done.

• Impact on other developers: If you have sent a PR touching a particular
  part of XLA, there's a good chance other people are looking at the same
  part. If we accept technical debt in your PR, then everyone who's looking at
  this file will be impacted by this debt until the follow-up is submitted.

• Reviewer bandwidth: Deferring a change to a follow-up imposes multiple
  costs on our already overloaded reviewers. Reviewers will probably forget
  what the first PR was about while waiting for the follow-up, making the next
  review more difficult. Also, reviewers will have to keep track of expected
  follow-ups, making sure that they actually happen. If the change can be made
  such that it is truly orthogonal to the original PR so that some other
  reviewer could review it, bandwidth would be less of a problem. In our
  experience, this is rarely the case.

What is the turn around time for code review and merge? Copy Section

The OpenXLA project has a large contributor base. While we strive for quick
reviews and code merges, there are delays due to peaks of work.

In order to empower our partner teams to contribute high quality features and
fixes to the codebase, and to get quicker reviews and merges, we have
established the co-maintainer program. In this program, selected trusted partner
contributors ensure that the PRs submitted by those partners meet OpenXLA
quality requirements as specified in this contributing guide.

If co-maintainers from external partner teams have approved a PR, the Google XLA
team commits to reviewing and merging (if approved) within a controlled SLO
(t_{Google Review} in the diagram below). If the PR is not approved,
the Google XLA team will provide a documented justification and where available
a reproducer of the failure.

At this moment the list of co-maintainers from our partner teams is:

• @Tixxx
• @sergachev
• @jreiffers
• @ezhulenev
• @mminutoli
• @pemeliya

When Google reviews the PR, it may fail due to an internal test. In this case,
Google will try to generate a reproducer and share it with the partner. In other
cases, the review may fail for technical, architectural or stylish reasons -
Google will share the feedback with the author. We expect that if the PR is
rejected in the first pass, we will enter a tight loop iteration between the
partner team and Google (“Iterative review”) to bring the PR to the desired
state.

The OpenXLA project is committed to strengthening its co-maintainer program and
adding more contributors to this list. We plan to review on a monthly basis the
target turnaround times and report back to the participating teams.

Copybara (`docs/copybara.md`) Copy Section

Copybara quirks

The purpose of this document is to describe oddities you might see while
contributing due to the tool that manages copying source back and forth from
Google's internal repository. This tool is called Copybara.

Internal source of truth Copy Chapter

Because the source of truth for the code in this repository is Google's internal
repo, Copybara does transformations to the code whenever the code is imported
and exported. This means that sometimes seemingly normal changes can break
internally in surprising ways.

PR merge status and diff inconsistencies Copy Chapter

Since the source of truth is internal, PRs are not merged directly, they are
imported to the Google internal repo where they undergo additional testing,
and then that internal change is submitted, and attributed to the PR author.
Because of the transformations that Copybara applies, there's no guarantee that
the diff will be identical (for example, Copybara applies formatting on import).

For this reason, Copybara won't mark the PR as merged, it will close the PR and
separately apply a commit that should map very closely to the PR.

Dependency on TSL by copy Copy Chapter

As implemented currently, to prevent any temporary broken commits, XLA
depends on TSL not by downloading a copy by using Bazel's http_archive, but by
having Copybara copy TSL into XLA's third_party directory.

Custom Call (`docs/custom_call.md`) Copy Section

XLA Custom Calls

This document describes how to write and use XLA custom calls using XLA FFI
library. Custom call is a mechanism to describe an external "operation" in the
HLO module to the XLA compiler (at compile time), and XLA FFI is a mechanism to
register implementation of such operations with XLA (at run time). FFI stands
for "foreign function interface" and it is a set of C APIs that define a binary
interface (ABI) for XLA to call into external code written in other programming
languages. XLA provides header-only bindings for XLA FFI written in C++, which
hides all the low level details of underlying C APIs from the end user.

Caution: The custom-call API/ABI uses PJRT-style versioning (major, minor),
however at this point it is still experimental and can be broken at any time.
Once API/ABI is finalized we intend to provide stability guarantees
similar to PJRT.

Caution The HLO-visible names of functions registered with the custom-call
macros API do not respect C++ namespaces. As a result, accidental collisions
from functions registered by different libraries are entirely possible! The
API will reject such duplicate registrations, but to avoid issues in large
projects the safest option is to either fully namespace-qualify all references
to the functions in both the XLA_REGISTER_CUSTOM_CALL registration macros
and custom call target references or to use C-style namespacing directly in
the function name.

JAX + XLA Custom Calls Copy Chapter

See JAX documentation for
end to end examples of integrating custom calls and XLA FFI with JAX.

XLA FFI Binding Copy Chapter

XLA FFI binding is a compile-time specification of the custom call signature:
custom call arguments, attributes and their types, and additional parameters
passed via the execution context (i.e., gpu stream for GPU backend). XLA FFI
binding can be bound to any C++ callable (function pointer, lambda, etc.) with
compatible operator() signature. Constructed handler decodes XLA FFI call
frame (defined by the stable C API), type check all parameters, and forward
decoded results to the user-defined callback.

XLA FFI binding heavily relies on template metaprogramming to be able to
compile constructed handler to the most efficient machine code. Run time
overheads are in order of a couple of nanoseconds for each custom call
parameter.

XLA FFI customization points implemented as template specializations, and
users can define how to decode their custom types, i.e., it is possible
to define custom decoding for user-defined enum class types.

Returning Errors From Custom Calls Copy Section

Custom call implementations must return xla::ffi::Error value to signal
success or error to XLA runtime. It is similar to absl::Status, and has
the same set of error codes. We do not use absl::Status because it does
not have a stable ABI and it would be unsafe to pass it between dynamically
loaded custom call library, and XLA itself.

// Handler that always returns an error.
auto always_error = Ffi::Bind().To(
    []() { return Error(ErrorCode::kInternal, "Oops!"); });

// Handler that always returns a success.
auto always_success = Ffi::Bind().To(
    []() { return Error::Success(); });

Buffer Arguments And Results Copy Section

XLA uses destination passing style for results: custom calls (or any other XLA
operations for that matter) do not allocate memory for results, and instead
write into destinations passed by XLA runtime. XLA uses static buffer
assignment, and allocates buffers for all values based on their live ranges at
compile time.

Results passed to FFI handlers wrapped into a Result<T> template, that
has a pointer-like semantics: operator-> gives access to the underlying
parameter.

AnyBuffer arguments and results gives access to custom call buffer parameters
of any data type. This is useful when custom call has a generic implementation
that works for multiple data types, and custom call implementation does run time
dispatching based on data type. AnyBuffer gives access to the buffer data
type, dimensions, and a pointer to the buffer itself.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  api_version = 4 : i32
} : (tensor<2x2xf32>) -> tensor<2x2xf32>

// Buffers of any number of dimensions and data type.
auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer arg, Result<AnyBuffer> res) -> Error {
      void* arg_data = arg.untyped_data();
      void* res_data = res->untyped_data();
      return Error::Success();
    });

Constrained Buffer Arguments And Results Copy Section

Buffer allows to add constraints on the buffer data type and number of
dimensions, and they will be automatically checked by the handler and return
an error to XLA runtime, if run time arguments do not match the FFI handler
signature.

// Buffers of any number of dimensions and F32 data type.
auto handler = Ffi::Bind().Arg<Buffer<F32>>().Ret<Buffer<F32>>().To(
    [](Buffer<F32> arg, Result<Buffer<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

// Buffers of number of dimensions 2 and F32 data type.
auto handler = Ffi::Bind().Arg<BufferR2<F32>>().Ret<BufferR2<F32>>().To(
    [](BufferR2<F32> arg, Result<BufferR2<F32>> res) -> Error {
      float* arg_data = arg.typed_data();
      float* res_data = res->typed_data();
      return Error::Success();
    });

Matching And Verifying Buffers Copy Section

Buffer patterns can express constraints that are more specific than the argument
and result types in an FFI binding. They are useful for refining an AnyBuffer
to a concrete buffer type, checking dimension sizes, and checking relationships
between multiple buffer shapes.

Patterns are immutable and can specify dtype and rank either directly or with
the corresponding modifiers:

namespace m = ::xla::ffi::match;

m::Buffer<F32, 2>();
m::Buffer().WithDType<F32>().WithRank<2>();

The Match function checks a pattern and returns its maximally constrained
buffer type. To refine an AnyBuffer, the pattern must specify exactly one data
type and one rank. Matching an already typed buffer verifies any additional
constraints and preserves its type.

namespace m = ::xla::ffi::match;

auto handler = Ffi::Bind().Arg<AnyBuffer>().Ret<AnyBuffer>().To(
    [](AnyBuffer input, Result<AnyBuffer> output) -> Error {
      int64_t rows;
      int64_t cols;

      ErrorOr<BufferR2<F32>> typed_input =
          Match("input", input,
                m::Buffer<F32>().WithDims(&rows, &cols));
      if (typed_input.has_error()) {
        return std::move(typed_input).error();
      }

      ErrorOr<Result<BufferR2<F32>>> typed_output =
          Match("output", output,
                m::Buffer<F32>().WithDims(rows, cols));
      if (typed_output.has_error()) {
        return std::move(typed_output).error();
      }

      float* input_data = typed_input->typed_data();
      float* output_data = (*typed_output)->typed_data();
      return Error::Success();
    });

In this example matching input captures its two dimension sizes, and matching
output verifies that it has the same shape. Captures are committed only when
the complete pattern matches successfully.

A dimension pattern can be unconstrained, fixed, or captured. Fixed values and
capture pointers are implicitly converted to dimension patterns:

m::Dim()       // Any dimension size.
m::Dim(16)     // A dimension whose size is 16.
m::Dim(&rows)  // Any dimension size, captured in `rows` on success.

m::Buffer<F32>().WithDims(&rows, 16);

WithDims describes a complete shape and sets its rank from the number of
arguments. WithDim or WithDim(index, ...) constrains dimensions
positionally without fixing the complete rank. A positional constraint requires
the referenced dimension to exist but leaves every other dimension
unconstrained:

ErrorOr<BufferR4<F32>> MatchActivations(AnyBuffer input) {
  return Match(
      "input", input,
      m::Buffer<F32, 4>().WithDim<0>(1).WithDim<3>(128));
}

Error VerifyDimension(AnyBuffer input, size_t index, int64_t size) {
  return Verify("input", input, m::Buffer().WithDim(index, size));
}

MatchActivations requires a rank-4 F32 buffer with dimensions 0 and 3 equal to
1 and 128, respectively; dimensions 1 and 2 are unrestricted. Without an
explicit WithRank, WithDim(index, ...) accepts any rank greater than
index.

Dimension captures can also express relationships between positions:

auto matrix = m::Buffer<F32>().WithDims(&rows, 128);

int64_t n;
ErrorOr<BufferR2<F32>> square =
    Match("matrix", buffer,
          m::Buffer<F32>().WithDims(&n, &n));

When the same capture pointer appears multiple times, all corresponding
dimensions must have the same size. The Match call above therefore accepts
only square matrices. Dimension captures are written only after the complete
pattern succeeds and are unchanged on failure.

Whole-buffer relationships do not require capturing every dimension.
WithShapeOf matches the complete runtime shape of another buffer while
allowing a different dtype. Like also requires the same dtype:

Error VerifySortBuffers(AnyBuffer keys, AnyBuffer values,
                        Result<AnyBuffer> keys_output) {
  Error error =
      Verify("values", values, m::Buffer().WithShapeOf(keys));
  if (error.failure()) {
    return error;
  }

  return Verify("keys output", *keys_output, m::Buffer().Like(keys));
}

Both modifiers copy the reference buffer's metadata when the pattern is built;
the pattern does not retain a reference to the buffer. Because these are runtime
constraints, they do not refine an AnyBuffer to a concrete return type. They
are primarily useful with Verify, or with Match when the input buffer
already has a concrete type.

Use Verify when the buffer should keep its existing type, or when a pattern
accepts more than one data type or rank. For example, an index buffer can accept
either S32 or S64 and dispatch on its data type after verification:

namespace m = ::xla::ffi::match;

Error VerifyIndices(AnyBuffer indices) {
  auto pattern = m::Buffer().WithDType<S32, S64>().WithRank<1, 2>();

  Error error = Verify("indices", indices, pattern);
  if (error.failure()) {
    return error;
  }

  // Dispatch an implementation based on `indices.element_type()`.
  return Error::Success();
}

Because this pattern permits multiple concrete buffer types, it cannot refine an
AnyBuffer with Match: there is no unique return type. It can still be passed
to Match with an already typed buffer, whose return type is already known. An
already typed buffer can also be verified; this is typically useful for checking
shape relationships:

Error VerifySameShape(BufferR2<F32> lhs, BufferR2<F32> rhs) {
  int64_t rows;
  int64_t cols;

  Error error = Verify(
      "lhs", lhs,
      m::Buffer().WithDims(&rows, &cols));
  if (error.failure()) {
    return error;
  }

  return Verify("rhs", rhs,
                m::Buffer().WithDims(rows, cols));
}

The name passed as the first argument to Match and Verify is required and is
included in errors, together with a description of the failed constraint.

The external FFI API in xla/ffi/api/ffi.h returns Error from Verify and
ErrorOr<T> from Match, as in the examples above. The internal FFI API in
xla/ffi/ffi.h provides the same matching interface using absl::Status and
absl::StatusOr<T>:

ABSL_ASSIGN_OR_RETURN(BufferR2<F32> input,
                      Match("input", buffer, m::Buffer<F32, 2>()));
ABSL_RETURN_IF_ERROR(Verify(
    "input", input,
    m::Buffer().WithDims(expected_rows, m::Dim())));

Variadic Arguments And Results Copy Section

If the number of arguments and result can be different in different instances of
a custom call, they can be decoded at run time using RemainingArgs and
RemainingRets.

auto handler = Ffi::Bind().RemainingArgs().RemainingRets().To(
    [](RemainingArgs args, RemainingRets results) -> Error {
      ErrorOr<AnyBuffer> arg = args.get<AnyBuffer>(0);
      ErrorOr<Result<AnyBuffer>> res = results.get<AnyBuffer>(0);

      if (!arg.has_value()) {
        return Error(ErrorCode::kInternal, arg.error());
      }

      if (!res.has_value()) {
        return Error(ErrorCode::kInternal, res.error());
      }

      return Error::Success();
    });

Variadic arguments and results can be declared after regular arguments and
results, however binding regular arguments and results after variadic one is
illegal.

auto handler =
    Ffi::Bind()
        .Arg<AnyBuffer>()
        .RemainingArgs()
        .Ret<AnyBuffer>()
        .RemainingRets()
        .To([](AnyBuffer arg, RemainingArgs args, AnyBuffer ret,
               RemainingRets results) -> Error { return Error::Success(); });

Attributes Copy Section

XLA FFI supports automatic decoding of mlir::DictionaryAttr passed as a
custom_call backend_config into FFI handler arguments.

Note: See stablehlo RFC
for details, and stablehlo.custom_call operation specification.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    i32 = 42 : i32,
    str = "string"
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

In this example custom call has a single buffer argument and two attributes, and
XLA FFI can automatically decode them and pass to the user-defined callable.

auto handler = Ffi::Bind()
  .Arg<BufferR0<F32>>()
  .Attr<int32_t>("i32")
  .Attr<std::string_view>("str")
  .To([](BufferR0<F32> buffer, int32_t i32, std::string_view str) {
    return Error::Success();
  });

User-Defined Enum Attributes Copy Section

XLA FFI can automatically decode integral MLIR attributes into user-defined
enums. Enum class must have the same underlying integral type, and decoding
has to be explicitly registered with XLA FFI.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    command = 0 : i32
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

enum class Command : int32_t {
  kAdd = 0,
  kMul = 1,
};

XLA_FFI_REGISTER_ENUM_ATTR_DECODING(Command);

auto handler = Ffi::Bind().Attr<Command>("command").To(
    [](Command command) -> Error { return Error::Success(); });

Binding All Custom Call Attributes Copy Section

It is possible to get access to all custom call attributes as a dictionary
and lazily decode only the attributes that are needed at run time.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<int32_t> i32 = attrs.get<int32_t>("i32");
  return Error::Success();
});

Per-execution Custom Options Copy Section

FFI handlers can access runtime custom options as a Dictionary by binding
Ctx<CustomOptions>(). In JAX, pass options with jax.execution_options:

with jax.execution_options(custom_options={"foo": 42}):
    result = compiled_fn(x)

Python integer options decode as int64_t:

auto handler =
    Ffi::Bind().Ctx<CustomOptions>().To([](Dictionary options) -> Error {
      ErrorOr<int64_t> foo = options.get<int64_t>("foo");
      return Error::Success();
    });

WARNING: GPU command buffers can replay recorded or captured work without
invoking the FFI handler again. Changing custom options does not trigger a
command buffer update or invalidate the stream-capture cache, so replay may use
work recorded with earlier option values. Custom calls whose recorded work
depends on options that change between executions must stay out of GPU command
buffers until this is supported.

User-defined Struct Attributes Copy Section

XLA FFI can decode dictionary attributes into user-defined structs.

%0 = "stablehlo.custom_call"(%arg0) {
  call_target_name = "foo",
  backend_config= {
    range = { lo = 0 : i64, hi = 42 : i64 }
  },
  api_version = 4 : i32
} : (tensor<f32>) -> tensor<f32>

In example above range is an mlir::DictionaryAttr attribute, and instead
of accessing dictionary fields by name, it can be automatically decoded as
a C++ struct. Decoding has to be explicitly registered with a
XLA_FFI_REGISTER_STRUCT_ATTR_DECODING macro (behind the scene it defines
a template specialization in ::xla::ffi namespace, thus macro must be added to
the global namespace).

struct Range {
  int64_t lo;
  int64_t hi;
};

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(Range, StructMember<int64_t>("lo"),
                                             StructMember<int64_t>("hi"));

auto handler = Ffi::Bind().Attr<Range>("range").To([](Range range) -> Error{
  return Error::Success();
});

Custom attributes can be loaded from a dictionary, just like any other
attribute. In example below, all custom call attributes decoded as a
Dictionary, and a range can be accessed by name.

auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error {
  ErrorOr<Range> range = attrs.get<Range>("range");
  return Error::Success();
});

Create a custom call on CPU Copy Chapter

You can create an HLO instruction that represents a custom call via XLA's client
API. For example, the following code uses a custom call to compute A[i] = B[i % 128]+ C[i] on the CPU. (Of course you could – and should! – do this
with regular HLO.)

/* Detailed source-code truncated for AI context efficiency. */

Create a custom call on GPU Copy Chapter

The GPU custom call registration with XLA FFI is almost identical, the only
difference is that for GPU you need to ask for an underlying platform stream
(CUDA or ROCM stream) to be able to launch kernel on device. Here is a CUDA
example that does the same computation (A[i] = B[i % 128] + C[i]) as the CPU
code above.

void do_it() { /* same implementation as above */ }

__global__ custom_call_kernel(const float* in0, const float* in1, float* out) {
  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
  out[idx] = in0[idx % 128] + in1[idx];
}

void do_custom_call(CUstream stream, BufferF32 in0, BufferF32 in1,
                    xla::ffi::Result<BufferF32> out) {
  size_t d0 = in0.dimensions[0];
  size_t d1 = in1.dimensions[0];
  size_t d2 = out->dimensions[0];

  assert(d0 == 128 && d1 == 2048 && d2 == 2048 && "unexpected dimensions");

  const int64_t block_dim = 64;
  const int64_t grid_dim = 2048 / block_dim;
  custom_call_kernel<<<grid_dim, block_dim, 0, stream>>>(
    in0.data, in1.data, out->data);
}

XLA_FFI_DEFINE_HANDLER(handler, do_custom_call,
                       ffi::Ffi::Bind()
                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
                           .Arg<BufferF32>()
                           .Arg<BufferF32>()
                           .Ret<BufferF32>());

XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "do_custom_call",
                         "CUDA", handler);

Notice first that the GPU custom call function is still a function executed on
the CPU. The do_custom_call CPU function is responsible for enqueueing work
on the GPU. Here it launches a CUDA kernel, but it could also do something else,
like call cuBLAS.

Arguments and results also live on the host, and data member contains a pointer
to device (i.e. GPU) memory. Buffers passed to custom call handler have the
shape of the underlying device buffers, so the custom call can compute kernel
launch parameters from them.

Passing tuples to custom calls Copy Chapter

Consider the following custom call.

using xla::ShapeUtil;
using xla::F32;
Shape p0_shape = ShapeUtil::MakeTuple({
    ShapeUtil::MakeShape(F32, {32}),
    ShapeUtil::MakeTuple({
        ShapeUtil::MakeShape(F32, {64}),
        ShapeUtil::MakeShape(F32, {128}),
    }),
    ShapeUtil::MakeShape(F32, {256}),
});
xla::XlaOp p0 = xla::Parameter(0, p0_shape, "p0");

Shape out_shape = ShapeUtil::MakeTuple({
  ShapeUtil::MakeShape(F32, {512}),
  ShapeUtil::MakeShape(F32, {1024}),
});
xla::CustomCall(&b, "do_custom_call", /*operands=*/{p0}, out_shape, ...);

On both CPU and GPU, a tuple is represented in memory as an array of pointers.
When XLA calls custom calls with tuple arguments or results it flattens them and
passes as regular buffer arguments or results.

Tuple outputs as temp buffers Copy Section

Tuple inputs to custom calls are a convenience, but they aren't strictly
necessary. If we didn't support tuple inputs to custom calls, you could always
unpack the tuples using get-tuple-element before passing them to the custom
call.

On the other hand, tuple outputs do let you do things you couldn't otherwise.

The obvious reason to have tuple outputs is that tuple outputs are how a custom
call (or any other XLA op) returns multiple independent arrays.

But less obviously, a tuple output is also a way to give your custom call temp
memory. Yes, an output can represent a temp buffer. Consider, an output buffer
has the property that the op can write to it, and it can read from it after it's
been written to. That's exactly what you want from a temp buffer.

In the example above, suppose we wanted to use the F32[1024] as a temp buffer.
Then we'd write the HLO just as above, and we'd simply never read tuple index 1
of the custom call's output.

Determinism (`docs/determinism.md`) Copy Section

Determinism (GPU)

Non-determinism on GPU in XLA stems from two distinct sources:

1. Compilation-Time Non-Determinism (Autotuning)
2. Execution-Time Non-Determinism (Non-Deterministic Ops)

Source 1: Compilation-Time Non-Determinism (Autotuning) Copy Chapter

Cause Copy Section

During compilation, XLA's autotuner profiles multiple candidate kernel
implementations (e.g., cuBLAS, cuDNN, Triton, native emitters) live on the
host's GPU to find the fastest algorithm for instructions like GEMMs and
convolutions.

Because live profiling relies on microsecond-level timing measurements, subtle
environmental hardware fluctuations (GPU thermal throttling, dynamic clock
frequency scaling, host OS scheduling) can cause different kernel binaries to
be selected/generated as the winner across separate compilation runs.

Because different kernel algorithms use distinct tile sizes, memory access
patterns, and reduction tree structures, they accumulate floating-point numbers
in different orders. Due to floating-point non-associativity, compiling the same
HLO graph twice can produce binaries that output slightly different numerical
calculation results.

How to Achieve Compilation-time Determinism Copy Section

• Persisting & Caching Autotuning Decisions: Run autotuning once and save
  the choices to a persistent cache (see
  persisted autotuning). Using the flag
  --xla_gpu_require_complete_aot_autotune_results enforces that XLA
  strictly reuses cached decisions and immediately fails compilation with a
  NotFound error if an entry is missing, completely avoiding live hardware
  benchmarking.
• Disabling Autotuning: Setting --xla_gpu_autotune_level=0 disables
  autotuning benchmarks at the compile time altogether and selects default /
  first valid configurations deterministically.

Source 2: Execution-Time Non-Determinism (Non-Deterministic Ops) Copy Chapter

Cause Copy Section

Even when the compiled binary is 100% identical across runs, running the exact
same executable multiple times can naturally yield different numerical outputs.

On GPUs, thousands of threads execute in parallel. Operations such as scatter,
accumulate values into memory buffers using atomic operations across
uncoordinated threads. Because thread scheduling order varies naturally from run
to run, floating-point additions occur in non-deterministic order. Because
floating-point addition is non-associative ((a + b) + c != a + (b + c)),
execution outputs vary between runs.

How to Achieve Determinism Copy Section

• --xla_gpu_exclude_nondeterministic_ops: Enforces execution determinism
  by instructing XLA to:
  1. Filter out non-deterministic algorithm/kernel/backend variants during
     config selection and enforce deterministic lowerings (for scatter, GEMM,
     convolution).
  2. Rewrite operations like select-and-scatter into a sequence of
     deterministic ops to eliminate non-deterministic atomic scatter
     operations.
  3. Disable live autotuning benchmarks (picking the first valid
     deterministic configuration).

Side Effects & Trade-Offs Copy Chapter

Enforcing determinism on GPU involves notable trade-offs:

• Runtime Performance Degradation (Slowdowns):
  • Deterministic implementations of operations like scatter (and expanded
    select-and-scatter) replace uncoordinated atomic updates with
    additional index sorting passes or deterministic reduction steps, which
    can lead to substantial runtime throughput loss compared to
    non-deterministic atomic fast paths.
  • Bypassing live autotuning (--xla_gpu_autotune_level=0) prevents XLA
    from discovering the fastest kernel for the specific GPU architecture,
    resulting in sub-optimal kernel choices.
• Compilation Hard-Failures:
  • If an HLO graph contains instructions without any deterministic
    implementation, enabling --xla_gpu_exclude_nondeterministic_ops causes
    compilation to halt with an error.
  • When using --xla_gpu_require_complete_aot_autotune_results,
    compilation will fail immediately with a NotFound error if any
    instruction lacks an entry in the pre-populated autotune cache.
• Disabled Optimizations & Emitters:
  • Certain high-performance kernel emitters and fusion transformations
    (such as Triton fusion autotuning) are automatically excluded from
    consideration when determinism is required.

Developer Environment (`docs/developer_environment.md`) Copy Section

Setting Up Developer Environment

Setting up LSP with clangd Copy Chapter

Background Copy Section

Editors such as Emacs, Vim, or VS Code support features like code navigation,
code completion, inline compiler error messages, and others, through
LSP, the Language
Server Protocol. A common language server with LSP support is
clangd, which relies on the presence of
compile_commands.json, a JSON file with a record of the compile commands for
each file in a project.

How do I generate `compile_commands.json` for XLA source code? Copy Section

Use the
build_tools/lint/generate_compile_commands.py
script. The following invocation from XLA repo root generates a
compile_commands.json file in place: bazel aquery "mnemonic(CppCompile, //xla/...)" --output=jsonproto | python3 build_tools/lint/generate_compile_commands.py

Build Cleaner Copy Chapter

XLA CI inside Google runs additional checks to verify that all targets correctly
list all dependencies in the BUILD files, which are not enabled by default in
OSS Bazel CI. Running following commands before sending PR to XLA team will
massively speedup the time it takes to merge the PR as otherwise you will need
some Googler to make this fixes for you internally, or a few more rounds of PR
reviews to fix them based on the Googlers feedback.

Layering Check Copy Section

Building with --features=layering_check makes sure that you don't accidentally
include a header via transitive dependency without listing it in your target
dependencies

Remove unused dependencies from BUILD files Copy Section

Install
Buildozer
tool:

sudo curl -fsSL -o /usr/bin/buildozer https://github.com/bazelbuild/buildtools/releases/download/6.0.0/buildozer-linux-amd64
sudo chmod 755 /usr/bin/buildozer

Install Bant tool:

# To some writable directory that does not require root access
bazel build -c opt //bant && install -D --strip bazel-bin/bant/bant ~/bin/bant

# For a system directory that requires root-access
sudo install -D --strip bazel-bin/bant/bant /usr/local/bin/bant

Use bant to generate buildozer commands to remove unused deps:

bant dwyu //xla/core/collectives:symmetric_memory

if you feel lucky, you can execute them directly:

. <(bant dwyu //xla/core/collectives:symmetric_memory)

Developer Guide (`docs/developer_guide.md`) Copy Section

XLA developer guide

This guide shows you how to get started developing the XLA project.

Before you begin, complete the following prerequisites:

1. Go to Contributing page and review the contribution
   process.
2. If you haven't already done so, sign the
   Contributor License Agreement.
3. Install or configure the following dependencies:
   • A GitHub account
   • Docker

Then follow the steps below to get the source code, set up an environment, build
the repository, and create a pull request.

Get the code Copy Chapter

1. Create a fork of the XLA repository.

2. Clone your fork of the repo, replacing {USER} with your GitHub username:

   sh Copy

   git clone https://github.com/{USER}/xla.git

3. Change into the xla directory: cd xla

4. Configure the remote upstream repo:

   sh Copy

   git remote add upstream https://github.com/openxla/xla.git

Set up an environment Copy Chapter

1. Install Bazel.

   To build XLA, you must have Bazel installed. The recommended way to install
   Bazel is using Bazelisk,
   which automatically downloads the correct Bazel version for XLA. If Bazelisk
   is unavailable, you can install Bazel
   manually.

2. Create and run the
   ml-build
   Docker container.

   To set up a Docker container for building XLA with support for both CPU and
   GPU, run the following command:

   sh Copy

   docker run -itd --rm \
     --name xla \
     -w /xla \
     -v $PWD:/xla \
     us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
     bash

   If building with GPU/CUDA support, add --gpus all to grant the container
   access to all available GPUs. This enables automatic detection of CUDA
   compute capabilities.

Build Copy Chapter

Configure for CPU:

docker exec xla ./configure.py --backend=CPU

Configure for GPU:

docker exec xla ./configure.py --backend=CUDA

CUDA compute capabilities will be detected automatically by running
nvidia-smi. If GPUs are not available during the build, you must specify
the compute capabilities manually. For example:

# Automatically detects compute capabilities (requires GPUs)
./configure.py --backend=CUDA

# Manually specify compute capabilities (for builds without GPUs)
./configure.py --backend=CUDA --cuda_compute_capabilities="9.0"

Build:

docker exec xla bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Note: You can build XLA on a machine without GPUs. In that case:

• Do not use --gpus all flag when starting the Docker container.
• During ./configure.py, manually specify the CUDA compute capabilities
  using the --cuda_compute_capabilities flag.

Note: Thanks to hermetic CUDA rules, you don't need to build XLA inside a
Docker container. You can build XLA for GPU directly on your machine - even if
it doesn't have a GPU or the NVIDIA driver installed.

# Automatically detects compute capabilities (requires GPUs)
./configure.py --backend=CUDA

# Manually specify compute capabilities (for builds without GPUs)
./configure.py --backend=CUDA --cuda_compute_capabilities="9.0"

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Your first build will take quite a while because it has to build the entire
stack, including XLA, MLIR, and StableHLO.

To learn more about building XLA, see Build from source.

Create a pull request Copy Chapter

When you're ready to send changes for review, create a
pull request.

To learn about the XLA code review philosophy, see
Review Process.

Static Analysis (Clang-Tidy) Copy Chapter

To maintain code quality, XLA uses clang-tidy for static analysis and include
verification.

How to Run Copy Section

There are two ways to execute checks. Running it against specific targets can be
done with:

bazel build --config=clang-tidy //path/to:target1 //path/to:target2

There is a helper script that is also used in CI workflows that runs it against
git diff from feature branch against the upstream main.

# Make sure the main is updated.
git fetch origin main
bazel run //build_tools/ci:run_clang_tidy

The helper script also allows to automagically fix clang-tidy errors where
possible.

bazel run //build_tools/ci:run_clang_tidy -- --fix

Developing New Backend (`docs/developing_new_backend.md`) Copy Section

Developing a new backend for XLA

This guide is for system engineers who want XLA to output programs that target
their hardware efficiently. The guide is not step-by-step and assumes knowledge
of LLVM, Bazel, and XLA.

XLA provides an abstract interface that a new architecture or accelerator can
implement to create a backend to run ML programs output by XLA. Retargeting XLA
should be significantly simpler and more scalable than implementing every
existing op from a frontend framework such as PyTorch or TensorFlow for new
hardware.

Most implementations will fall into one of the following scenarios:

1. Existing CPU architecture not yet officially supported by XLA, with or
   without an existing LLVM backend.
2. Non-CPU-like hardware with an existing LLVM backend.
3. Non-CPU-like hardware without an existing LLVM backend.

Note: An LLVM backend can mean either one of the officially released LLVM
backends or a custom LLVM backend developed in-house.

Scenario 1: Existing CPU architecture not yet officially supported by XLA Copy Chapter

In this scenario, start by looking at the existing
XLA CPU backend. XLA
makes it easy to target different CPUs by using LLVM, because the main
difference between XLA backends for CPUs is the code generated by LLVM.

If the hardware vendor has an LLVM backend for their hardware, it is simple to
link the backend with the LLVM built with XLA. In JIT mode, the XLA CPU backend
emits code for the host CPU. For ahead-of-time compilation,
xla::AotCompilationOptions
can provide an LLVM triple to configure the target architecture.

If there is no existing LLVM backend, but another kind of code generator exists,
it should be possible to reuse most of the existing CPU backend.

Scenario 2: Non-CPU-like hardware with an existing LLVM backend Copy Chapter

It is possible to model a new
xla::Compiler
implementation on the existing
xla::CPUCompiler
and
xla::GPUCompiler
classes, since these already emit LLVM IR. Depending on the nature of the
hardware, it is possible that many aspects of the LLVM IR generation will have
to be changed, but a lot of code can be shared with the existing backends.

A good example to follow is the
GPU backend
of XLA. The GPU backend targets a non-CPU-like ISA, and therefore some aspects
of its code generation are unique to the GPU domain. Other kinds of hardware,
e.g. DSPs like Hexagon (which has an upstream LLVM backend), can reuse parts of
the LLVM IR emission logic, but other parts will be unique.

Scenario 3: Non-CPU-like hardware without an existing LLVM backend Copy Chapter

If it is not possible to utilize LLVM, then the best option is to implement a
new backend for XLA for the desired hardware. This option requires the most
effort. The classes that need to be implemented are as follows:

• StreamExecutor:
  For many devices not all methods of StreamExecutor are needed. See
  existing StreamExecutor implementations for details.
• xla::Compiler:
  This class encapsulates the compilation of an HLO computation into an
  xla::Executable.
• xla::Executable:
  This class is used to launch a compiled computation on the platform.
• xla::TransferManager:
  This class enables backends to provide platform-specific mechanisms for
  constructing XLA literal data from given device memory handles. In other
  words, it helps encapsulate the transfer of data from the host to the device
  and back.

Effort Levels (`docs/effort_levels.md`) Copy Section

Effort Levels

XLA provides options to control the amount of effort the compiler will expend to

• optimize for runtime performance, and
• make the program "fit in memory" (which has a platform-dependent meaning)

Optimization Level Copy Chapter

Similar to the -O flags in gcc or clang, this field allows the user to influence
how much work the compiler does in optimizing for execution time. It can be set
via the
optimization_level
field of the ExecutableBuildOptionsProto message, or the
optimization_level
field of the ExecutionOptions message.

Lower optimization levels will cause various HLO passes to behave differently,
typically doing less work, or may disable certain HLO passes entirely. The
optimization level may also influence the compiler backend, such that the exact
effect of this field has a dependence on the target platform. However, as a
general guideline, the following table describes the expected overall effect of
each value:

Use in XLA:GPU Copy Section

Pipelining can be configured separately for all-gather, all-reduce, and
reduce-scatter. The xla_gpu_pipeline_all_gather,
xla_gpu_pipeline_all_reduce, and xla_gpu_pipeline_reduce_scatter flags
accept default, off, on, or explicit. default follows optimization
effort: it selects on at O1 or above, or when execution-time optimization
effort is at least 0.2; otherwise, it selects off. Any other mode overrides
optimization effort. All-gather and all-reduce default to default, while
reduce-scatter defaults to on. See
GPU XLA flags for the mode semantics.

In XLA:GPU, there are several passes that we disable by default because they
significantly increase compilation time by increasing the HLO size. For
convenience, we consolidate them under the optimization level option, such that
setting optimization_level to O1 or above will lead to the following behavior:

• Collective pipelining uses on by default for operations commonly used for
  data-parallel communication. The off and explicit modes still override
  this default.
  • xla_gpu_pipeline_all_gather
  • xla_gpu_pipeline_all_reduce
  • xla_gpu_pipeline_reduce_scatter
• Unrolling while loops by a factor of two. Breaks down the loop-barrier
  potentially leading to a better compute-communication overlap and less
  copies.
  • xla_gpu_enable_while_loop_double_buffering
• Latency Hiding Scheduler will do most of the work to hide the communication
  latency.
  • xla_gpu_enable_latency_hiding_scheduler
• To maximize networking bandwidth, combiner passes will combine pipelined
  collectives to the maximum available memory. The optimization does not kick
  in if the loop is already unrolled in the input HLO.
  • all_gather_combiner
  • all_reduce_combiner
  • reduce_scatter_combiner

Memory Fitting Level Copy Chapter

Another effort level option controls the degree to which the compiler will
attempt to make the resulting program "fit in memory", where "fit" and "memory"
have backend-dependent meanings (for example, in XLA:TPU, this option controls
the degree to which the compiler works to keep the TPU's high-bandwidth memory
(HBM) usage below the HBM capacity). It can be set via the
memory_fitting_level
field of the ExecutableBuildOptionsProto message, or the
memory_fitting_level
field of the ExecutionOptions message.

As with optimization level, the exact meaning of each effort level value is
backend-dependent, but the following table describes the expected effect as a
general guideline:

Emitters (`docs/emitters.md`) Copy Section

XLA:GPU Emitters

There are three ways how to generate code for HLO in XLA:GPU.

1. Replacing HLO with custom calls to external libraries, e.g. NVidia cuBLAS, cuDNN.
2. Tiling HLO to block-level and then using OpenAI Triton.
3. Using XLA Emitters to progressively lower HLO to LLVM IR.

This document is focused on XLA:GPU Emitters.

Hero-based codegen Copy Chapter

There are 7 emitter types in XLA:GPU. Each emitter type corresponds to a "hero"
of the fusion, i.e. the most important op in the fused computation that shapes
the code generation for the whole fusion.

For example, the tranpose emitter will be selected if there is a
HloTransposeInstruction within the fusion that requires using shared memory to
improve the memory reading and writing patterns. The reduction emitter generates
reductions using shuffles and shared memory. The loop emitter is the default
emitter. If a fusion does not have a hero for which we have a special emitter,
then the loop emitter will be used.

High-level overview Copy Chapter

The code consists of the following big building blocks:

• Computation partitioner - splitting an HLO fusion computation into functions
• Emitters - converting partitioned HLO fusion to MLIR (xla_gpu, tensor, arith, math, scf dialects)
• Compilation pipeline - optimizes and lowers IR to LLVM

Partitioning Copy Chapter

See computation_partitioner.h.

Non-elementwise HLO instructions cannot always be emitted together. Consider the
following HLO graph:

param
       |
      log
      |  \
      |  transpose
      |  /
      add

If we emit this in a single function, the log will be accessed at two
different indices for each element of the add. The old emitters solve this
problem by generating the log twice. For this particular graph, this is not
a problem, but when there are multiple splits, the code size grows
exponentially.

Here, we solve this problem by partitioning the graph into pieces that can be
safely emitted as one function. The criteria are:

• Instructions that have only one user are safe to emit together with their
  user.
• Instructions that have multiple users are safe to emit together with their
  users if they are accessed through the same indices by all users.

In the example above, the add and tranpose access different indices of the
log, so it is not safe to emit it together with them.

The graph is therefore partitioned into three functions (each containing just
one instruction).

The same is applicable to the following example with slice and pad of add.

Elemental emission Copy Chapter

See elemental_hlo_to_mlir.h.

Elemental emission creates loops and math/arith ops for HloInstructions. For
the most part, this is straightforward, but there are some interesting things
going on here.

Indexing transformations Copy Section

Some instructions (transpose, broadcast, reshape, slice, reverse and
a few more) are purely transformations on indices: to produce an element of the
result, we need to produce some other element of the input. For this, we can
reuse XLA's indexing_analysis, which has
functions to produce the output to input mapping for an instruction.

For example, for a transpose from [20,40] to [40,20], it will produce the
following indexing map (one symbolic expression per input dimension; d0 and d1
are the output dimensions):

(d0, d1) -> d1
  (d0, d1) -> d0

So for these pure index transformation instructions, we can simply get the map,
apply it to the output indices, and produce the input at the resulting index.

Similarly, the pad op uses indexing maps and constraints for most of the
implementation. pad is also an indexing transformation with some added checks
to see if we return an element of the input or the padding value.

Tuples Copy Section

We do not support internal tuples. We also do not support nested tuple
outputs. All XLA graphs that use these features can be converted to graphs that
do not.

Gather Copy Section

We only support canonical gathers as produced by gather_simplifier.

Subgraph functions Copy Chapter

For a subgraph of a computation with parameters %p0 to %p_n, and subgraph
roots with r dimensions and element types (e0 to e_m), we use the
following MLIR function signature:

(%p0: tensor<...>, %p1: tensor<...>, ..., %pn: tensor<...>,
 %i0: index, %i1: index, ..., %i_r-1: index) -> (e0, ..., e_m)

That is, we have one tensor input per computation parameter, one index input per
dimension of the output, and one result per output.

To emit a function, we simply use the elemental emitter above, and recursively
emit its operands until we reach the edge of the subgraph. Then, we:emit a
tensor.extract for parameters or emit a func.call for other subgraphs

Entry function Copy Chapter

Each emitter type differs in how it generates the entry function, i.e. the
function for the hero. The entry function is different from the functions above,
since it has no indices as inputs (just the thread and block IDs) and actually
needs to write the output somewhere. For the loop emitter, this is fairly
straightforward, but the transpose and reduction emitters have non-trivial write
logic.

The signature of the entry computation is:

(%p0: tensor<...>, ..., %pn: tensor<...>,
 %r0: tensor<...>, ..., %rn: tensor<...>) -> (tensor<...>, ..., tensor<...>)

Where like before, the %pns are the parameters of the computation, and the
%rns are the results of the computation. The entry computation takes the
results as tensors, tensor.inserts updates into them, and then returns them.
No other uses of the output tensors are allowed.

Compilation pipeline Copy Chapter

Loop emitter Copy Section

See loop.h.

Let's study the most important passes of the MLIR compilation pipeline using the
HLO for the GELU function.

This HLO computation only has elementwise ops, constants and broadcasts. It will
be emitted using the loop emitter.

MLIR Conversion

After conversion to MLIR we get an xla_gpu.loop that depends on
%thread_id_x and %block_id_x and defines the loop that traverses all
elements of the output linearly to guarantee coalesced writes.

On every iteration of this loop we call

%pure_call = xla_gpu.pure_call @gelu(%input, %dim0, %dim1, %dim2)
      : (tensor<6x512x4096xbf16>, index, index, index) -> bf16

to compute elements of the root operation. Note, that we have only one outlined
function for @gelu, because the partitioner did not detect a tensor that has 2
or more various access patterns.

/* Detailed source-code truncated for AI context efficiency. */

Inliner

After @gelu is inlined, we get a single @main function. It can happen that
the same function is called twice or more. In this case we don't inline. More
details on the inlining rules can be found in
xla_gpu_dialect.cc.

func.func @main(%arg0: tensor<6x512x4096xbf16>, %arg1: tensor<6x512x4096xbf16>) -> tensor<6x512x4096xbf16> {
 ...
  %thread_id_x = gpu.thread_id  x {xla.range = [0 : index, 127 : index]}
  %block_id_x = gpu.block_id  x {xla.range = [0 : index, 24575 : index]}

  %xla_loop = xla_gpu.loop (%thread_id_x, %block_id_x)[%vector_index] -> (%dim0, %dim1, %dim2)
      in #map iter_args(%iter = %output) -> (tensor<6x512x4096xbf16>) {
    %extracted = tensor.extract %input[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
    %0 = arith.mulf %extracted, %extracted : bf16
    %1 = arith.mulf %0, %extracted : bf16
    %2 = arith.mulf %1, %cst : bf16
    %3 = arith.addf %extracted, %2 : bf16
    %4 = arith.mulf %3, %cst_0 : bf16
    %5 = math.tanh %4 : bf16
    %6 = arith.addf %5, %cst_1 : bf16
    %7 = arith.mulf %6, %cst_2 : bf16
    %8 = arith.mulf %extracted, %7 : bf16
    %inserted = tensor.insert %8 into %iter[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
    xla_gpu.yield %inserted : tensor<6x512x4096xbf16>
  }
  return %xla_loop : tensor<6x512x4096xbf16>
}

`xla_gpu` to `scf` conversion

See lower_xla_gpu_to_scf.cc.

xla_gpu.loop represents a loop nest with a boundary check inside. If the loop
inductions variables are out of bounds of the indexing map domain, then this
iteration is skipped. It means, that the loop is converted to 1 or more nested
scf.for ops with an scf.if inside.

%xla_loop = scf.for %vector_index = %c0 to %c4 step %c1 iter_args(%iter = %output) -> (tensor<6x512x4096xbf16>) {
   %2 = arith.cmpi sge, %thread_id_x, %c0 : index
   %3 = arith.cmpi sle, %thread_id_x, %c127 : index
   %4 = arith.andi %2, %3 : i1
   %5 = arith.cmpi sge, %block_id_x, %c0 : index
   %6 = arith.cmpi sle, %block_id_x, %c24575 : index
   %7 = arith.andi %5, %6 : i1
   %inbounds = arith.andi %4, %7 : i1
   %9 = scf.if %inbounds -> (tensor<6x512x4096xbf16>) {
     %dim0 = xla_gpu.apply_indexing #map(%thread_id_x,  %block_id_x)[%vector_index]
     %dim1 = xla_gpu.apply_indexing #map1(%thread_id_x, %block_id_x)[%vector_index]
     %dim2 = xla_gpu.apply_indexing #map2(%thread_id_x, %block_id_x)[%vector_index]
     %extracted = tensor.extract %input[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
     // ... more arithmetic operations
     %29 = arith.mulf %extracted, %28 : bf16
     %inserted = tensor.insert %29 into %iter[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
     scf.yield %inserted : tensor<6x512x4096xbf16>
   } else {
     scf.yield %iter : tensor<6x512x4096xbf16>
   }
   scf.yield %9 : tensor<6x512x4096xbf16>
 }

Flatten tensors

See flatten_tensors.cc.

The N-d tensors are projected onto 1D. This will simplify the vectorization and
the lowering to LLVM because every tensor access now corresponds to how the data
is aligned in memory.

#map = #xla_gpu.indexing_map<"(th_x, bl_x, vector_index) -> (th_x * 4 + bl_x * 512 + vector_index),"
 "domain: th_x in [0, 127], bl_x in [0, 24575], vector_index in [0, 3]">

func.func @main(%input: tensor<12582912xbf16>, %output: tensor<12582912xbf16>) -> tensor<12582912xbf16> {
 %xla_loop = scf.for %vector_index = %c0 to %c4 step %c1 iter_args(%iter = %output) -> (tensor<12582912xbf16>) {
   %dim = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x, %vector_index)
   %extracted = tensor.extract %input[%dim] : tensor<12582912xbf16>
   %2 = arith.mulf %extracted, %extracted : bf16
   %3 = arith.mulf %2, %extracted : bf16
   %4 = arith.mulf %3, %cst_2 : bf16
   %5 = arith.addf %extracted, %4 : bf16
   %6 = arith.mulf %5, %cst_1 : bf16
   %7 = math.tanh %6 : bf16
   %8 = arith.addf %7, %cst_0 : bf16
   %9 = arith.mulf %8, %cst : bf16
   %10 = arith.mulf %extracted, %9 : bf16
   %inserted = tensor.insert %10 into %iter[%dim] : tensor<12582912xbf16>
   scf.yield %inserted : tensor<12582912xbf16>
 }
 return %xla_loop : tensor<12582912xbf16>
}

Vectorization

See vectorize_loads_stores.cc.

The pass analyses the indices in the tensor.extract and tensor.insert ops
and if they are produced by xla_gpu.apply_indexing that accesses the elements
contiguously w.r.t. to the %vector_index and the access is aligned, then
tensor.extract is converted to vector.transfer_read and hoisted out of the
loop.

In this particular case, there is an indexing map
(th_x, bl_x, vector_index) -> (th_x * 4 + bl_x * 512 + vector_index) used to
compute elements to extract and insert in a scf.for loop from 0 to 4.
Therefore, both tensor.extract and tensor.insert can be vectorized.

func.func @main(%input: tensor<12582912xbf16>, %output: tensor<12582912xbf16>) -> tensor<12582912xbf16> {
 %vector_0 = arith.constant dense<0.000000e+00> : vector<4xbf16>
 %0 = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x, %c0)
 %2 = vector.transfer_read %input[%0], %cst {in_bounds = [true]} : tensor<12582912xbf16>, vector<4xbf16>
 %xla_loop:2 = scf.for %vector_index = %c0 to %c4 step %c1
     iter_args(%iter = %output, %iter_vector = %vector_0) -> (tensor<12582912xbf16>, vector<4xbf16>) {
   %5 = vector.extract %2[%vector_index] : bf16 from vector<4xbf16>
   %6 = arith.mulf %5, %5 : bf16
   %7 = arith.mulf %6, %5 : bf16
   %8 = arith.mulf %7, %cst_4 : bf16
   %9 = arith.addf %5, %8 : bf16
   %10 = arith.mulf %9, %cst_3 : bf16
   %11 = math.tanh %10 : bf16
   %12 = arith.addf %11, %cst_2 : bf16
   %13 = arith.mulf %12, %cst_1 : bf16
   %14 = arith.mulf %5, %13 : bf16
   %15 = vector.insert %14, %iter_vector [%vector_index] : bf16 into vector<4xbf16>
   scf.yield %iter, %15 : tensor<12582912xbf16>, vector<4xbf16>
 }
 %4 = vector.transfer_write %xla_loop#1, %output[%0] {in_bounds = [true]}
     : vector<4xbf16>, tensor<12582912xbf16>
 return %4 : tensor<12582912xbf16>
}

Loop unrolling

See optimize_loops.cc.

The loop unrolling finds scf.for loops that can be unrolled. In this case, the
loop over the elements of the vector disappears.

func.func @main(%input: tensor<12582912xbf16>, %arg1: tensor<12582912xbf16>) -> tensor<12582912xbf16> {

  %cst_0 = arith.constant dense<0.000000e+00> : vector<4xbf16>
  %dim = xla_gpu.apply_indexing #map(%thread_id_x, %block_id_x, %c0)
  %2 = vector.transfer_read %input[%dim], %cst {in_bounds = [true]} : tensor<12582912xbf16>, vector<4xbf16>
  %3 = vector.extract %2[%c0] : bf16 from vector<4xbf16>
  ...
  %13 = vector.insert %12, %cst_0 [%c0] : bf16 into vector<4xbf16>
  %14 = vector.extract %2[%c1] : bf16 from vector<4xbf16>
  ...
  %24 = vector.insert %23, %13 [%c1] : bf16 into vector<4xbf16>
  %25 = vector.extract %2[%c2] : bf16 from vector<4xbf16>
  ...
  %35 = vector.insert %34, %24 [%c2] : bf16 into vector<4xbf16>
  %36 = vector.extract %2[%c3] : bf16 from vector<4xbf16>
  ...
  %46 = vector.insert %45, %35 [%c3] : bf16 into vector<4xbf16>
  %47 = vector.transfer_write %46, %arg1[%dim] {in_bounds = [true]} : vector<4xbf16>, tensor<12582912xbf16>
  return %47 : tensor<12582912xbf16>
}

Conversion to LLVM

We mostly use the standard LLVM lowerings, but there are a few special passes.
We cannot use the memref lowerings for tensors, since we don't bufferize the
IR and our ABI is not compatible with the memref ABI. Instead, we have a
custom lowering directly from tensors to LLVM.

• The lowering of tensors is done in lower_tensors.cc. tensor.extract is
  lowered to llvm.load, tensor.insert to llvm.store, in the obvious way.
• propagate_slice_indices and merge_pointers_to_same_slice together
  implement a detail of buffer assignment and XLA's ABI: if two tensors share
  the same buffer slice, they are only passed once. These passes deduplicate the
  function arguments.

llvm.func @__nv_tanhf(f32) -> f32
llvm.func @main(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
  %11 = nvvm.read.ptx.sreg.tid.x : i32
  %12 = nvvm.read.ptx.sreg.ctaid.x : i32
  %13 = llvm.mul %11, %1 : i32
  %14 = llvm.mul %12, %0 : i32
  %15 = llvm.add %13, %14 : i32
  %16 = llvm.getelementptr inbounds %arg0[%15] : (!llvm.ptr, i32) -> !llvm.ptr, bf16
  %17 = llvm.load %16 invariant : !llvm.ptr -> vector<4xbf16>
  %18 = llvm.extractelement %17[%2 : i32] : vector<4xbf16>
  %19 = llvm.fmul %18, %18  : bf16
  %20 = llvm.fmul %19, %18  : bf16
  %21 = llvm.fmul %20, %4  : bf16
  %22 = llvm.fadd %18, %21  : bf16
  %23 = llvm.fmul %22, %5  : bf16
  %24 = llvm.fpext %23 : bf16 to f32
  %25 = llvm.call @__nv_tanhf(%24) : (f32) -> f32
  %26 = llvm.fptrunc %25 : f32 to bf16
  %27 = llvm.fadd %26, %6  : bf16
  %28 = llvm.fmul %27, %7  : bf16
  %29 = llvm.fmul %18, %28  : bf16
  %30 = llvm.insertelement %29, %8[%2 : i32] : vector<4xbf16>
  ...
}

Transpose emitter Copy Section

Let's consider a slightly more involved example.

The transpose emitter differs from the loop emitter only in how the entry
function is generated.

func.func @transpose(%arg0: tensor<20x160x170xf32>, %arg1: tensor<170x160x20xf32>) -> tensor<170x160x20xf32> {
  %thread_id_x = gpu.thread_id  x {xla.range = [0 : index, 127 : index]}
  %block_id_x = gpu.block_id  x {xla.range = [0 : index, 959 : index]}

  %shmem = xla_gpu.allocate_shared : tensor<32x1x33xf32>
  %xla_loop = xla_gpu.loop (%thread_id_x, %block_id_x)[%i, %j]
      -> (%input_dim0, %input_dim1, %input_dim2, %shmem_dim0, %shmem_dim1, %shmem_dim2)
      in #map iter_args(%iter = %shmem) -> (tensor<32x1x33xf32>) {
    %extracted = tensor.extract %arg0[%input_dim0, %input_dim1, %input_dim2] : tensor<20x160x170xf32>
    %0 = math.exp %extracted : f32
    %inserted = tensor.insert %0 into %iter[%shmem_dim0, %shmem_dim1, %shmem_dim2] : tensor<32x1x33xf32>
    xla_gpu.yield %inserted : tensor<32x1x33xf32>
  }

  %synced_tensor = xla_gpu.sync_threads %xla_loop : tensor<32x1x33xf32>

  %xla_loop_0 = xla_gpu.loop (%thread_id_x %block_id_x)[%i, %j] -> (%dim0, %dim1, %dim2)
      in #map1 iter_args(%iter = %arg1) -> (tensor<170x160x20xf32>) {
    // indexing computations
    %extracted = tensor.extract %synced_tensor[%0, %c0, %1] : tensor<32x1x33xf32>
    %2 = math.absf %extracted : f32
    %inserted = tensor.insert %2 into %iter[%3, %4, %1] : tensor<170x160x20xf32>
    xla_gpu.yield %inserted : tensor<170x160x20xf32>
  }
  return %xla_loop_0 : tensor<170x160x20xf32>
}

In this case, we generate two xla_gpu.loop ops. The first one performs
coalesced reads from the input and writes the result to the shared memory.

The shared memory tensor is created using xla_gpu.allocate_shared op.

After the threads are synchronized using xla_gpu.sync_threads, the second
xla_gpu.loop reads the elements from the shared memory tensor and performs
coalesced writes to the output.

Reproducer Copy Section

In order to see the IR after every pass of the compilation pipeline, one can
launch run_hlo_module with the --xla_dump_emitter_re=mlir-fusion flag.

run_hlo_module --platform=CUDA --xla_disable_all_hlo_passes --reference_platform="" /tmp/gelu.hlo --xla_dump_emitter_re=mlir-fusion --xla_dump_to=<some_directory>

where /tmp/gelu.hlo contains

HloModule m:

gelu {
  %param = bf16[6,512,4096] parameter(0)
  %constant_0 = bf16[] constant(0.5)
  %bcast_0 = bf16[6,512,4096] broadcast(bf16[] %constant_0), dimensions={}
  %constant_1 = bf16[] constant(1)
  %bcast_1 = bf16[6,512,4096] broadcast(bf16[] %constant_1), dimensions={}
  %constant_2 = bf16[] constant(0.79785)
  %bcast_2 = bf16[6,512,4096] broadcast(bf16[] %constant_2), dimensions={}
  %constant_3 = bf16[] constant(0.044708)
  %bcast_3 = bf16[6,512,4096] broadcast(bf16[] %constant_3), dimensions={}
  %square = bf16[6,512,4096] multiply(bf16[6,512,4096] %param, bf16[6,512,4096] %param)
  %cube = bf16[6,512,4096] multiply(bf16[6,512,4096] %square, bf16[6,512,4096] %param)
  %multiply_3 = bf16[6,512,4096] multiply(bf16[6,512,4096] %cube, bf16[6,512,4096] %bcast_3)
  %add_1 = bf16[6,512,4096] add(bf16[6,512,4096] %param, bf16[6,512,4096] %multiply_3)
  %multiply_2 = bf16[6,512,4096] multiply(bf16[6,512,4096] %add_1, bf16[6,512,4096] %bcast_2)
  %tanh_0 = bf16[6,512,4096] tanh(bf16[6,512,4096] %multiply_2)
  %add_0 = bf16[6,512,4096] add(bf16[6,512,4096] %tanh_0, bf16[6,512,4096] %bcast_1)
  %multiply_1 = bf16[6,512,4096] multiply(bf16[6,512,4096] %add_0, bf16[6,512,4096] %bcast_0)
  ROOT %multiply_0 = bf16[6,512,4096] multiply(bf16[6,512,4096] %param, bf16[6,512,4096] %multiply_1)
}

ENTRY main {
  %param = bf16[6,512,4096] parameter(0)
  ROOT fusion = bf16[6,512,4096] fusion(%param), kind=kLoop, calls=gelu
}

Links to code Copy Chapter

• Compilation pipeline: emitter_base.h
• Optimization and conversion passes: backends/gpu/codegen/emitters/transforms
• Partition logic: computation_partitioner.h
• Hero-based emitters: backends/gpu/codegen/emitters
• XLA:GPU ops: xla_gpu_ops.td
• Correctness and lit tests: backends/gpu/codegen/emitters/tests

Error Codes (`docs/error_codes.md`) Copy Section

XLA Error codes

This page is a list of all error codes emitted by the XLA compiler.

• E0100 - Runtime: Buffer Allocation Failure
• E0101 - Runtime: Program Allocation Failure
• E0102 - Runtime: Program Input Buffer Mismatch
• E0200 - Runtime: Core Halted Unexpectedly
• E1000 - Compile Time: HBM OOM
• E1001 - Compile Time: Scoped Vmem OOM
• E1200 - Compile Time: Host Offload Output Mismatch
• E2001 - Compile Time: Unsupported RHS DataType on
  Hardware
• E2002 - Compile Time: Mosaic Input/Output
  Misaligned Block and Tiling
• E2003 - Compile Time: Mosaic Unproven Memory
  Access Alignment
• E3000 - Compile Time: SparseCore Allocation
  Failure
• E3001 - Compile Time: SparseCore No Viable Logical
  Replica Count

Flags Guidance (`docs/flags_guidance.md`) Copy Section

XLA Flags Guidance

This guide offers a curated selection of key XLA flags to assist users
in effectively navigating and utilizing XLA's capabilities. The following
sections detail flags that can significantly impact runtime performance and
memory utilization. Should any issues, such as crashes, arise after enabling a
flag, it is recommended to revert to the default setting and create a
GitHub issue.

Correctness Flags Copy Chapter

Performance Flags Copy Chapter

The following flags are instrumental in enhancing runtime performance.
Experimenting with these settings may lead to considerable performance gains.

CPU Performance Flags Copy Section

Memory Flags Copy Chapter

The flags listed below are provided to address HBM-related issues. These
should only be adjusted if you encounter HBM "out of memory" errors during model
compilation. In all other scenarios, the default values are recommended, as
altering them could adversely affect performance.

Other commonly used flags Copy Chapter

TPU XLA flags Copy Section

GPU XLA flags Copy Section

The -O1 optimization level enables advanced compiler passes for improved GPU
performance, including data-parallel collective pipelining
(xla_gpu_pipeline_all_gather, xla_gpu_pipeline_all_reduce, and
xla_gpu_pipeline_reduce_scatter), while loop unrolling
(xla_gpu_enable_while_loop_double_buffering), latency hiding scheduling
(xla_gpu_enable_latency_hiding_scheduler), and SOL latency estimation on
Hopper/Blackwell (xla_gpu_enable_analytical_sol_latency_estimator). See
GPU Effort Levels for details.

Gpu Aot (`docs/gpu_aot.md`) Copy Section

XLA:GPU ahead-of-time (AOT) compilation

Ahead-of-time (AOT) compilation compiles and serializes an XLA computation in
one environment, allowing the resulting binary artifact to be loaded and
executed later by a separately deployed runtime.

Unlike just-in-time (JIT) compilation, which compiles computations in the
running process upon first execution, AOT removes compilation from the serving
path, which reduces startup latency and lets compilation run in a dedicated
build environment.

Note: If you are here because a reviewer flagged a backward or forward
compatibility problem with your change, go to the
GPU AOT compatibility guide.

Workflow and APIs Copy Chapter

The standard AOT pipeline consists of compile, serialize, and load/deserialize
phases:

• PjRt entry points: In the C++ PjRt interface
  (xla/pjrt/pjrt_client.h),
  compilation is initiated via PjRtClient::Compile (or CompileAndLoad).
  The resulting executable is serialized into a string via
  PjRtExecutable::SerializeExecutable. On the serving host, the runtime
  reconstructs the executable using PjRtClient::DeserializeExecutable or
  loads it directly for execution with PjRtClient::LoadSerializedExecutable.

• Compiler-level artifacts: At the compiler layer, ahead-of-time
  compilation produces a
  GpuAotCompilationResult
  (a CompiledModule;
  AotCompilationResult
  is a deprecated alias). This object encapsulates the serialized
  GpuExecutableProto
  and exposes SerializeAsString() and LoadExecutable().

• Command-line tooling: The
  xla_compile
  tool compiles HLO or StableHLO modules into serialized executables:

  bash Copy

  xla_compile --module_file=module.hlo --output_file=output.bin \
      --platform=gpu --gpu_target_config=gpu_target_config.pbtxt

  Compilation can run on a host without a GPU by providing a target device
  configuration (see the
  xla/backends/gpu/target_config README).

What is serialized Copy Chapter

The top-level container for serialized GPU programs is
GpuExecutableProto:

• thunks: The sequence of execution thunks. Kernel thunks (for example
  CustomKernelThunk) embed their own compiled kernel binary (cubin for
  NVIDIA, hsaco for ROCm), so most of the device code lives here.
• binary: A module-level device binary. Today it carries only the module's
  constant buffers; kernels are no longer stored here.
• buffer_allocations: Buffer sizes and allocation slices.
• gpu_compute_capability: Hardware architecture capability of the compile
  target.
• executable_abi_version: Toolchain and runtime ABI versions.
• hlo_module_with_config: Embedded serialized HloModuleProto and debug
  metadata.

Every proto message and field transitively reachable from GpuExecutableProto
is part of the compatibility contract described in the
GPU AOT compatibility guide.

How an artifact is loaded and executed Copy Chapter

1. Reconstruct the thunk graph: Deserialize ThunkProto entries into the
   in-memory thunk sequence, resolving custom call, FFI, and kernel symbols
   along the way. The embedded HLO module is parsed as well, so an HLO opcode
   unknown to the loading binary fails the load.
2. Run thunk graph conversion passes: Execute load-time optimization and
   conversion passes on the in-memory thunks (for example, constructing runtime
   CommandBufferThunks
   to submit work via CUDA Graphs). CommandBufferThunk is reconstructed
   dynamically at load time and is never directly serialized.
3. Allocate buffers: Initialize memory spaces and buffer allocations
   according to the serialized buffer assignment slices.
4. Execute: Dispatch the prepared thunk sequence on device streams.

Pass, fusion, codegen, and autotuning changes that keep the serialized schema
and runtime semantics intact do not affect AOT compatibility.

Troubleshooting: load failures that are not versioning bugs Copy Chapter

Not all executable load failures are caused by serialization regressions. Two
common environment mismatches produce load-time errors:

• Compute capability mismatch:
  GpuExecutable::FromProto
  validates that the serialized gpu_compute_capability matches the target
  device. An executable compiled for NVIDIA Hopper cannot execute on NVIDIA
  Ampere.
• CUDA toolchain ABI mismatch:
  CudaRuntimeAbiVersion::IsCompatibleWith
  checks that the runtime host's CUDA toolkit, cuDNN, and CUB versions satisfy
  or exceed the minimum versions recorded in ExecutableAbiVersionProto. If
  the host driver or library environment is older than the compile target,
  loading fails with FailedPreconditionError.

Gpu Aot Compatibility (`docs/gpu_aot_compatibility.md`) Copy Section

XLA:GPU AOT compatibility guide

This guide defines the compatibility contract between the XLA:GPU compiler that
produces a serialized executable and the runtime that loads it, and tells you
how to land a change without breaking that contract. For what AOT compilation is
and how an artifact is structured and loaded, see
XLA:GPU ahead-of-time compilation.

The contract applies exclusively to the XLA:GPU backend
(GpuExecutableProto).
CPU and TPU AOT executables use different serialization mechanisms and are out
of scope.

Compatibility windows {#compatibility} Copy Chapter

XLA:GPU provides two compatibility windows:

• Backward compatibility (6 months): A runtime containing your change must
  load and execute, with identical semantics, every artifact emitted by
  compilers over the preceding 6 months. This allows loading models compiled
  some time ago on an updated runtime.
• Forward compatibility (2 weeks): A compiler containing your change must
  not emit any serialization construct that a runtime up to 2 weeks older
  cannot deserialize or would misinterpret. This allows loading models
  compiled with the latest compiler on older, long-running inference servers
  during rolling deployments.

Every proto message and field transitively reachable from GpuExecutableProto
is part of this contract.

Change triage: is my change safe? {#change-triage} Copy Chapter

Follow this decision sequence to determine whether your change requires a staged
rollout:

1. Does your change touch any of the following?

   • a proto reachable from GpuExecutableProto
   • ToProto/FromProto
   • HLO opcodes or HLO proto that can end up in a GPU executable
   • names or semantics of registered kernel symbols, custom calls, or FFI
     handlers
   • load-time thunk passes or flags the runtime reads
   • execution semantics of existing thunks

   No: The change is compiler-internal (e.g. optimizations, fusion logic,
   codegen rewrites) or a purely internal runtime refactor with unchanged
   behavior for existing thunks. It is safe to land directly.

   Yes: Proceed to step 2.

2. Does the change alter how the runtime interprets existing artifacts
   (execution of existing thunks, load-time passes, defaults of runtime-read
   flags)?

   Yes: Allowed only if correct for every artifact emitted in the last 6
   months; if the semantic meaning changes, introduce a new field or flag
   instead.

   No: Proceed to step 3.

3. Does the change remove an existing thunk kind, proto field, or
   deserialization fallback branch?

   Yes: Violates backward compatibility unless compiler emission ceased
   at least 6 months ago. Keep the deserialization fallback until the 6-month
   window has elapsed, and mark deleted proto tags and names as reserved.

   No: Proceed to step 4.

4. Does the change emit a new thunk kind, new proto field required for
   correctness, or new HLO opcode into the serialized artifact?

   Yes: Violates forward compatibility if released immediately. Follow
   the staged rollout.

   No: Proceed to step 5.

5. Does the change rename or delete a registered kernel symbol, custom call,
   or FFI handler name?

   Yes: Violates backward compatibility. Retain backwards-compatible
   registrations under legacy names for at least 6 months.

   No: Proceed to step 6.

6. Does the change add an optional proto field that older runtimes can safely
   ignore without correctness regressions?

   Yes: Safe to land, provided unset proto fields preserve legacy runtime
   semantics.

   No: Redesign the change to preserve backward and forward compatibility.

Safe vs. breaking changes Copy Chapter

Staged rollout for breaking changes {#staged-rollout} Copy Chapter

When introducing a new serialization construct, land it in three phases so that
both compatibility windows are satisfied:

1. Runtime support. Add deserialization and execution logic to the runtime.
   If compiler emission is included in the same PR, keep emission disabled by
   default behind an experimental flag. Then wait at least 2 weeks from the
   merge of this change so the updated runtime propagates across deployment
   environments.
2. Compiler emission. Enable compiler emission by default (or flip the
   default value of the feature flag). Then keep the legacy deserialization
   fallback in the runtime for at least 6 months, counted from this phase
   (when the compiler stopped emitting the legacy form), not from phase 1.
3. Cleanup. Delete the legacy deserialization path and mark the deprecated
   proto tag numbers and field names as reserved.

Phases 1 and 2 can be two pull requests (the second sent once the first is 2
weeks old) or a single PR where emission is disabled by default behind a flag
and flipped in a follow-up.

Removal comments and examples Copy Section

Every compatibility fallback in the runtime should document its expiration date:

// Backward-compatibility fallback for legacy AOT-compiled kernels without
// an explicit CollectiveKernelSpec.
// Can be removed in <month year> (6 months backward compatibility window).

Public OpenXLA rollout examples:

• PR #49046: Introduced dedicated
  stream assignment in the runtime first, keeping emission disabled by default
  during the rollout window.
• PR #46865: Added runtime thunk
  support for CollectiveReduce, with compiler emission guarded behind an
  experimental flag.

Authoring runtime code for deserialized executables Copy Chapter

When implementing runtime thunks and deserialization routines:

• No compiler-only types: Runtime thunks must not depend on compiler data
  structures such as HloInstruction* or BufferAssignment.
• Validate deserialized bounds: Always validate buffer allocation indices,
  slice offsets, and stream IDs against allocation bounds before indexing.
• Do not duplicate payloads: Do not store the same large payload (for
  example constants) more than once in the artifact; it multiplies load-time
  memory.
• Guard initialization allocations: Serving runtimes operate under strict
  memory limits. Avoid unconditional heap allocations during FromProto or
  thunk initialization.

Testing Copy Chapter

All serialization changes must be covered by automated tests:

• Round-trip unit tests: Verify serialization and deserialization fidelity
  via ToProto and FromProto tests (see
  ThunkProtoDeserializationTest).
• Golden AOT tests: End-to-end compatibility tests live under
  xla/tests/aot_compatibility/gpu
  (such as collective_ops_aot_test.cc), running against serialized golden
  executables in executables/<target>/v<N>/.
• Testing all historical versions: By default, golden tests validate
  against the boundary versions (oldest and second-newest). Set
  XLA_AOT_TEST_ALL_VERSIONS=1 in your test environment to execute against
  all historical golden versions.

Gpu Architecture (`docs/gpu_architecture.md`) Copy Section

XLA:GPU Architecture Overview

Introduction

XLA is a hardware- and framework- domain-specific compiler for linear algebra,
offering best-in-class performance. JAX, TF, Pytorch and others use XLA by
converting the user input to
StableHLO (“high-level
operation”: a set of ~100 statically shaped instructions like addition,
subtraction, matmul, etc) operation set, from which XLA produces optimized code
for a variety of backends:

During the execution, the frameworks invoke the
PJRT runtime
API, which lets the frameworks perform the operation “populate the specified
buffers using a given StableHLO program on a specific device”.

XLA:GPU Pipeline

XLA:GPU uses a combination of “native” (PTX, via LLVM) emitters and TritonIR
emitters to generate high-performance GPU kernels (blue color indicates 3P
components):

Running Example: JAX Copy Chapter

To illustrate the pipeline, let’s start with a running example in JAX, which
computes a matmul combined with multiplication by a constant and negation:

def f(a, b):
    return -((a @ b) * 0.125)

We can inspect the HLO generated by the function:

M = 1024
K = 512
N = 2048
key = jax.random.PRNGKey(1701)
a = jax.random.randint(key, (M, K), dtype=jax.numpy.int8, minval=0, maxval=255)
b = jax.random.normal(key, (K, N), dtype=jax.dtypes.bfloat16)

print(jax.xla_computation(f)(a, b).as_hlo_text())

which generates:

HloModule xla_computation_f, entry_computation_layout={(s8[1024,512]{1,0}, bf16[512,2048]{1,0})->(bf16[1024,2048]{1,0})}

ENTRY main.10 {
  Arg_0.1 = s8[1024,512]{1,0} parameter(0)
  convert.5 = bf16[1024,512]{1,0} convert(Arg_0.1)
  Arg_1.2 = bf16[512,2048]{1,0} parameter(1)
  dot.6 = bf16[1024,2048]{1,0} dot(convert.5, Arg_1.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}
  constant.3 = bf16[] constant(0.125)
  broadcast.4 = bf16[1024,2048]{1,0} broadcast(constant.3), dimensions={}
  multiply.7 = bf16[1024,2048]{1,0} multiply(dot.6, broadcast.4)
  ROOT negate.8 = bf16[1024,2048]{1,0} negate(multiply.7)
}

We can visualize the input HLO computation as well, using
jax.xla_computation(f)(a, b).as_hlo_dot_graph():

Optimizations on HLO: Key Components Copy Chapter

A number of notable optimization passes happen on HLO, as HLO->HLO rewrites.

SPMD Partitioner Copy Section

The XLA SPMD partitioner, as described in
GSPMD: General and Scalable Parallelization for MLComputation Graphs,
consumes HLO with sharding annotations (produced e.g. by jax.pjit), and
produces a sharded HLO which can then run on a number of hosts and devices.
Apart from partitioning, the SPMD attempts to optimize HLO for an optimal
execution schedule,
overlapping computation
and communication between the nodes.

Example

Consider starting from a simple JAX program sharded across two devices:

# Defines a mesh with two axes called ‘x’ and ‘y’,
# sharded across two devices: first and second CPU.
with jax.sharding.Mesh(
      [['cpu:0', 'cpu:1']], ('x', 'y')):

    @pjit
    def f(a, b):
        out = -((a @ b) * 0.125)
        # Shard output matrix access across ‘x’
        # and ‘y’ respectively. Generates ‘Sharding’
        # custom call.
        out = with_sharding_constraint(
          out, jax.lax.PartitionSpec('x', 'y'))
        return out

# Random inputs to call our function.
a = jax.random.randint(key, (1024, 512), jnp.int8)
b = jax.random.normal(key, (512, 2048), jnp.float32)

print(f.lower(a, b).compiler_ir())

Visualizing it, the sharding annotations are presented as custom calls:

To check how the SPMD partitioner expands the custom call, we can look at HLO
after optimizations:

print(f.lower(np.ones((8, 8)).compile().as_text())

Which generates HLO with a collective:

Layout Assignment Copy Section

HLO decouples logical shape and physical layout (how tensors are laid out in
memory). For example, a matrix f32[32, 64] can be represented either in
row-major or column-major order, represented as {1,0} or {0,1} respectively.
In general, layout is represented as a part of shape, showing a permutation over
the number of dimensions indicating physical layout in memory.

For each operation present in the HLO, the Layout Assignment pass chooses an
optimal layout (e.g. NHWC for a convolution on Ampere). For example, an
int8xint8->int32  matmul operation prefers {0,1} layout for the RHS of the
computation. Similarly, “transposes” inserted by the user are ignored, and
encoded as a layout change.

The layouts are then propagated through the graph, and conflicts between layouts
or at graph endpoints are materialized as copy operations, which perform the
physical transposition. For example, starting from the graph

Running the layout assignment we see the following layouts and copy operation
inserted:

Fusion Copy Section

Fusion is XLA’s single most important optimization, which groups multiple
operations (e.g. addition into exponentiation into matmul) to a single kernel.
Since many GPU workloads tend to be memory-bound, fusion dramatically speeds up
the execution by avoiding the writing of intermediate tensors to HBM and then
reading them back, and instead passes them around in either registers or shared
memory.

Fused HLO instructions are blocked together in a single fusion computation,
which establishes the following invariants:

• No intermediate storage inside the fusion is materialized in HBM (it has to
  be all passed through either registers or shared memory).

• A fusion is always compiled to exactly one GPU kernel

HLO Optimizations on Running Example Copy Chapter

We can inspect the post-optimization HLO using jax.jit(f).lower(a, b).compile().as_text(), and verify that a single fusion got generated:

/* Detailed source-code truncated for AI context efficiency. */

Note that the fusion backend_config tells us that Triton will be used as a
code generation strategy, and it specifies the chosen tiling.

We can also visualize the resulting module:

Buffer Assignment and Scheduling Copy Chapter

A buffer assignment pass takes into account the shape information, and aims to
produce an optimal buffer allocation for the program, minimizing the amount of
intermediate memory consumed. Unlike TF or PyTorch immediate-mode (non-compiled)
execution, where the memory allocator does not know the graph in advance, the
XLA scheduler can “look into the future” and produce an optimal computation
schedule.

Compiler Backend: Codegen and Library Selection Copy Chapter

For every HLO instruction in the computation, XLA chooses whether to run it
using a library linked into a runtime, or to codegen it to PTX.

Library Selection Copy Section

For many common operations, XLA:GPU uses fast-performance libraries from NVIDIA,
such as cuBLAS, cuDNN, and NCCL. The libraries have an advantage of verified
fast performance, but often preclude complex fusion opportunities.

Direct code generation Copy Section

The XLA:GPU backend generates high-performance LLVM IR directly for a number of
operations (reductions, transposes, etc).

Triton code generation Copy Section

For more advanced fusions which include matrix multiplication or softmax,
XLA:GPU uses Triton as a code-generation
layer. HLO Fusions are converted to TritonIR (an MLIR dialect which serves as an
input to Triton), selects tiling parameters and invokes Triton for PTX
generation:

We have observed the resulting code to perform very well on Ampere, at
near-roofline performance with properly tuned tile sizes.

Runtime Copy Chapter

XLA Runtime converts the resulting sequence of CUDA kernel calls and library
invocations into a RuntimeIR (an MLIR dialect in XLA), on which CUDA graph
extraction is performed. CUDA graph is still work in progress, only some nodes
are currently supported. Once CUDA graph boundaries are extracted, RuntimeIR is
compiled via LLVM to a CPU executable, which can then be stored or transferred
for Ahead-Of-Time compilation.

Hlo Dumps (`docs/hlo_dumps.md`) Copy Section

Dump HLO Computations

An HLO dump is a textual representation of the HLO modules at different stages
of the computation. It is useful for debugging, and you often need to include it
in bug reports. This is typically a human-readable text file that lists the
HLO instructions and their properties. Sometimes, HLO modules are dumped as:

• HloProto: Protocol buffer files, which are a more structured,
  machine-readable format.
• HloSnapshot: HLO module plus its inputs. For replaying HLOs, you
  sometimes require the actual inputs fed to a given computation rather than
  random data.

You can use XLA flags to specify and get dumps. In most cases, you can set it
with an environment variable. JAX also offers a programmatic way to print the
HLO dump.

Local Execution Copy Chapter

Using Environment Variables Copy Section

You can set the XLA_FLAGS environment variable with the necessary flags to get
dumps. This works for JAX, TensorFlow, and PyTorch/XLA.

To dump HLO modules and other debugging information to a specific directory, run
your program with the --xla_dump_to flag:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

For example, you can use /tmp or /tmp/xladump as the paths.

By default, this dumps HLO modules as text, at the very beginning and end of the
optimization pipeline.

You can also explicitly specify the format:

1. Text dumps

XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"

1. HLO protos

XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"

1. HLO Snapshots

XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"

1. Graph render with graphviz server (only works well for small graphs)

XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"

1. Graph render to HTML file (only works well for small graphs)

XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

For larger graphs, you can use interactive_graphviz to visualize parts of the
graph.

Note: If --xla_dump_to is not specified but another dumping flag is
specified, it will dump to stdout. But the dump will not include binary data,
e.g., proto files, to stdout.

You can also set the HLO Dumps to use syntactic sugar wrappers as op names, by
setting the --xla_syntax_sugar_async_ops flag to true. This can reduce the
dump by about 20%. By default, this flag is set to false, and actual op names
are used in the dump.

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_syntax_sugar_async_ops=true"

Dump Specific Intermediate Passes Copy Chapter

In addition to the standard pre-optimized / final-optimized HLOs, you can also
dump the state of HLOs after a particular compiler pass.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

HLO modules will be dumped for the passes whose names match the regular
expression (regex). For example, you can observe the HLOs resulting from passes
related to SPMD partitioning with:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

To dump the result after every XLA pass (this will result in a lot of files),
you can set:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

JAX-specific Options Copy Section

Programmatically in JAX

Instead of passing flags or environment variables, you can also programmatically
dump HLO using JAX’s lower and compile APIs.

Locally fetch the unoptimized original lowered HLO with:

jax.jit(f).lower(*args).as_text('hlo')

For dumping to files during HLO compilation passes, specify:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

Dump jaxprs

jaxprs are JAX's intermediate
representation for program traces. To dump this, set the environment variables:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Learn more in JAX documentation on
Exporting and serializing staged-out computations: Debugging.

Google Colab Copy Chapter

Environment variables Copy Section

In the first executed cell of your notebook (because environment variables and
command-line flags are usually only processed once, e.g., at module-import time
or XLA backend initialization time), add the XLA_FLAGS detailed above with
os.environ, for example:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

This will dump the computation to DIRECTORY_PATH, for example /tmp. On
Colab, navigate to the "Files" browser in the left sidebar, to view and access
this directory.

You can use all the flags mentioned in the Local Execution section.

JAX-specific options Copy Section

Similar to local execution; for live, interactive introspection you can directly
print a computation’s pre-optimized HLO:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

You can also directly print a computation’s optimized HLO:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

Dumping All/Small Computations

If you want to see everything in a dump including all small compilations, set
the JAX environment variable:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mosaic

Mosaic is a compiler for the Pallas TPU backend, and the experimental Pallas GPU
backend. To dump mosaic computation, set the following flag:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Or, set TPU init arguments as an environment variable:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

Check out the
JAX documentation on Pallas and Mosaic
to learn more.

More with HLO Dumps Copy Chapter

Finding the right computation Copy Section

Usually, many computations get dumped. The dumped files are explicitly named
with the JAX, Tensorflow, or PyTorch/XLA "computation name” that are called out
in the logs, making it easy to identify the relevant HLO files. For example:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Otherwise, you can use ripgrep to quickly identify which module holds
particular symbols or computations.

Tip: Include the 3 dumped before/after/buffer-assignment files of interest
in your bug reports.

HLO Conversion Copy Section

A tool called hlo-opt that can translate between HLOProto and text formats.
It's useful in cases where you have one format, but need the other for
debugging.

Learn to use it:
XLA Tooling documentation: hlo-opt.

Replay Copy Section

You can run (replay) the dumped computations on a specified XLA backend with
fake data or input snapshots. This is a convenient way to reproduce, iterate,
and debug issues in XLA.

The following commands use fake data. If you have saved HLO Snapshots, you can
pass those in instead, and the data from the snapshot will be used. To still use
fake data while running the snapshot, pass the flag --force_fake_data.

CPU backend:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

GPU backend:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt

Hlo Isolation (`docs/hlo_isolation.md`) Copy Section

HLO Isolation User Guide

This document explains how to install and use the HLO Isolation API and CLI.
The HLO isolation tooling helps developers and researchers isolate, verify, and
debug numeric mismatches and stability issues in compiled HLO modules.

Installation Copy Chapter

You can use HLO isolation components through the standard OpenXLA/TensorFlow
build mechanism (Bazel) or via binary distribution.

Source and Bazel Setup Copy Section

When building the compiler and tools stack from source, include or depend on the
following libraries:

• API Target:
  //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_api
• CLI Target:
  //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test

To build the standalone CLI tool:

bazel build -c opt //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test

Command-Line Interface (CLI) Copy Chapter

The hlo_isolation_test CLI allows you to isolate and run numeric mismatch and
stability checks against compiled HLO modules directly from a terminal. It
compares the execution results across different environments (e.g., TPU vs.
Defused TPU / CPU / Interpreter).

Flag Reference Copy Section

The CLI supports the following flags:

• --hlo_file: Path to the input .hlo or .pbtxt file to load. Can be text
  or proto format (required).
• --test_platform: Target platform to run the primary test on (e.g., cpu,
  gpu, tpu). Defaults to cpu.
• --reference_platform: Reference platform for baseline comparison (e.g.,
  interpreter). If empty, reference comparison is disabled.
• --filter_by_name: Regular expression to match the module name. Only
  matching modules will be run. Defaults to .*.
• --skip_by_name: Regular expression to match the module name. Matching
  modules will be skipped.
• --filter_by_opcode: Regular expression to match instruction opcodes. Only
  modules containing at least one matching opcode will be run. Defaults to
  .*.
• --skip_by_opcode: Regular expression to match instruction opcodes. Modules
  containing any matching opcode will be skipped.
• --abs_error_bound: Absolute error bound used for comparison. Defaults to
  0.01.
• --rel_error_bound: Relative error bound used for comparison. Defaults to
  0.1.
• --run_hlo_passes: Boolean flag to determine whether to run standard HLO
  passes on the submodules. Defaults to false.
• --shard_index: The specific shard index to run (zero-based). Defaults to
  -1 (disabled).
• --num_shards: The total number of shards. Defaults to 1.

Basic Invocation Copy Section

./hlo_isolation_test \
  --hlo_file=/path/to/failing_fusion.hlo \
  --test_platform=gpu \
  --reference_platform=interpreter

Result and Artifact Dumps Copy Chapter

When a submodule encounters a numeric mismatch or other failure during isolation
testing, the tool automatically serializes debug artifacts to disk for deeper
inspection.

Dump Contents Copy Section

On numeric mismatch, the tool writes the following debug artifacts:

1. The interactive HTML mismatch report (failed-module-<module_name>.html).
2. The failed HLO submodule text (failed-module-<module_name>.txt).
3. The expected output literal (failed-<module_name>-expected.txt).
4. The actual output literal (failed-<module_name>-actual.txt).
5. The mismatching elements summary (failed-<module_name>-mismatches.txt).

Interactive HTML Graph UI Copy Section

When numeric mismatches are detected during isolation testing, an interactive
HTML report (failed-module-<module_name>.html) is generated for each mismatch
detected. It provides a graphical view of the failing HLO module by rendering
the HLO graph and highlighting mismatching instructions with annotations and
tooltips that display detailed numeric mismatch statistics (including actual
and expected values, relative and absolute error bounds, and element mismatch
percentages).

Above: Example HTML
output for a fusion containing a numeric mismatch.

Dump Target Location Copy Section

• Test Environment: If run via bazel test or an environment defining the
  TEST_UNDECLARED_OUTPUTS_DIR environment variable, the results are placed
  directly in that directory with the exact file names listed above (e.g.,
  failed-<module_name>-expected.txt).
• Standard/Manual Run: When executed manually via the command line, the
  artifacts are written to the operating system's temporary directory (e.g.,
  /tmp), preserving the exact same unified file naming conventions (e.g.,
  /tmp/failed-<module_name>-expected.txt).

C++ Integration and API Copy Chapter

For developers building custom compiler passes, testing rigs, or automated
pipelines, the C++ API provides a direct way to integrate isolation tests.

Using the API Directly Copy Section

The core API provides functional interfaces to run modules and fetch structured
reports:

#include "third_party/tensorflow/compiler/xla/tools/hlo_isolation/hlo_isolation_api.h"

xla::hlo_isolation::PipelineIsolationOptions options;
options.module_options.abs_error_bound = 0.01;
options.module_options.rel_error_bound = 0.1;
// Filter specific opcodes programmatically
options.filter_by_opcode = "exponential";

absl::StatusOr<std::vector<xla::HloIsolationTestResult>> results =
    xla::hlo_isolation::RunIsolationPipeline(
        input_hlo_module,
        &my_test_runner,
        &my_reference_runner,
        options);

Using the Test Mixin Copy Section

When writing GoogleTest C++ test suites, you can inherit from
HloIsolationTestMixin for built-in assertion handling. The base class must
provide both a test runner and a reference runner (e.g., via
HloPjRtInterpreterReferenceMixin):

#include "third_party/tensorflow/compiler/xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
#include "third_party/tensorflow/compiler/xla/tools/hlo_isolation/hlo_isolation_test_base.h"

class MyCustomPassIsolationTest : public xla::hlo_isolation::HloIsolationTestMixin<
    xla::HloPjRtInterpreterReferenceMixin<xla::HloPjRtTestBase>> {};

TEST_F(MyCustomPassIsolationTest, ChecksMyFusionSanity) {
  RunAndVerifyIsolationTest(my_failing_module);
}

Sharded Execution (K8s/Slurm) Copy Chapter

For large modules or heavy test matrices, you can partition execution across
multi-device clusters (such as Google Kubernetes Engine or Slurm) using the
sharding flags. Each shard_index deterministically runs an isolated subset of
the decomposed submodules. This allows for reproducible distributed verification
and targeted re-execution of failing partitions.

Example: Kubernetes Job Copy Section

Each test shard is executed as a separate Kubernetes pod using completionMode: Indexed. The JOB_COMPLETION_INDEX is passed directly to the CLI's
--shard_index flag.

apiVersion: batch/v1
kind: Job
metadata:
  name: hlo-isolation-job
spec:
  completions: 50
  parallelism: 50
  completionMode: Indexed
  template:
    spec:
      containers:
      - name: test-runner
        image: gcr.io/my-project/hlo-isolation-tools:latest
        command: ["/bin/sh", "-c"]
        args:
        - |
          ./hlo_isolation_test \
            --num_shards=50 \
            --shard_index=$JOB_COMPLETION_INDEX \
            --hlo_file=/data/path/to/hlo.hlo
        volumeMounts:
        - name: hlo-data-volume
          mountPath: /data
      volumes:
      - name: hlo-data-volume
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: my-xla-debug-bucket

Key Capabilities Copy Chapter

• Portability: Decouples internal test wrappers from the standalone API,
  making it easy to debug HLO mismatches locally.
• Granularity: Granular opcode and name filtering improves the debugging
  loop when interacting with massive HLO dumps.
• Interactive Visualization: Dumps interactive HTML reports with graph
  visualization and per-instruction numeric mismatch annotations for fast
  root-cause diagnosis.
• Extensibility: Custom runner execution callbacks and data injectors
  (make_fake_arguments_fn) permit full customization for advanced
  verification workflows.

Hlo Passes (`docs/hlo_passes.md`) Copy Section

HLO Passes

This document outlines the HLO
optimizations and transformations passes in the
XLA compiler.

Introduction Copy Chapter

A single HLO Pass can be comprised of one or many compiler optimizations and
transformations, and XLA provides several hundred such passes. HLO focuses only
on the shape (e.g. a 3x4 matrix) and the
operation semantics of the arrays
to make the optimization or transformation easier.

For example:

• AlgebraicSimplifier:
  A pass that performs a number of mostly arithmetic simplifications and
  optimizations. Including:

  • When dividing by a constant, an optimization is performed to transform
    the operation to multiplication by the inversion of the constant.

• HloRematerialization:
  A pass that recomputes selected expressions in the computation to reduce
  memory pressure caused by long live ranges of array-shaped values.

Developer details Copy Chapter

The base class for HLO passes can be found in
xla/hlo/pass/hlo_pass_interface.h.
HLO pass should not extend this class directly but instead should extend
HloModulePass.

See also
XLA HLO Pass Framework.

Tooling and Testing Copy Section

XLA comes with multiple command line tools, including the hlo-opt tool. This
tool allows execution of an individual pass independent of the given platform
compilation stages. For more information see
Tooling.

For information on writing unit tests for HLO Passes see
Testing HLO Passes.

Hardware-independent HLO Pass Examples Copy Chapter

This section describes a few examples of passes shared across XLA backends. Some
passes may be specialized for specific backends, but the high-level
functionality is similar.

Shared passes or hardware-independent passes can be found in
xla/hlo/transforms.

Rematerialization Copy Section

See also
HloRematerialization.

Selectively recomputes expressions within the HLO graph to reduce memory usage.
Trades off higher compute for lower memory usage. Can reduce memory usage by
tens of percent and is required to run many large models.

Algebraic Simplifier Copy Section

See also
AlgebraicSimplifier.

A grab bag of simplifications, optimizations, and canonicalizations. Analogous
to
LLVM’s instcombine pass.

Constant Folding Copy Section

See also
HloConstantFolding.

Replaces expressions which can be evaluated at compile time with their constant
equivalent.

Dead Code Elimination Copy Section

See also
HloDCE
.

Removes operations with unused results (fast implementation).

Call Graph Flattening Copy Section

See also
FlattenCallGraph.

A legalization pass which converts the HLO call graph into a tree by cloning
computations. Required because memory is statically assigned to HLO operations
and not based on dynamic call context.

Reshape Mover Copy Section

See also
ReshapeMover.

Reshapes and transposes can be expensive, especially on TPU. This pass moves and
reshapes and transposes across elementwise operations enabling the operations to
be merged or eliminated.

Zero-sized HLO Elimination Copy Section

See also
ZeroSizedHloElimination.

HLO supports arrays of zero size (one or more dimensions has a bound of zero).
This pass simplifies the graph by replacing zero-sized operations with
zero-sized constants.

TPU-specific HLO Pass Examples Copy Chapter

Passes specific to the TPU backend.

Model parallelism Copy Section

The partitioning of an XLA program across multiple cores is performed at the HLO
level and the TPU HLO pipeline includes a number of passes for supporting
multi-core execution.

Spatial partitioning

See also
ShardingPropagation.

Pass to support dividing operations across devices along non-batch dimensions.

Handling of bfloat16 Copy Section

See also
BFloat16ConversionFolding,
BFloat16MixedPrecisionRemoval,
and
BFloat16Propagation.

TPUs support bfloat16 as a lower-precision, more compact floating-point
representation than 32-bit floats. Using bfloat16 reduces memory footprint and
memory bandwidth. The TPU HLO pipeline includes various passes for replacing
floats with bfloat16 into the program and propagating the precision through the
graph.

Legalization passes Copy Section

See also
GatherExpander,
and
BatchNormExpander.

Passes which transform unsupported HLO into a form which the backend can emit or
for which the backend produces a more efficient lowering.

GPU-specific HLO Pass Example Copy Chapter

Passes specific to the GPU backend are found in
xla/service/gpu.
These passes can be identified as classes defined in namespace gpu.

cuDNN Rewriter Copy Section

See also
CudnnFusedConvRewriter
and
CudnnNormRewriter.

Rewrites fused convolution and norm operations into their respective library
calls in cuDNN.

CPU-specific HLO Pass Examples Copy Chapter

Passes specific to the CPU backend are found in
xla/service/cpu.
These passes can be identified as classes defined in namespace cpu.

Convolution Canonicalization Copy Section

See also
ConvCanonicalization.

Canonicalizes convolutions so that they can be lowered to a fast implementation
in Eigen.

Operation Parallelization Copy Section

See also
ParallelTaskAssigner.

Partitions HLOs into tasks to run on separate threads.

Analysis passes Copy Chapter

Analysis passes are not considered "HLO passes" since they do not transform HLO
and may not extend HloModulePass. Shared analyses are found in
xla/hlo/analysis.

Analysis Pass Examples Copy Section

Dataflow Analysis

See also
HloDataflowAnalysis.

Identifies all HLO values in the graph and their uses.

Alias Analysis

See also
HloAliasAnalysis.

Identifies must-alias relationships between values in the program.

Computation Cost Analysis

See also
HloCostAnalysis.

Computes FLOP count and memory usage for all operations in the program.

HLO Verification

See also
HloVerifier.

Verifies various invariants of the HLO graph.

Hlo To Thunks (`docs/hlo_to_thunks.md`) Copy Section

From HLO to Thunks

This document outlines the journey of an XLA High Level Optimizer (HLO) module
from its initial state to a final executable. Sometimes we will omit the
"module" and refer to it just as "HLO".

Pre-optimization HLO Copy Chapter

We start with pre-optimization HLO module. Pre-optimization HLO does not contain
operations (ops) that are considered internal to XLA, such as fusion or
bitcast. Ops don't have a layout at this stage, or if they do, it will be
ignored. Pre-optimization HLO is usually produced by higher-level frameworks
like TensorFlow and JAX. When using the XLA flag -xla_dump_to, the
pre-optimization HLO is dumped to a file with file name suffix
“before_optimizations.txt”.

Optimize HLO Module Copy Chapter

The XLA:GPU pipeline turns the pre-optimization HLO into optimized HLO by
running a sequence of passes. The passes can be grouped together semantically
and run in the following order:

Sharding related passes Copy Section

This includes passes like the Shardy
Partitioner or those for SPMD sharding.

Optimization passes Copy Section

This can include both legalization passes and simplification passes.

Collective optimization passes Copy Section

Similar to Optimization passes, but focuses on collective ops.

Layout assignment passes Copy Section

Each HLO op is assigned a layout which is a part of the instruction shape. The
layout controls how the tensor is laid out physically in memory.

Example of a shape with layout:

f32[10,20,30]{2,0,1}

After the element type, there are the logical dimensions of the shape, followed
by the layout permutation in minor to major order. In this example, the most
minor dimension is 30, the second most minor dimension is 10, and the major
dimension is 20.

The goal of layout assignment is to minimize the number of required physical
transpositions using a greedy strategy. It starts with certain layout
constraints (e.g., cuDNN/cuBLAS libraries expect consecutive dimensions) and
propagates layouts “down” and then “up” the HLO graph. At the
end of layout propagation, some instructions may have conflicting layouts, one
propagated from an operand, one propagated from a user. To resolve this
conflict, a copy HLO instruction is inserted that changes the layout from the
operand layout to the instruction layout.

Layout normalization passes Copy Section

Given that it is somewhat difficult to figure out the physical shape, layout
normalization attempts to rewrite the shape such that it uses the default layout
{rank-1, rank-2, …, 0}. In the example above, the normalized shape would be
f32[20,10,30]{2,1,0}. Copy ops that change layouts are rewritten as a
combination of transpose and bitcast. Given that currently we cannot
normalize all ops, there are still some ops that may have non-default layouts,
most notably gather and dot. At the boundaries between normalized ops and
non-normalized ops there will be bitcast ops that represent a transpose, i.e.
a transpose with a layout assigned that makes it a no-op physically.

Layout normalization also makes some implicit transposes explicit which is
important because codegen can handle explicit transposes with a dedicated
emitter. For example, a reshape is technically allowed to have a different
physical layout between operand and result (e.g. due to different rank). The
ReshapeDecomposer pass that runs as part of the layout normalization passes
turns a reshape into a sequence of transpose, reshape bitcast and
transpose.

Post layout assignment optimization passes Copy Section

The most important passes here are Triton fusions (GEMM fusions +
Softmax/Layernorm fusions) or rewrites to library calls. Autotuning also runs in
this step, where XLA chooses chooses between different emitters, picks the best
algorithm for convolutions or dots, finds the best tiling for fusions handled by
the Triton emitter etc.

Fusion passes Copy Section

The two main passes are PriorityFusion and Multi-Output fusion.

In PriorityFusion, we form fusions guided by the cost model. When fusing we
would allow duplicating ops with several users if the op can be fused into all
users. We would also allow extending existing Triton Softmax fusions if
possible.

Multi-Output fusion is a separate pass that allows fusing ops/fusions that
share an operand. It can also fuse operands/operand fusions into users without
duplication by adding extra output(s), so other users of the op to be fused can
be redirected to these outputs. This pass needs to be careful not to introduce
cycles into the HLO graph.

After Multi-Output fusion, common subexpression elimination (HloCSE pass)
runs, potentially merging previously duplicated ops back together if they ended
up in the same fusion.

Several post-fusion passes Copy Section

Several passes related to collectives (like turning them to async, or enforcing
a certain relative order of collectives).

Finally we run CopyInsertion where copies are added to ensure that in-place
operations don't overwrite data that is still needed elsewhere.

At the end of optimization, the optimized HLO is dumped if using the flag
-xla_dump_to to a file that has the file name suffix
"after_optimizations.txt". If you want to dump the HLO after intermediate
passes that actually change the HloModule, you can use the flag
-xla_dump_hlo_pass_re=.* (or a specific regular expression to restrict it to
certain passes).

Scheduling Copy Chapter

An HLO Module without a schedule still has some degree of freedom in the order
in which ops are processed. Any topological sort respecting operand/result
relationships and control dependencies is valid. Scheduling determines what
specific order to use. The main concern at this stage is the maximum memory
consumption that depends on the lifetime of tensors. In an initial step, we try
different scheduler algorithms and pick the schedule that should minimize the
peak memory consumption. Note that at this point we don't work with a physical
buffers yet (that will happen in "Buffer Assignment") and simulate the memory
usage.

Then LatencyHidingScheduler pass runs and tries to maximize
compute-communication overlap. But that may increase memory usage again.

Finally, in case peak memory consumption is higher than the amount of memory we
have available we run HloRematerialization. This pass attempts to reduce
memory usage at the cost of performance, as e.g. some fusions might be split and
some ops might be duplicated to have shorter buffer lifetimes. If
rematerialization occurs, it might be beneficial to investigate ways to reduce
memory requirements on the model side (e.g., using smaller batch sizes).

Buffer Assignment Copy Chapter

Immediately before we lower to LLVM IR, we run the buffer assignment passes that
will assign buffer slices to each instruction in the HLO graph. The buffer
assignment runs in several steps:

1. HloDataflowAnalysis assigns HloValues (essentially logical buffers) to
   instructions. For in-place ops, the HloValue of an operand can be reused. An
   op may define more than one HloValue (e.g. with a tuple result shape).

2. HloAliasAnalysis attempts to combine buffers for aliasing operations, and
   computes a mapping from HloValue to HloBuffer.

3. BufferAssignment computes a mapping of HloBuffers to buffer slices inside
   a big buffer in such a way that the same buffer slice is not used for different
   HloBuffers with overlapping life times. For ops that may alias, it is ok that
   there is a slight overlap (the end time of the one HloBuffer may coincide with
   the start time of the other HloBuffer). When using the flag -xla_dump_to,
   some information about buffer assignment is dumped to a file with the name
   suffix "after_optimizations-buffer-assignment.txt".

Thunks Copy Chapter

After an HLO graph is optimized and scheduled, it is lowered into a
linear sequence of thunks for a specific backend (CPU or GPU).

In XLA, a Thunk is an abstraction of a self-contained unit of work that the
runtime executes. It might be a compiled kernel launch, specific operation,
library call, control-flow construct, collective communication, and so on. A
Thunk Sequence represents the entire executable for a specific backend.

Thunk Emission Copy Section

The process of converting a scheduled HLO computation into a thunk sequence is
called "thunk emission". This is handled by a dedicated emitter class in each
backend.

For the GPU Backend, this is handled by
IrEmitterUnnested.
EmitHloComputation iterates through the scheduled list of HLO Instructions in
a computation and dispatches to a specialized Emit... method (e.g.,
EmitFusion, EmitConvolutionThunk, EmitWhile). Each of these methods
constructs the appropriate Thunk object(s) and appends them to the
thunk sequence.

For the CPU Backend,
ThunkEmitter
performs this role and is organized in a similar manner. Final ThunkSequence
is embedded in the CpuExecutable.

Note that each instruction in the entry computation of an HLO module might
correspond to no (kTuple, kConstant, ..), one, or multiple (for example
sort instruction) thunks in the final thunk sequence.

Command Buffers: Optimizing Execution on the GPU Copy Section

Modern GPU hardware allows recording a sequence of GPU operations (kernel
launches, memory copies, etc.) once and then replaying the sequence multiple
times with minimal CPU overhead. This is a critical performance optimization,
especially for workloads with many small, fast-launching kernels. XLA uses
Command Buffer as an abstraction of CUDA Graphs or HIP Graphs. The core
interface is defined in
GpuCommandBuffer.

A command buffer is represented in a thunk sequence by
CommandBufferThunk.

The emitter does not produce this thunk directly from HLO instructions. Instead,
this is done by
CommandBufferConversionPass
that runs on the ThunkSequence itself.

The pass identifies contiguous sub-sequences of compatible thunks (e.g., a
series of KernelThunks and GemmThunks). It then replaces the found
sub-sequence with a single CommandBufferThunk. The new thunk encapsulates the
logic of the original thunks as a list of lightweight Command objects. When a
CommandBufferThunk executes for the first time on a given GPU stream, it
"records" its sequence of commands into a hardware command buffer. On all
subsequent executions, it simply issues a single command to the GPU to "replay"
the recorded sequence. This avoids the CPU overhead of launching each individual
kernel.

Executable Copy Chapter

The final product of the XLA compilation pipeline is a self-contained,
platform-specific
Executable.
This object encapsulates all the information needed to run the compiled program
on a target device, such as a CPU or GPU. It is the bridge between the compiler
and the runtime. Modern runtimes like PJRT use slightly higher-level
abstractions (see
PjRtExecutable),
but these ultimately wrap a backend-specific executable.

An Executable contains several key pieces of information generated
during compilation. While the exact contents vary by backend, they generally
include:

• Compiled Code: This is the low-level machine code that will run on the device.
  For CPUs, this is typically one or more object files. For GPUs, this is the
  compiled device code in PTX or HSACO format, which is loaded onto the GPU at
  runtime.

• Execution Plan (ThunkSequence): The core of the runtime logic. This is a
  linear sequence of Thunk objects. Each thunk represents a single unit of work,
  such as launching a kernel, calling a library function (e.g., cuBLAS), or
  handling control flow. The runtime executes the program by iterating through
  this sequence.

• Memory Layout (BufferAssignment): This critical piece of metadata, produced by
  the BufferAssigner, describes the complete memory layout for the computation.
  It specifies the size of every buffer and how memory is allocated and reused
  for parameters, outputs, and temporary values. The runtime uses this to
  allocate device memory and pass the correct pointers to each thunk.

• (optional) HLO Module: For debugging and profiling, the executable often
  retains a reference to the final, optimized HloModule that it was compiled
  from.

The creation of the final executable is orchestrated by the compiler for each
specific backend. The RunBackend method of a Compiler implementation is the
final step in the compilation process, which packages all the compiled artifacts
into an Executable object.
GpuCompiler
and
CpuCompiler
target GPU and CPU respectively.

When a user calls Execute... on an executable, the runtime uses the
BufferAssignment to allocate memory, and then invokes the ThunkSequence to
launch the operations on the device using the compiled code.

Indexing (`docs/indexing.md`) Copy Section

Indexing Analysis

HLO indexing analysis is a dataflow analysis
that describes how elements of one tensor relate to another via "indexing
maps". For example, how indices of an HLO instruction output map to indices of
HLO instruction operands.

Example

For a broadcast from tensor<20xf32> to tensor<10x20x30xf32>

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

the indexing map from the output to input is (i, j, k) -> (j) for i in [0, 10], j in [0, 20] and k in [0, 30].

Motivation Copy Chapter

XLA uses several bespoke solutions to reason about coalescing, operand
utilization, and tiling schemes (more details below). The goal of indexing
analysis is providing a reusable component for such use cases. Indexing analysis
is built on
XLA's custom SymbolicExpr and SymbolicMap
infrastructure and adds HLO semantics.

Coalescing Copy Section

Reasoning about memory coalescing becomes feasible for non-trivial cases, when
we know what elements/slices of the inputs are read to compute an element of the
output.

Operand Utilization Copy Section

Operand utilization in XLA indicates how much each input of the instruction is
used assuming its output is fully used. Currently, utilization is also not
computed for a generic case. Indexing analysis allows us to compute utilization
precisely.

Tiling Copy Section

A tile/slice is hyper-rectangular subset of a tensor parameterized by offsets,
sizes and strides. Tile propagation is a way to compute tile parameters of the
producer/consumer of the op using the tiling parameters of the op itself. There
is already a
library
that does it for softmax and dot. Tile propagation can be made more generic and
robust if it is expressed via indexing maps.

Indexing map Copy Chapter

An indexing map is a combination of

• a symbolically expressed function that maps every element of one tensor A to
  ranges of elements in tensor B;
• constraints on valid function arguments, including function's domain.

Function arguments are split into 3 categories to better communicate their
nature:

• dimension variables of the tensor A or a GPU grid we are mapping from;
  values are known statically. Index elements are also called dimension
  variables.

• range variables. They define a one-to-many mapping and specify a set of
  elements in B used to compute a single value of A; values are known
  statically. The contracting dimension of a matrix multiplication is an example
  of a range variable.

• runtime variables that are only known at during execution. For example,
  indices argument of gather
  operation.

Result of the function is an index of the target B tensor.

In short, an indexing function from tensor A to tensor B for operation x
is

map_ab(index in A, range variables, runtime variables) -> index in B.

To better separate the types of mapping arguments we write them as:

map_ab(index in A)[range variables]{runtime variables} -> (index in B)

For example, let's look at the indexing maps for the reduce operation
f32[4, 8] out = reduce(f32[2, 4, 8, 16] in, 0), dimensions={0,3}:

• to map elements of in to out our function can be expressed as
  (d0, d1, d2, d3) -> (d1, d2). The constraints of the variables
  d0 in [0, 1], d1 in [0, 3], d2 in [0, 7], d3 in [0, 15] are defined by the
  shape of in.

• To map elements of out to in: out has only two dimensions, and reduction
  introduces two range variables that cover reducing dimensions. Thus the
  mapping function is (d0, d1)[s0, s1] -> (s0, d0, d1, s1), where (d0, d1)
  is index of out. s0, s1 are ranges defined by operation's semantics and
  span dimension 0 and 3 of the in tensor. The constraints are
  d0 in [0, 3], d1 in [0, 7], s0 in [0,1], s1 in [0, 15].

It's important to note that in most scenarios we are interested in mapping from
the elements of the output. For computation

C = op1(A, B)
E = op2(C, D)

we can talk about "indexing of B" meaning "mapping of elements of E into
the elements of B". This might be counter-intuitive compared to other types of
data-flow analysis that work from input toward outputs.

Constraints on variables enable optimization opportunities and help with
code generation. In the documentation and implementation constraints are also
referred to as domain as they define all valid combinations or argument values
of the mapping function. For many operation, constraints simply describe the
dimensions of tensors but for some operations they might be more complicated;
see examples below.

By having functions and argument constraints expressed symbolically and being
able to combine functions and constraints we can compute a compact indexing
mapping for an arbitrary large computation (fusion).

Expressiveness of symbolic function and constraints is a balance between
implementation complexity and optimization gains we get from having a more
precise representation. For some HLO operations we capture access patterns only
approximately.

Implementation Copy Chapter

To minimize recomputation, we need a framework for symbolic computation. This is
implemented as SymbolicExpr and SymbolicMap.

Note: Previously, mlir::AffineMap was used for symbolic computation.

A typical SymbolicMap looks like:

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

SymbolicMap has two types of parameters: dimensions and symbols.
Dimensions correspond to the dimension variables d; symbols correspond to
the range variables r and runtime variables rt. SymbolicMap does not
contain any metadata about constraints of the parameters, so we have to provide
them separately.

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    // struct Interval represents a closed interval [lower_bound, upper_bound].
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  SymbolicMap symbolic_map_;

  // A dimension variable represents a dimension of a tensor or a GPU grid.
  // Dimension variables correspond to the dimensions of the `symbolic_map_`.
  std::vector<Variable> dim_vars_;

  // A range variable represents a range of values, e.g. to compute a single
  // element of the reduction's result we need a range of values from the input
  // tensor. Range variables correspond to the front portion of the
  // symbols in `symbolic_map_`.
  std::vector<Variable> range_vars_;

  // A runtime variable represents a runtime symbol, e.g. a dynamic offset in of
  // a HLO dynamic-update-slice op. Runtime variables correspond to the back
  // portion of the symbols in `symbolic_map_`.
  std::vector<Variable> rt_vars_;

   // Inequality constraints for symbolic expressions. They restrict the feasible
  // set for the domain of the indexing map. It contains symbolic expressions
  // other than SymbolicDimExpr and SymbolicSymbolExpr.
  llvm::MapVector<SymbolicExpr, Interval> constraints_;
};

Code reference:
indexing_map.h#L114

dim_vars_ encode the inclusive box constraints for the dimension
variables d of the indexing map, which usually coincide with the
shape of the output tensor for ops like transpose, reduce, elementwise, dot, but
there are some exceptions like
HloConcatenateInstruction.

range_vars_ all values that range variables s take. The range variables
are needed when multiple values are necessary to compute a single element of the
tensor we are mapping from, e.g. for output->input indexing map of reductions or
input->output map for broadcasts.

rt_vars_ encode the feasible values in runtime. For example, the offset is
dynamic for a 1D HloDynamicSliceInstruction. The corresponding RTVar will
have feasible values between 0 and tensor_size - slice_size - 1.

constraints_ capture relations between values in form
<expression> in <range>, e.g. d0 + s0 in [0, 20]. Together with
Variable.bounds they define the "domain" of indexing function.

Let's study-by-example to understand what all of the above actually means.

Indexing Maps for Unfused Ops Copy Chapter

Elementwise Copy Section

For elementwise ops the indexing map is an identity.

p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  output = f32[10, 20] add(p0, p1)

The output to input map output -> p0:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

The input to output map p0 -> output:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

[Broadcast](https://openxla.org/xla/operation_semantics#broadcastindim) Copy Section

Broadcasting means that some of the dimensions will be removed when we map
output to input and added when we map input to output.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

The output to input map:

(d0, d1, d2) -> (d1),
domain:
d0 in [0, 9],
d1 in [0, 19],
d2 in [0, 29]

The input to output map:

(d0)[s0, s1] -> (s0, d0, s1),
domain:
d0 in [0, 19],
s0 in [0, 9],
s1 in [0, 29]

Note that now we have range variables s on the right side for the
input-to-output mapping. Those are the symbols that represent ranges of values.
For example, in this particular case every element of input with index d0 is
mapped to a 10x1x30 slice of the output.

[Iota](https://openxla.org/xla/operation_semantics#iota) Copy Section

Iota has no input tensor operand, so there is no input index arguments.

iota = f32[2,4] iota(), dimensions={1}

Output to input map:

(d0, d1) -> ()
domain:
d0 in [0, 1]
d1 in [0, 3]

Input to output map:

()[s0, s1] -> (s0, s1)
domain:
s0 in [0, 1]
s1 in [0, 3]

[DynamicSlice](https://openxla.org/xla/operation_semantics#dynamicslice) Copy Section

DynamicSlice has offsets known only at runtime.

src = s32[2, 2, 258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = s32[1, 2, 32] dynamic-slice(src, of1, of2, of3), dynamic_slice_sizes={1, 2, 32}

The output to input map from ds to src:

(d0, d1, d2){rt0, rt1, rt2} -> (d0 + rt0, d1 + rt1, d2 + rt2),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31],
rt0 in [0, 1],
rt1 in [0, 0],
rt2 in [0, 226]

Note that now we have rt on the right side for the input-to-output mapping.
Those are the symbols that represent runtime values. For example, in this
particular case for every element of the output with indices d0, d1, d2 we
access slice offsets of1, of2 and of3 to compute the index of the input.
The intervals for the runtime variables are derived by assuming that the entire
slice stays in bounds.

The output to input map for of1, of2 and of3:

(d0, d1, d2) -> (),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31]

[DynamicUpdateSlice](https://openxla.org/xla/operation_semantics#dynamicupdateslice) Copy Section

src = s32[20,30] parameter(0)
upd = s32[5,10] parameter(1)
of1 = s32[] parameter(2)
of2 = s32[] parameter(3)
dus = s32[20,30] dynamic-update-slice(
    s32[20,30] src, s32[5,10] upd, s32[] of1, s32[] of2)

The output to input map for src is trivial. It can be made more precise by
restricting the domain to the not-updated indices, but right now indexing maps
do not support inequality constraints.

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 19],
d1 in [0, 29]

The output to input map for upd:

(d0, d1){rt0, rt1} -> (d0 - rt0, d1 - rt1),
domain:
d0 in [0, 19],
d1 in [0, 29],
rt0 in [0, 15],
rt1 in [0, 20]

Note that now we have rt0 and rt1 that represent runtime values. In
this particular case for every element of the output with indices d0, d1 we
access slice offsets of1 and of2 to compute the index of the input. The
intervals for the runtime variables are derived by assuming that the entire
slice stays in bounds.

The output to input map for of1 and of2:

(d0, d1) -> (),
domain:
d0 in [0, 19],
d1 in [0, 29]

[Gather](https://openxla.org/xla/operation_semantics#gather) Copy Section

Only the simplified gather is supported. See gather_simplifier.h.

operand = f32[33,76,70] parameter(0)
indices = s32[1806,2] parameter(1)
gather = f32[1806,7,8,4] gather(operand, indices),
  offset_dims={1,2,3},
  collapsed_slice_dims={},
  start_index_map={0,1},
  index_vector_dim=1,
  slice_sizes={7,8,4}

The output to input map for operand:

(d0, d1, d2, d3){rt0, rt1} -> (d1 + rt0, d2 + rt1, d3),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
rt0 in [0, 26],
rt1 in [0, 68]

Note that now we have rt symbols that represent runtime values.

The output to input map for indices:

(d0, d1, d2, d3)[s0] -> (d0, s0),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
s0 in [0, 1]

The range variable s0 shows that we need the entire row (d0, *) of the
indices tensor to compute an element of the output.

[Transpose](https://openxla.org/xla/operation_semantics#transpose) Copy Section

Indexing map for transpose is a permutation of input/output dimensions.

p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}

The output to input map:

(d0, d1, d2, d3) -> (d0, d3, d1, d2),
domain:
d0 in [0, 2],
d1 in [0, 5],
d2 in [0, 127],
d3 in [0, 12287],

The input to output map:

(d0, d1, d2, d3) -> (d0, d2, d3, d1),
domain:
d0 in [0, 2],
d1 in [0, 12287],
d2 in [0, 5],
d3 in [0, 127]

[Reverse](https://openxla.org/xla/operation_semantics#rev_reverse) Copy Section

Indexing map for reverse changes the reverted dimensions to upper_bound(d_i) - d_i:

p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}

The output to input map:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3),
domain:
d0 in [0, 0],
d1 in [0, 16],
d2 in [0, 8],
d3 in [0, 8]

The input to output map:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3),
domain:
d0 in [0, 0],
d1 in [0, 16],
d2 in [0, 8],
d3 in [0, 8]

[(Variadic)Reduce](https://openxla.org/xla/operation_semantics#reduce) Copy Section

Variadic reduction have several inputs and several initial values, the map from
output to input adds the reduced dimensions.

p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
out = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=max

The output to input maps:

• out[0] -> p0:

(d0)[s0] -> (s0, d0),
domain:
d0 in [0, 9],
s0 in [0, 255]

• out[0] -> p0_init:

(d0) -> (),
domain:
d0 in [0, 9]

The input to output maps:

• p0 -> out[0]:

(d0, d1) -> (d1),
domain:
d0 in [0, 255],
d1 in [0, 9]

• p0_init -> out[0]:

()[s0] -> (s0),
domain:
s0 in [0, 9]

[Slice](https://openxla.org/xla/operation_semantics#slice) Copy Section

Indexing from output to input for slice results in a strided indexing map which
is valid for every element of the output. Mapping from the input to output is
restricted to a strided range of the elements in the input.

p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
  slice={[5:10:1], [3:20:7], [0:50:2]}

The output to input map:

(d0, d1, d2) -> (d0 + 5, d1 * 7 + 3, d2 * 2),
domain:
d0 in [0, 4],
d1 in [0, 2],
d2 in [0, 24]

The input to output map:

(d0, d1, d2) -> (d0 - 5, (d1 - 3) floordiv 7, d2 floordiv 2),
domain:
d0 in [5, 9],
d1 in [3, 17],
d2 in [0, 48],
(d1 - 3) mod 7 in [0, 0],
d2 mod 2 in [0, 0]

[Reshape](https://openxla.org/xla/operation_semantics#reshape) Copy Section

Reshapes come in different flavors.

Collapse shape

This is a "linearizing" reshape from N-D to 1D.

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

The output to input map:

(d0) -> (d0 floordiv 8, d0 mod 8),
domain:
d0 in [0, 31]

The input to output map:

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

Expand shape

This is an inverse "collapse shape" op, it reshapes a 1D input into N-D output.

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

The output to input map:

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

The input to output map:

(d0) -> (d0 floordiv 8, d0 mod 8),
domain:
d0 in [0, 31]

Generic reshape

These are the reshape ops that cannot be represented as a single expand or
collapse shape. They can be only represented as a composition of 2 or more
expand or collapse shapes.

Example 1: Linearization-delinearization.

p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

This reshape can be represented as a composition of collapse shape of
tensor<4x8xf32> to tensor<32xf32> and then a shape expansion to
tensor<2x4x4xf32>.

The output to input map:

(d0, d1, d2) -> (d0 * 2 + d1 floordiv 2, d2 + (d1 mod 2) * 4),
domain:
d0 in [0, 1],
d1 in [0, 3],
d2 in [0, 3]

The input to output map:

(d0, d1) -> (d0 floordiv 2, d1 floordiv 4 + (d0 mod 2) * 2, d1 mod 4),
domain:
d0 in [0, 3],
d1 in [0, 7]

Example 2: Expanded and collapsed subshapes

p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)

This reshape can be represented as a composition of two reshapes. The first one
collapses the outermost dimensions tensor<4x8x12xf32> to tensor<32x12xf32>
and the second one expand the innermost dimension tensor<32x12xf32> into
tensor<32x3x4xf32>.

The output to input map:

(d0, d1, d2) -> (d0 floordiv 8, d0 mod 8, d1 * 4 + d2),
domain:
d0 in [0, 31],
d1 in [0, 2]
d2 in [0, 3]

The input to output map:

(d0, d1, d2) -> (d0 * 8 + d1, d2 floordiv 4, d2 mod 4),
domain:
d0 in [0, 3],
d1 in [0, 7],
d2 in [0, 11]

Bitcast Copy Section

A bitcast op can be represented as a
sequence of transpose-reshape-transpose.
Therefore, its indexing maps are just a composition of indexing maps for this
sequence.

[Concatenate](https://openxla.org/xla/operation_semantics#concatenate) Copy Section

Output-to-input mapping for concat is defined for all inputs, but with
non-overlapping domains, i.e. only one of the inputs will be used at a time.

p0 = f32[2, 5, 7] parameter(0)
p1 = f32[2, 11, 7] parameter(1)
p2 = f32[2, 17, 7] parameter(2)
ROOT output = f32[2, 33, 7] concatenate(f32[2, 5, 7] p0, f32[2, 11, 7] p1, f32[2, 17, 7] p2), dimensions={1}

The output to inputs maps:

• output -> p0:

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]

• output -> p1:

(d0, d1, d2) -> (d0, d1 - 5, d2),
domain:
d0 in [0, 1],
d1 in [5, 15],
d2 in [0, 6]

• output -> p2:

(d0, d1, d2) -> (d0, d1 - 16, d2),
domain:
d0 in [0, 1],
d1 in [16, 32],
d2 in [0, 6]

The inputs to output maps:

• p0 -> output:

(d0, d1, d2) -> (d0, d1, d2),
domain:
d0 in [0, 1],
d1 in [0, 4],
d2 in [0, 6]

• p1 -> output:

(d0, d1, d2) -> (d0, d1 + 5, d2),
domain:
d0 in [0, 1],
d1 in [0, 10],
d2 in [0, 6]

• p2 -> output:

(d0, d1, d2) -> (d0, d1 + 16, d2),
domain:
d0 in [0, 1],
d1 in [0, 16],
d2 in [0, 6]

[Dot](https://openxla.org/xla/operation_semantics#dot) Copy Section

Indexing maps for dot are very similar to the ones of reduce.

p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
output = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

The output to inputs maps:

• output -> p0:

(d0, d1, d2)[s0] -> (d0, d1, s0),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]

• output -> p1:

(d0, d1, d2)[s0] -> (d0, s0, d2),
domain:
d0 in [0, 3],
d1 in [0, 127],
d2 in [0, 63],
s0 in [0, 255]

The inputs to output maps:

• p0 -> output: