# Technical Documentation: openxla/xla > ℹ️ **Provenance:** Multi-Tier Ingestion · [CodeWiki Reference](https://codewiki.google/github.com/openxla/xla) ### Overview (README.md) # XLA XLA (Accelerated Linear Algebra) is an open-source machine learning (ML) compiler for GPUs, CPUs, and ML accelerators. OpenXLA Ecosystem 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](https://openxla.org/) is the project's website. ## Get started If you want to use XLA to compile your ML project, refer to the corresponding documentation for your ML framework: * [PyTorch](https://pytorch.org/xla) * [TensorFlow](https://www.tensorflow.org/xla) * [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) 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 If you'd like to contribute to XLA, review [How to Contribute](docs/contributing.md) and then see the [developer guide](docs/developer_guide.md). ## Contacts * For questions, contact the maintainers - maintainers at openxla.org ## Resources * [Community Resources](https://github.com/openxla/community) ## Code of Conduct While under TensorFlow governance, all community spaces for SIG OpenXLA are subject to the [TensorFlow Code of Conduct](https://github.com/tensorflow/tensorflow/blob/master/CODE_OF_CONDUCT.md). --- ### Errors Overview (`docs/errors_overview.md`) # 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 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 Here is an index list with all [error codes](error_codes.md). --- ### Index (`docs/index.md`) # 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 - **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](https://mlir.llvm.org/) 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 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](architecture.md) and then read [Contributing](contributing.md). --- ### Megascale Overview (`docs/megascale/overview.md`) # 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](./debugging_workflow.md) 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](./performance_tuning.md) details how to configure and tune Megascale XLA flags to optimize collective communication performance and alleviate bottlenecks in multi-slice workloads. ## Terminology - **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`) # PJRT C++ Device API Overview ## Background [PJRT](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h) 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 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 _Full reference at [`pjrt_client.h > PjRtClient`](https://github.com/openxla/xla/blob/924b74d84de3760cc589fd1525c7346691d51df5/xla/pjrt/pjrt_client.h#L486)._ 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 _Full references at [`pjrt_client.h > PjRtDevice`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L102), and [`pjrt_device_description.h`](https://github.com/openxla/xla/blob/main/xla/pjrt/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 _Full reference at [`pjrt_client.h > PjRtMemorySpace`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L72)._ 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 _Full reference at [`pjrt_client.h > PjRtBuffer`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L1111)._ 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: ```cpp // Literal to Buffer absl::StatusOr> BufferFromHostBuffer(...) {...} // Buffer to Literal xla::Future<> ToLiteral(xla::MutableLiteralBase* literal) override {...} ``` APIs for creating a buffer have [Buffer Semantics](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L858) 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 _Full reference at [`pjrt_compiler.h > PjRtCompiler`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_compiler.h#L157)._ 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 _Full reference at [`pjrt_executable.h > PjRtExecutable`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_executable.h#L306), and [`pjrt_client.h > PjRtLoadedExecutable`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L1506)._ 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: ```cpp // Execute on addressable devices absl::StatusOr>>> Execute(absl::Span> argument_handles, ...) {...} // Execute assigned replica/partition on the specified device absl::StatusOr>> ExecuteSharded(absl::Span argument_handles, PjRtDevice* device, ...) {...} // Execute on specified device, single replica / partition absl::StatusOr>> ExecutePortable(absl::Span 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 ### Futures & Async Computations If any part of a plugin is implemented asynchronously, it _must_ properly implement futures. Consider the following program: ```py @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 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 ### Example Log 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`) 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 * Issues and feature requests can be filed in the [OpenXLA/xla repo](https://github.com/openxla/xla). * Questions regarding PJRT can be asked on the [OpenXLA Discord][discord]. ## Resources * [PJRT C API header](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) * [PJRT C API changelog](https://github.com/openxla/xla/blob/main/xla/pjrt/c/CHANGELOG.md) * [PJRT integration guide](https://github.com/openxla/xla/blob/main/xla/pjrt/c/docs/pjrt_integration_guide.md) * [PJRT design docs](https://drive.google.com/drive/folders/18M944-QQPk1E34qRyIjkqDRDnpMa3miN) * [PJRT API ABI versioning and compatibility](https://docs.google.com/document/d/1TKB5NyGtdzrpgw5mpyFjVAhJjpSNdF31T6pjPl_UT2o/edit) * [PJRT Plugin Mechanism design doc](https://docs.google.com/document/d/1Qdptisz1tUPGn1qFAVgCV2omnfjN01zoQPwKLdlizas/edit) * [OpenXLA 2024 Fall DevLab PJRT plugin tutorial slides](https://drive.google.com/file/d/1epUJkMONG2t06WOeMHz4Oi3F_-8cTuz-/view) ([recording](https://www.youtube.com/watch?v=2GlMqaNxP_w&list=PLlFotmaRrOzv2OIEpijqiHGmY7rpscFcj)) * [IREE PJRT plugin implementation](https://github.com/iree-org/iree/tree/main/integrations/pjrt) [discord]: https://discord.gg/ZKXq7b3V8A "Join on Discord" --- ### Tf2xla Index (`docs/tf2xla/index.md`) # XLA: Optimizing Compiler for Machine Learning [OpenXLA](https://openxla.org) is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes. ## Introduction 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 ### Explicit compilation with `tf.function(jit_compile=True)` 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](./tutorials/jit_compile.ipynb) for a more detailed usage example, and a [tutorial video](https://www.youtube.com/watch?v=cPAD9vLKE0c) on `jit_compile=True` usage. ### Usage with Keras For Keras models, `jit_compile=True` can be set as an argument to [`model.compile`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile): ``` model.compile(optimizer="adam", jit_compile=True) ``` ### Usage with distributed strategy XLA:GPU can be used with TF distributed strategy ([`MirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy) or [`MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/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 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](./tutorials/autoclustering_xla.ipynb). ## Inspect compiled programs 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](./operation_semantics.md), one per each compiled cluster. Attaching those when submitting XLA bug reports is extremely helpful! - `module_XXXX.ir-*.ll` Generated files in [LLVM](https://llvm.org/docs/LangRef.html) intermediate representation, with [NVPTX](https://llvm.org/docs/NVPTXUsage.html) intrinsics. - `module_XXXX.ptx` Generated [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html) 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 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`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/tools/run_hlo_module_main.cc) and iteratively running it on generated programs. ## Further reading - [OpenXLA Documentation](https://openxla.org) OpenXLA Documentation - [Known Issues](./known_issues.md) List of known issues with XLA+TF - [XLA - TensorFlow, Compiled](https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html): Read on Google Developers Blog - Check out the [XLA source](https://github.com/openxla/xla) on Github! ## XLA Frontends Apart from TensorFlow, XLA programs can be generated by: - [JAX](https://github.com/google/jax): Composable transformations of Python+NumPy programs - [Julia](https://github.com/JuliaTPU/XLA.jl): The Julia language for scientific computing - [PyTorch](https://github.com/pytorch/xla): PyTorch framework - [Nx](https://github.com/elixir-nx/nx): Numerical computing library for the Elixir programming language ## Talks ### Using XLA from TF using `jit_compile=True` ### XLA Overview --- ### Aliasing (`docs/aliasing.md`) # 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 For example, consider a trivial HLO module which simply adds `1` to its input: ```mlir 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: ```mlir 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`](https://github.com/openxla/xla/blob/main/xla/client/xla_builder.h) API. ## Defining aliasing at runtime The aliasing defined in the previous step is specified during *compilation*. During execution, you can use the [`LocalClient::RunAsync`](https://github.com/openxla/xla/blob/main/xla/client/local_client.h) API to choose whether to donate the buffer. Input buffers to the program are wrapped in [`ExecutionInput`](https://github.com/openxla/xla/blob/main/xla/service/executable.h)s, 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`) # 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 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 The XLA compiler takes model graphs from ML frameworks defined in [StableHLO](https://github.com/openxla/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](https://en.wikipedia.org/wiki/Common_subexpression_elimination), 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](http://llvm.org) 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](./developing_new_backend.md). 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`) # 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 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 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 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 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` 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 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`) # The Hitchhiker's Guide to fewer `#ifdef`s in XLA **TL;DR:** `#ifdef`s are bad for maintainability. Please try to avoid them. Scroll down for details on how. [TOC] ## Introduction 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 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:** `#ifdef`s 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 `#include`s 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 ### Category I - Skipping test cases in unit tests 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: ```cpp 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`](https://github.com/openxla/xla/blob/main/xla/stream_executor/device_description.h) has the information needed to make the same decision at runtime: ```cpp 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`](https://github.com/openxla/xla/blob/main/xla/stream_executor/semantic_version.h): * **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 Another class of preprocessor conditionals guards code that would also compile fine without the conditional: ```cpp 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: ```cpp 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 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): ```cpp #if TENSORFLOW_USE_ROCM #include #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) `#ifdef`s. #### 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: ```cpp // feature.h // Do NOT include any platform-specific headers (CUDA/ROCm) here. #include "absl/status/status.h" absl::Status Foo(); ``` ```cpp // feature_rocm.cc #include "feature.h" #include "something/something/rocm.h" absl::Status Foo() { // Do something ROCm-specific. } ``` ```cpp // feature_stub.cc #include "feature.h" #include "absl/status/status.h" absl::Status Foo() { return absl::UnimplementedError("This is a ROCm-only feature."); } ``` ```python # 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`) # Broadcasting This document describes the broadcasting semantics of XLA. ## What is broadcasting? Broadcasting is the process of making arrays with different shapes have compatible shapes for arithmetic operations. The terminology is borrowed from [NumPy broadcasting](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html). 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](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html). ## Principles 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 *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](https://github.com/openxla/xla/blob/main/xla/hlo/builder/xla_builder.cc). In the XLA source code, this type of broadcasting is sometimes called "InDim" broadcasting. ### Formal definition 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 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](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html). ## Broadcast composition 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`) # 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](developer_guide.md). ## Linux ### Configure 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 We recommend using a suitable Docker image - such as [ml-build](https://us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/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. ```sh 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: ```sh 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: ```sh ./configure.py --backend=CPU bazel build \ --spawn_strategy=sandboxed \ --test_output=all \ //xla/... ``` ### GPU support 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: ```sh 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: ```sh 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.](https://www.tensorflow.org/install/source#gpu_support_2) You can build XLA targets with GPU support without Docker as well. Configure and build targets using the following commands: ```sh ./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.](https://github.com/google-ml-infra/rules_ml_toolchain/blob/main/gpu) ### Build XLA with CUDA/cuDNN Support Using the JAX CI/Release Container XLA is a compiler used internally by JAX. JAX is distributed via PyPI wheels. The [JAX Continuous Integration documentation](https://github.com/jax-ml/jax/tree/main/ci#running-these-scripts-locally-on-your-machine) explains how to build JAX wheels using the [tensorflow/ml-build:latest](https://us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build) 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 ```bash git clone https://github.com/jax-ml/jax.git cd jax ``` 2. Start JAX CI/Release Docker container by running: ```bash ./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: ```bash 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: ```bash 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.: ```bash /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.: ```bash --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 ```bash cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \ ./dist/ ``` 7. Exit the interactive shell: ```bash exit ``` ## Windows 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 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 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: ```bash # 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`) # 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 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: ```bash # 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 Command buffers are already enabled for several operation types. You do not need an extra flag to enable the default behavior. | Setting | Default | What it controls | | --- | --- | --- | | `xla_gpu_enable_command_buffer` | See the command-type table below | Which operations may be captured. | | `xla_gpu_enable_collectives_command_buffer_filter` | `ALLCOLLECTIVES` | Which supported library collective types may be captured when `COLLECTIVES` is enabled. | | `xla_gpu_graph_min_graph_size` | `5` | Minimum eligible region size for capture. | | `xla_gpu_command_buffer_scheduling_mode` | `LHS` | How operations inside a graph may overlap. | | `xla_gpu_command_buffer_update_mode` | `ALWAYS_UPDATE` | How buffer addresses are managed to reduce graph updates. | | `xla_cmd_buffer_trace_cache_size` | `16` | Number of cached captures per traced-command cache. | | `xla_gpu_command_buffer_unroll_loops` | `false` | Whether eligible loops with known trip counts are unrolled during graph lowering. | | `xla_enable_command_buffers_during_profiling` | `false` | Whether graph execution remains enabled during an XLA-detected profiling session. | ## Choose which operations to capture `--xla_gpu_enable_command_buffer` accepts a list of operation categories, not a boolean. It has three useful forms: ```bash # 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. | Category | Enabled by default? | Operations it allows | | --- | --- | --- | | `FUSION` | Yes | GPU kernels, including fusion/custom kernels, and supported copies within one device. | | `CUBLAS` | Yes | Operations using the GEMM thunk path. Retained as a category even when a build lowers most matrix multiplications through other paths. | | `CUBLASLT` | Yes | Matrix multiplication through the cuBLASLt thunk path. | | `CUDNN` | Yes | cuDNN thunk operations, such as supported library fusions. | | `CUSTOM_CALL` | Yes | Compatible custom calls. An FFI handler must advertise command buffer compatibility. | | `DYNAMIC_SLICE_FUSION` | Yes | Supported fusions combining dynamic slicing with other operations. | | `CONDITIONAL` | Yes | Supported conditional branches. | | `COLLECTIVES_KERNEL` | Yes | Collective-kernel thunks, such as eligible one-shot collective kernels. Separate from library collectives. | | `COLLECTIVES` | No | Supported library collective operations, groups, and send/receive paths. | | `WHILE` | No | Supported while loops whose condition and body can be captured. | | `CONVOLUTION` | No | The convolution thunk path, including supported ROCm convolutions. | | `HOST_EXECUTE` | No | Supported host-execution start/done operations; support is still initial. | 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 To capture only selected library collective types, enable `COLLECTIVES` and set `--xla_gpu_enable_collectives_command_buffer_filter`: ```bash 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: ```bash XLA_FLAGS="--xla_gpu_enable_command_buffer=-COLLECTIVES,-COLLECTIVES_KERNEL" \ python train.py ``` ## Adjust the minimum graph size `--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: ```bash 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 `--xla_gpu_command_buffer_scheduling_mode` controls concurrency inside command buffers. Start with `LHS`, then compare one alternative at a time. | Value | Meaning | When to consider it | | --- | --- | --- | | `LHS` | Uses the latency-hiding schedule's overlap decisions. | Default starting point, especially for workloads that overlap communication with computation. | | `SERIALIZE` | Serializes commands within each command buffer. | Diagnose whether concurrency affects performance or execution behavior. It does not serialize the entire application. | | `CONCURRENT_REGIONS` | Allows concurrency in regions of small, latency-bound kernels while serializing larger kernels. Also changes buffer assignment to limit reuse within concurrent regions. | Workloads with independent small kernels that individually underutilize the GPU. Peak memory can increase compared with `LHS` or `SERIALIZE`. | | `CONCURRENT` | Allows independent operations to overlap and changes buffer assignment to support that overlap. | An aggressive experiment when the workload has available GPU capacity and memory headroom. Memory use can increase enough to cause an out-of-memory error. | ```bash 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 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. | Value | Behavior | When to consider it | | --- | --- | --- | | `ALWAYS_UPDATE` | Uses normal allocation addresses and updates commands when their addresses change. Despite the name, it does not unconditionally rebuild the graph on every execution. | Default baseline. | | `SKIP_TEMP` | Gives command-buffer-referenced preallocated temporary buffers fixed VMM addresses. Other changing allocations can still require updates. | First experiment for reducing update overhead from temporary storage. | | `SKIP_PROFILED` | Includes those temporary buffers and observes other eligible allocations during initial executions. Selected stable candidates then use fixed VMM addresses. | Repeated workloads where input/output address patterns are stable and graph-update overhead is significant. | ```bash 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 `--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. ```bash 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 `--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: ```bash 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 `--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: ```bash 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 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 | Observation | What to check | | --- | --- | | No speedup | The workload may already be dominated by long kernels, or capture/update costs may offset launch savings. Compare warmed-up, synchronized measurements. | | Eligible operations remain outside graphs | Check command categories, graph-size threshold, dependencies, FFI compatibility, and GPU runtime support. An unsupported operation can split a region. | | Behavior changes under profiling | Check `xla_enable_command_buffers_during_profiling` and whether the profiler activates XLA's fallback. | | Higher memory use | Compare against `LHS`, the default trace-cache capacity, loop unrolling disabled, and `ALWAYS_UPDATE`, one change at a time. | | Unknown flag or enum value | Check the XLA version bundled with your framework. A newer upstream flag may not yet be in a release. | 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](https://github.com/openxla/xla/blob/main/xla/debug_options_flags.cc), [option enums](https://github.com/openxla/xla/blob/main/xla/xla.proto), [capture eligibility checks](https://github.com/openxla/xla/blob/main/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc), and [VMM allocation policy](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_executable_va_remap_allocator.cc). --- ### Contributing (`docs/contributing.md`) # 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](https://opensource.google/conduct/). ## Before you begin ### Sign the Contributor License Agreement Contributions to this project must be accompanied by a [Contributor License Agreement](https://cla.developers.google.com/about) (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 to see your current agreements or to sign a new one. ### Review the Code of Conduct This project follows [Tensorflow's Code of Conduct](https://github.com/tensorflow/tensorflow/blob/master/CODE_OF_CONDUCT.md). ## Contribution process ### Developer Guide 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](./developer_guide.md). ### Contribution guide 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](https://github.com/openxla/xla/blob/main/xla/backends/autotuner/codegen_backend.h) 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>> GetSupportedConfigs(const HloInstruction& instr); // Returns a default config for the given HLO instruction. virtual absl::StatusOr> 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](gpu_aot_compatibility.md). ### Code standards * *Coding style*: We follow [Google's code style guide](https://google.github.io/styleguide/). Specifically see the [C/C++](https://google.github.io/styleguide/cppguide.html) and [Python](https://google.github.io/styleguide/pyguide.html) guides. All code submitted must strictly conform to these style guides. * *Compact changes*: We follow [Google's engineering practices](https://google.github.io/eng-practices/). In particular, please observe the [guide on writing compact changes](https://google.github.io/eng-practices/review/developer/small-cls.html). 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 `#ifdef`s in XLA](./avoiding_ifdefs.md) 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 All submissions, including submissions by project members, require review. We use GitHub pull requests for this purpose. Consult [GitHub Help](https://help.github.com/articles/about-pull-requests/) 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 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: > > ``` > 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) ### "This infrastructure change is not related to my PR. Why should I do it?" 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?" 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? 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 (**tGoogle 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](https://github.com/Tixxx) * [@sergachev](https://github.com/sergachev) * [@jreiffers](https://github.com/jreiffers) * [@ezhulenev](https://github.com/ezhulenev) * [@mminutoli](https://github.com/mminutoli) * [@pemeliya](https://github.com/pemeliya) | Percentile | Target for **tGoogle Review** (co-maintainer approval to PR merged) | | --- | --- | | 50 percentile of PRs | 3 business days | | 80 percentile of PRs | 5 business days | 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`) # 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](https://github.com/google/copybara). ## Internal source of truth 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 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 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`) # 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 See [JAX documentation](https://jax.readthedocs.io/en/latest/ffi.html) for end to end examples of integrating custom calls and XLA FFI with JAX. ## XLA FFI Binding 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 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. ```c++ // 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 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` 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. ```mlir %0 = "stablehlo.custom_call"(%arg0) { call_target_name = "foo", api_version = 4 : i32 } : (tensor<2x2xf32>) -> tensor<2x2xf32> ``` ```c++ // Buffers of any number of dimensions and data type. auto handler = Ffi::Bind().Arg().Ret().To( [](AnyBuffer arg, Result res) -> Error { void* arg_data = arg.untyped_data(); void* res_data = res->untyped_data(); return Error::Success(); }); ``` ### Constrained Buffer Arguments And Results `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. ```c++ // Buffers of any number of dimensions and F32 data type. auto handler = Ffi::Bind().Arg>().Ret>().To( [](Buffer arg, Result> res) -> Error { float* arg_data = arg.typed_data(); float* res_data = res->typed_data(); return Error::Success(); }); ``` ```c++ // Buffers of number of dimensions 2 and F32 data type. auto handler = Ffi::Bind().Arg>().Ret>().To( [](BufferR2 arg, Result> res) -> Error { float* arg_data = arg.typed_data(); float* res_data = res->typed_data(); return Error::Success(); }); ``` ### Matching And Verifying Buffers 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: ```c++ namespace m = ::xla::ffi::match; m::Buffer(); m::Buffer().WithDType().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. ```c++ namespace m = ::xla::ffi::match; auto handler = Ffi::Bind().Arg().Ret().To( [](AnyBuffer input, Result output) -> Error { int64_t rows; int64_t cols; ErrorOr> typed_input = Match("input", input, m::Buffer().WithDims(&rows, &cols)); if (typed_input.has_error()) { return std::move(typed_input).error(); } ErrorOr>> typed_output = Match("output", output, m::Buffer().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: ```c++ 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().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: ```c++ ErrorOr> MatchActivations(AnyBuffer input) { return Match( "input", input, m::Buffer().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: ```c++ auto matrix = m::Buffer().WithDims(&rows, 128); int64_t n; ErrorOr> square = Match("matrix", buffer, m::Buffer().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: ```c++ Error VerifySortBuffers(AnyBuffer keys, AnyBuffer values, Result 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: ```c++ namespace m = ::xla::ffi::match; Error VerifyIndices(AnyBuffer indices) { auto pattern = m::Buffer().WithDType().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: ```c++ Error VerifySameShape(BufferR2 lhs, BufferR2 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` 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`: ```c++ ABSL_ASSIGN_OR_RETURN(BufferR2 input, Match("input", buffer, m::Buffer())); ABSL_RETURN_IF_ERROR(Verify( "input", input, m::Buffer().WithDims(expected_rows, m::Dim()))); ``` ### Variadic Arguments And Results 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 arg = args.get(0); ErrorOr> res = results.get(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. ```c++ auto handler = Ffi::Bind() .Arg() .RemainingArgs() .Ret() .RemainingRets() .To([](AnyBuffer arg, RemainingArgs args, AnyBuffer ret, RemainingRets results) -> Error { return Error::Success(); }); ``` ### Attributes XLA FFI supports automatic decoding of `mlir::DictionaryAttr` passed as a `custom_call` `backend_config` into FFI handler arguments. Note: See [stablehlo RFC](https://github.com/openxla/stablehlo/blob/main/rfcs/20240312-standardize-customcallop.md) for details, and `stablehlo.custom_call` operation specification. ```mlir %0 = "stablehlo.custom_call"(%arg0) { call_target_name = "foo", backend_config= { i32 = 42 : i32, str = "string" }, api_version = 4 : i32 } : (tensor) -> tensor ``` 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. ```c++ auto handler = Ffi::Bind() .Arg>() .Attr("i32") .Attr("str") .To([](BufferR0 buffer, int32_t i32, std::string_view str) { return Error::Success(); }); ``` ### User-Defined Enum Attributes 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. ```mlir %0 = "stablehlo.custom_call"(%arg0) { call_target_name = "foo", backend_config= { command = 0 : i32 }, api_version = 4 : i32 } : (tensor) -> tensor ``` ```c++ enum class Command : int32_t { kAdd = 0, kMul = 1, }; XLA_FFI_REGISTER_ENUM_ATTR_DECODING(Command); auto handler = Ffi::Bind().Attr("command").To( [](Command command) -> Error { return Error::Success(); }); ``` ### Binding All Custom Call Attributes 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. ```c++ auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error { ErrorOr i32 = attrs.get("i32"); return Error::Success(); }); ``` ### Per-execution Custom Options FFI handlers can access runtime custom options as a `Dictionary` by binding `Ctx()`. In JAX, pass options with `jax.execution_options`: ```python with jax.execution_options(custom_options={"foo": 42}): result = compiled_fn(x) ``` Python integer options decode as `int64_t`: ```c++ auto handler = Ffi::Bind().Ctx().To([](Dictionary options) -> Error { ErrorOr foo = options.get("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 XLA FFI can decode dictionary attributes into user-defined structs. ```mlir %0 = "stablehlo.custom_call"(%arg0) { call_target_name = "foo", backend_config= { range = { lo = 0 : i64, hi = 42 : i64 } }, api_version = 4 : i32 } : (tensor) -> tensor ``` 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). ```c++ struct Range { int64_t lo; int64_t hi; }; XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(Range, StructMember("lo"), StructMember("hi")); auto handler = Ffi::Bind().Attr("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. ```c++ auto handler = Ffi::Bind().Attrs().To([](Dictionary attrs) -> Error { ErrorOr range = attrs.get("range"); return Error::Success(); }); ``` ## Create a custom call on CPU 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 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. ```c++ 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 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<<>>( in0.data, in1.data, out->data); } XLA_FFI_DEFINE_HANDLER(handler, do_custom_call, ffi::Ffi::Bind() .Ctx>() .Arg() .Arg() .Ret()); 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 Consider the following custom call. ```c++ 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 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`) # 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) ### Cause 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 * **Persisting & Caching Autotuning Decisions**: Run autotuning once and save the choices to a persistent cache (see [persisted autotuning](./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) ### Cause 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 * **`--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 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`) # Setting Up Developer Environment ## Setting up LSP with clangd ### Background Editors such as Emacs, Vim, or VS Code support features like code navigation, code completion, inline compiler error messages, and others, through [LSP](https://en.wikipedia.org/wiki/Language_Server_Protocol), the Language Server Protocol. A common language server with LSP support is [clangd](https://clangd.llvm.org), 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? Use the [build_tools/lint/generate_compile_commands.py](https://github.com/openxla/xla/blob/main/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 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 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 Install [Buildozer](https://github.com/bazelbuild/buildtools/blob/main/buildozer/README.md) 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](https://github.com/hzeller/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`) # 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](contributing.md) and review the contribution process. 2. If you haven't already done so, sign the [Contributor License Agreement](https://cla.developers.google.com/about). 3. Install or configure the following dependencies: - A [GitHub](https://github.com/) account - [Docker](https://www.docker.com/) 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 1. Create a fork of the [XLA repository](https://github.com/openxla/xla). 2. Clone your fork of the repo, replacing `{USER}` with your GitHub username: ```sh git clone https://github.com/{USER}/xla.git ``` 3. Change into the `xla` directory: `cd xla` 4. Configure the remote upstream repo: ```sh git remote add upstream https://github.com/openxla/xla.git ``` ## Set up an environment 1. Install [Bazel](https://bazel.build/install). To build XLA, you must have Bazel installed. The recommended way to install Bazel is using [Bazelisk](https://github.com/bazelbuild/bazelisk#readme), which automatically downloads the correct Bazel version for XLA. If Bazelisk is unavailable, you can [install Bazel](https://bazel.build/install) manually. 2. Create and run the [ml-build](https://us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/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 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 Configure for CPU: ```sh docker exec xla ./configure.py --backend=CPU ``` Configure for GPU: ```sh 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: ```sh # 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: ```sh 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. ```sh # 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](build_from_source.md). ## Create a pull request When you're ready to send changes for review, create a [pull request](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/about-pull-requests). To learn about the XLA code review philosophy, see [Review Process](contributing.md#review-process). ## Static Analysis (Clang-Tidy) To maintain code quality, XLA uses `clang-tidy` for static analysis and include verification. ### How to Run There are two ways to execute checks. Running it against specific targets can be done with: ```sh 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. ```sh # 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. ```sh bazel run //build_tools/ci:run_clang_tidy -- --fix ``` --- ### Developing New Backend (`docs/developing_new_backend.md`) # 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](http://llvm.org), [Bazel](https://bazel.build/), 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](http://llvm.org) 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 In this scenario, start by looking at the existing [XLA CPU backend](https://github.com/openxla/xla/tree/main/xla/service/cpu). 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`](https://github.com/openxla/xla/tree/main/xla/service/compiler.h) 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 It is possible to model a new [`xla::Compiler`](https://github.com/openxla/xla/tree/main/xla/service/compiler.h) implementation on the existing [`xla::CPUCompiler`](https://github.com/openxla/xla/tree/main/xla/service/cpu/cpu_compiler.cc) and [`xla::GPUCompiler`](https://github.com/openxla/xla/tree/main/xla/service/gpu/nvptx_compiler.cc) 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](https://github.com/openxla/xla/tree/main/xla/service/gpu/) 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 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`](https://github.com/openxla/xla/tree/main/xla/stream_executor/stream_executor.h): For many devices not all methods of `StreamExecutor` are needed. See existing `StreamExecutor` implementations for details. * [`xla::Compiler`](https://github.com/openxla/xla/tree/main/xla/service/compiler.h): This class encapsulates the compilation of an HLO computation into an `xla::Executable`. * [`xla::Executable`](https://github.com/openxla/xla/tree/main/xla/service/executable.h): This class is used to launch a compiled computation on the platform. * [`xla::TransferManager`](https://github.com/openxla/xla/tree/main/xla/service/transfer_manager.h): 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`) # 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 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](https://github.com/openxla/xla/blob/f4d624b6811c28925c3006f5b779f1149b3b39ac/xla/pjrt/proto/compile_options.proto#L71) field of the ExecutableBuildOptionsProto message, or the [optimization_level](https://github.com/openxla/xla/blob/f4d624b6811c28925c3006f5b779f1149b3b39ac/xla/xla.proto#L1580) 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: | Level | Use Case | | :-------- | :---------------------------------------------------------------------- | | EFFORT_O0 | Fastest compilation, slowest runtime | | EFFORT_O1 | Faster compilation with reasonable runtime | | EFFORT_O2 | Strongly prioritize runtime (suitable default for production workloads) | | EFFORT_O3 | Expensive or experimental optimizations | ### Use in XLA:GPU 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](flags_guidance.md#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](https://github.com/openxla/xla/blob/5b54d0e9cf34f4e5ab05b3752ecb390145ca5716/xla/service/gpu/transforms/collectives/all_gather_combiner.cc#L78) * [all_reduce_combiner](https://github.com/openxla/xla/blob/5b54d0e9cf34f4e5ab05b3752ecb390145ca5716/xla/service/gpu/transforms/collectives/all_reduce_combiner.cc#L76) * [reduce_scatter_combiner](https://github.com/openxla/xla/blob/5b54d0e9cf34f4e5ab05b3752ecb390145ca5716/xla/service/gpu/transforms/collectives/reduce_scatter_combiner.cc#L76) ## Memory Fitting Level 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](https://github.com/openxla/xla/blob/5b54d0e9cf34f4e5ab05b3752ecb390145ca5716/xla/pjrt/proto/compile_options.proto#L79) field of the ExecutableBuildOptionsProto message, or the [memory_fitting_level](https://github.com/openxla/xla/blob/f4d624b6811c28925c3006f5b779f1149b3b39ac/xla/xla.proto#L1588) 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: | Level | Use Case | | :-------- | :---------------------------------------------------------------------- | | EFFORT_O0 | Minimal effort to fit (fail compilation as quickly as possible instead) | | EFFORT_O1 | Reduced effort to fit | | EFFORT_O2 | Significant effort to fit (suitable default for production workloads) | | EFFORT_O3 | Expensive or experimental algorithms to reduce memory usage | --- ### Emitters (`docs/emitters.md`) # 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](https://developer.nvidia.com/cublas), [cuDNN](https://developer.nvidia.com/cudnn). 2. Tiling HLO to block-level and then using [OpenAI Triton](https://openai.com/index/triton/). 3. Using XLA Emitters to progressively lower HLO to LLVM IR. This document is focused on XLA:GPU Emitters. ## Hero-based codegen 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 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 See [computation_partitioner.h](https://github.com/openxla/xla/blob/ca62f3e1bc9ea1d808c3a4de0a78bae7453389eb/xla/codegen/emitters/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 See [elemental_hlo_to_mlir.h](https://github.com/openxla/xla/blob/ca62f3e1bc9ea1d808c3a4de0a78bae7453389eb/xla/codegen/emitters/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 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](https://openxla.org/xla/indexing), 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 We do not support internal `tuple`s. We also do not support nested tuple outputs. All XLA graphs that use these features can be converted to graphs that do not. ### Gather We only support canonical gathers as produced by [`gather_simplifier`]( https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/gather_simplifier.h). ## Subgraph functions 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 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 `%pn`s are the parameters of the computation, and the `%rn`s are the results of the computation. The entry computation takes the results as tensors, `tensor.insert`s updates into them, and then returns them. No other uses of the output tensors are allowed. ## Compilation pipeline ### Loop emitter See [loop.h](https://github.com/openxla/xla/blob/cfd16b7f21feff17635c782f4489c0f478178eb9/xla/backends/gpu/codegen/emitters/loop.h#L4). 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](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/ir/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](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/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](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/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](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/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](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/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](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/lower_tensors.cc). `tensor.extract` is lowered to `llvm.load`, `tensor.insert` to `llvm.store`, in the obvious way. - [propagate_slice_indices](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/propagate_slice_indices.cc) and [merge_pointers_to_same_slice](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms/merge_pointers_to_same_slice.cc) 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 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 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= ``` 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 * Compilation pipeline: [emitter_base.h](https://github.com/openxla/xla/blob/cfd16b7f21feff17635c782f4489c0f478178eb9/xla/backends/gpu/codegen/emitters/emitter_base.h) * Optimization and conversion passes: [backends/gpu/codegen/emitters/transforms](https://github.com/openxla/xla/tree/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/transforms) * Partition logic: [computation_partitioner.h](https://github.com/openxla/xla/blob/ca62f3e1bc9ea1d808c3a4de0a78bae7453389eb/xla/codegen/emitters/computation_partitioner.h) * Hero-based emitters: [backends/gpu/codegen/emitters](https://github.com/openxla/xla/tree/cfd16b7f21feff17635c782f4489c0f478178eb9/xla/backends/gpu/codegen/emitters) * XLA:GPU ops: [xla_gpu_ops.td](https://github.com/openxla/xla/blob/39fc9c0c3e06b29e27f2aebdb21fca15754de928/xla/backends/gpu/codegen/emitters/ir/xla_gpu_types.td) * Correctness and lit tests: [backends/gpu/codegen/emitters/tests](https://github.com/openxla/xla/tree/30229c9836cafa6b05c6d42f0d918e5f8dc0b2dd/xla/backends/gpu/codegen/emitters/tests) --- ### Error Codes (`docs/error_codes.md`) # XLA Error codes This page is a list of all error codes emitted by the XLA compiler. - [E0100](./errors/error_0100.md) - Runtime: Buffer Allocation Failure - [E0101](./errors/error_0101.md) - Runtime: Program Allocation Failure - [E0102](./errors/error_0102.md) - Runtime: Program Input Buffer Mismatch - [E0200](./errors/error_0200.md) - Runtime: Core Halted Unexpectedly - [E1000](./errors/error_1000.md) - Compile Time: HBM OOM - [E1001](./errors/error_1001.md) - Compile Time: Scoped Vmem OOM - [E1200](./errors/error_1200.md) - Compile Time: Host Offload Output Mismatch - [E2001](./errors/error_2001.md) - Compile Time: Unsupported RHS DataType on Hardware - [E2002](./errors/error_2002.md) - Compile Time: Mosaic Input/Output Misaligned Block and Tiling - [E2003](./errors/error_2003.md) - Compile Time: Mosaic Unproven Memory Access Alignment - [E3000](./errors/error_3000.md) - Compile Time: SparseCore Allocation Failure - [E3001](./errors/error_3001.md) - Compile Time: SparseCore No Viable Logical Replica Count --- ### Flags Guidance (`docs/flags_guidance.md`) # 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 Flag | Description | Default Values | Suggested Values | Candidate Values :------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :----------------------------------- | :----------------------------------- | :--------------- `xla_mosaic_on_device_checks` | This flag enables on-device checks for Mosaic codegen. Currently, the supported checks are on bounds, i.e., if an out-of-bounds memory is touched, the compilation/execution would catch it. | `xla_mosaic_on_device_checks=bounds` | `xla_mosaic_on_device_checks=bounds` | `xla_mosaic_on_device_checks=bounds` ## Performance Flags The following flags are instrumental in enhancing runtime performance. Experimenting with these settings may lead to considerable performance gains. Flag | Description | Default Values | Suggested Values | Candidate Values :----------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :--------------- **Pipelining**
1. `xla_should_allow_loop_variant_parameter_in_chain`
2. `xla_should_add_loop_invariant_op_in_chain`
3. `xla_tpu_enable_ici_ag_pipelining` | These 3 flags should be used in conjunction to enable collective pipelining of ICI(Interchip-Interconnect) all-gather operations, which creates more opportunities for overlapping execution. | 1. `xla_should_allow_loop_variant_parameter_in_chain=kDisabled`
2. `xla_should_add_loop_invariant_op_in_chain=kDisabled`
3. `xla_tpu_enable_ici_ag_pipelining=false` | 1. `xla_should_allow_loop_variant_parameter_in_chain=kEnabled`
2. `xla_should_add_loop_invariant_op_in_chain=kEnabled`
3. `xla_tpu_enable_ici_ag_pipelining=true` | 1. `xla_should_allow_loop_variant_parameter_in_chain=kDisabled/kEnabled/kAuto`
2. `xla_should_add_loop_invariant_op_in_chain=kDisabled/kEnabled/kAuto`
3. `xla_tpu_enable_ici_ag_pipelining=true/false` **v5e/Async**
`xla_enable_async_all_gather`
`xla_tpu_enable_async_collective_fusion`
`xla_tpu_enable_async_collective_fusion_fuse_all_gather` | These 3 flags should be used in conjunction to activate asynchronous all-gather operations on v5e. | `xla_enable_async_all_gather=kAuto`
`xla_tpu_enable_async_collective_fusion=true`
`xla_tpu_enable_async_collective_fusion_fuse_all_gather=true` | `xla_enable_async_all_gather=kAuto`
`xla_tpu_enable_async_collective_fusion=true`
`xla_tpu_enable_async_collective_fusion_fuse_all_gather=true` | `xla_enable_async_all_gather=kDisabled/kEnabled/kAuto`
`xla_tpu_enable_async_collective_fusion=true/false`
`xla_tpu_enable_async_collective_fusion_fuse_all_gather=true/false` **v5e/Async**
`xla_tpu_enable_async_collective_fusion`
`xla_tpu_enable_async_collective_fusion_fuse_all_reduce` | These 2 flags should be used in conjunction to activate asynchronous all-reduce operations on v5e. | `xla_tpu_enable_async_collective_fusion=true`
`xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false` | `xla_tpu_enable_async_collective_fusion=true`
`xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true` | `xla_tpu_enable_async_collective_fusion=true/false`
`xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true/false` **Async**
`xla_tpu_enable_async_all_to_all` | This flag enables asynchronous all-to-all communication. | `xla_tpu_enable_async_all_to_all=false` | `xla_tpu_enable_async_all_to_all=true` | `xla_tpu_enable_async_all_to_all=true/false` **Latency-bound**
`xla_all_gather_latency_bound_threshold_in_bytes` | This flag is intended for latency-bound (i.e., small-sized) all-gather operations. Enabling this triggers specific optimizations that can reduce execution time for latency-bound all-gathers. Typically it’s used in inference workloads. | `xla_all_gather_latency_bound_threshold_in_bytes=-1`
(which is not enabled) | `4~16Mb(i.e. 4~16 * 1024 * 1024)` | `[0, 9223372036854775807]` **Latency-bound**
`xla_all_reduce_latency_bound_threshold_in_bytes` | This flag is intended for latency-bound (i.e., small-sized) all-gather operations. Enabling this triggers specific optimizations that can reduce execution time for latency-bound all-reduces. Typically it’s used in inference workloads. | `xla_all_reduce_latency_bound_threshold_in_bytes=-1`
(which is not enabled) | `4~16Mb(i.e. 4~16 * 1024 * 1024)` | `[0, 9223372036854775807]` **Latency-bound**
`xla_collective_permute_latency_bound_threshold_in_bytes` | This flag is intended for latency-bound (i.e., small-sized) all-gather operations. Enabling this triggers specific optimizations that can reduce execution time for latency-bound collective-permutes. Typically it’s used in inference workloads. | `xla_collective_permute_latency_bound_threshold_in_bytes=-1`
(which is not enabled) | `4~16Mb(i.e. 4~16 * 1024 * 1024)` | `[0, 9223372036854775807]` **Latency-bound**
`xla_all_to_all_latency_bound_threshold_in_bytes` | This flag is intended for latency-bound (i.e., small-sized) all-gather operations. Enabling this triggers specific optimizations that can reduce execution time for latency-bound all-to-all. Typically it’s used in inference workloads. | `xla_all_to_all_latency_bound_threshold_in_bytes=-1`
(which is not enabled) | `4~16Mb(i.e. 4~16 * 1024 * 1024)` | `[0, 9223372036854775807]` `xla_enable_async_collective_permute` | Rewrites all collective-permute operations to their asynchronous variants. When set to `auto`, XLA can turn on async collective based on other configurations or conditions automatically. | `xla_enable_async_collective_permute=kAuto` | `xla_enable_async_collective_permute=kAuto` | `xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled` **Compute centric**
`xla_tpu_enable_dot_strength_reduction` | This flag rewrites non-compute intensive dots as multiply + reduce operations. | **Compute centric**
`xla_tpu_enable_dot_strength_reduction=true` | `xla_tpu_enable_dot_strength_reduction=true` | `xla_tpu_enable_dot_strength_reduction=true/false` **Compute centric**
`xla_tpu_dot_dot_fusion` | This flag enables dot-dot fusion, which fuses a producer-dot operation with a consumer-dot operation. On doing so, the producer-dot's output is not manifested in slow/main memory driving down memory footprint. | `xla_tpu_dot_dot_fusion=true` | `xla_tpu_dot_dot_fusion=true` | `xla_tpu_dot_dot_fusion=true/false` **Compute centric**
`xla_jf_enable_multi_output_fusion` | This flag enables fusions that fuse multiple consumers (i.e. the resultant fusion will have multiple outputs) | `xla_jf_enable_multi_output_fusion=true` | `xla_jf_enable_multi_output_fusion=true` | `xla_jf_enable_multi_output_fusion=true/false` **Compute centric**
`xla_tpu_scoped_vmem_limit_kib` | This flag sets the amount of scratchpad VMEM available to per op for local usage in KiloBytes. Rest of the VMEM is used as buffer space. | `xla_tpu_scoped_vmem_limit_kib=16384` | `xla_tpu_scoped_vmem_limit_kib=16384` | `xla_tpu_scoped_vmem_limit_kib=[4096, VMEM size of the architecture - 1024]` **Compute centric**
`xla_tpu_async_copy_bandwidth_scaling_factor` | Scales effective bandwidth for async copies. This is used when making prefetch decisions and deciding which tensors should live in VMEM. | `xla_tpu_async_copy_bandwidth_scaling_factor=1` | `xla_tpu_async_copy_bandwidth_scaling_factor=1` | `xla_tpu_async_copy_bandwidth_scaling_factor=(0, 1]` **Compute centric**
`xla_msa_enable_cross_program_prefetch_freeing` | Enables freeing optimization for cross-program-prefetched buffers. | `xla_msa_enable_cross_program_prefetch_freeing=enabled` | `xla_msa_enable_cross_program_prefetch_freeing=enabled` | `xla_msa_enable_cross_program_prefetch_freeing=enabled/disabled` **Compute centric**
`xla_tpu_msa_inefficient_use_to_copy_ratio` | The ratio of use bytes to copy bytes for a given allocation site below which we consider the site to be inefficient. This is used while making VMEM placement decisions. A value of 0 would treat all sites as efficient and a value of 1 would require the amount of bytes used at the site to be at least as much as the async copy bytes. | `xla_tpu_msa_inefficient_use_to_copy_ratio=0.5` | `xla_tpu_msa_inefficient_use_to_copy_ratio=0.5` | `xla_tpu_msa_inefficient_use_to_copy_ratio=[0, 1]` ### CPU Performance Flags Flag | Description | Default Values | Suggested Values | Candidate Values :--- | :--- | :--- | :--- | :--- `xla_cpu_opt_preset` | Sets the CPU optimization preset. `FAST_COMPILE` packages up a series of compiler tweaks that trade off small amounts of runtime performance for large returns in compile time. `FAST_RUNTIME` is the default and doesn't need to be specified. | `FAST_RUNTIME` | `FAST_COMPILE` (for development) | `FAST_RUNTIME`, `FAST_COMPILE` ## Memory Flags 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. Flag | Description | Default Values | Suggested Values | Candidate Values :------------------------------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | :----------------------------------------------------- | :------------------------------------------------------ | :--------------- **Scheduler**
`xla_latency_hiding_scheduler_rerun` | This setting adjusts the behavior of the latency-hiding scheduler. It works by incrementally reducing the memory limit allocated for scheduling with each "rerun" of the process. | `xla_latency_hiding_scheduler_rerun=1` | `xla_latency_hiding_scheduler_rerun=5` | `0~10(it doesn’t make much sense beyond 10 reruns)` **Fusion**
`xla_tpu_rwb_fusion` | This flag enables reduce+broadcast type of fusions, and may decrease memory usage. | `xla_tpu_rwb_fusion=true` | `xla_tpu_rwb_fusion=false` | `xla_tpu_rwb_fusion=true/false` **Scheduler**
`xla_memory_scheduler` | This flag specifies the algorithm the memory scheduler will use to minimize memory consumption. Using a more advanced algorithm might get a less memory-consuming schedule, at the cost of longer compilation time. | `xla_memory_scheduler=kDefault` | `xla_memory_scheduler=kBrkga` | `xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga` **Scheduler**
`xla_tpu_enable_latency_hiding_scheduler` | This flag enables the latency-hiding scheduler, which allows us to perform asynchronous collective instead of synchronous ones. Disabling it reduces memory usage at the cost of losing the performance gains from these asynchronous operations. | `xla_tpu_enable_latency_hiding_scheduler=true` | `xla_tpu_enable_latency_hiding_scheduler=false` | `xla_tpu_enable_latency_hiding_scheduler=true/false` **SPMD**
`xla_jf_spmd_threshold_for_windowed_einsum_mib` | This flag sets the lower threshold of the minimum size of the dot to trigger collective matmul. Setting it to a higher value would save memory at the cost of losing opportunities to perform collective matmul. | `xla_jf_spmd_threshold_for_windowed_einsum_mib=-1` | `10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024)` | `[0, 9223372036854775807]` **Scheduler**
`xla_gpu_enable_analytical_sol_latency_estimator` | This flag enables the analytical estimator which maximizes compute-communication overlap on GPUs. | `xla_gpu_enable_analytical_sol_latency_estimator=true` | `xla_gpu_enable_analytical_sol_latency_estimator=false` | `true/false` ## Other commonly used flags | Flag | Type | Notes | | :---- | :---- | :----- | | `xla_dump_to` | String (filepath) | The folder where pre-optimization HLO files and other artifacts will be placed (see [XLA Tools](https://openxla.org/xla/tools)). | ### TPU XLA flags Flag | Type | Notes :---------------------------------------------- | :------------------- | :---- `xla_tpu_enable_data_parallel_all_reduce_opt` | Boolean (true/false) | Optimization to increase overlap opportunities for DCN (data center networking) all-reduces used for data parallel sharding. `xla_tpu_data_parallel_opt_different_sized_ops` | Boolean (true/false) | Enables pipelining of data parallel ops across multiple iterations even if their output sizes don't match what can be saved in place in the stacked variables. Can increase memory pressure. `xla_tpu_spmd_rng_bit_generator_unsafe` | Boolean (true/false) | Whether to run RngBitGenerator HLO in a partitioned way, which is unsafe if deterministic results are expected with different shardings on different parts of the computation. `xla_tpu_megacore_fusion_allow_ags` | Boolean (true/false) | Allows fusing all-gathers with convolutions/all-reduces. `xla_tpu_enable_ag_backward_pipelining` | Boolean (true/false) | Pipelines all-gathers (currently megascale all-gathers) backwards through scan loops. ### GPU XLA flags 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](https://openxla.org/xla/effort_levels) for details. Flag | Type | Notes :------------------------------------------------ | :----------------------------- | :---- `xla_gpu_enable_latency_hiding_scheduler` | Boolean (true/false) | This flag enables latency hiding schedulers to overlap asynchronous communication with computation efficiently. The default value is False. `xla_gpu_enable_analytical_sol_latency_estimator` | Boolean (true/false) | Enables platform specific scheduling decisions, which in turn improve compute-communication overlap. The default value is true. `xla_gpu_analytical_latency_estimator_options` | Structured string | Configures parameters for the `xla_gpu_enable_analytical_sol_latency_estimator`. Adjust by setting `nic_speed_gbps=$NIC_SPEED,nccl_op_launch_us=$LAUNCH_OVERHEAD,chunk_prep_us=$CHUNK_PREP,rtt_us=$RTT,chunk_size_bytes=$CHUNK_SIZE,gpus_per_node=$GPUS_PER_NODE`. The default value depends on a detected platform. `xla_gpu_enable_triton_gemm` | Boolean (true/false) | Use Triton-based matrix multiplication. `xla_gpu_enable_command_buffer` | List of CommandBufferCmdType | Which kind of commands should be captured in command buffers. See [Tuning GPU command buffers](https://openxla.org/xla/command_buffer_tutorial) for configuration and tuning guidance. `xla_gpu_all_reduce_combine_threshold_bytes` | Integer (bytes) | These flags tune when to combine multiple small AllGather / ReduceScatter / AllReduce into one big AllGather / ReduceScatter / AllReduce to reduce time spent on cross-device communication. For example, for the AllGather / ReduceScatter thresholds on a Transformer-based workload, consider tuning them high enough so as to combine at least a Transformer Layer’s weight AllGather / ReduceScatter. By default, the combine_threshold_bytes is set to 256. `xla_gpu_all_gather_combine_threshold_bytes` | Integer (bytes) | See xla_gpu_all_reduce_combine_threshold_bytes above. `xla_gpu_reduce_scatter_combine_threshold_bytes` | Integer (bytes) | See xla_gpu_all_reduce_combine_threshold_bytes above. `xla_gpu_pipeline_all_gather` | Enum (default/off/on/explicit) | Control all-gather pipelining. `default` follows optimization effort and is the flag default: it selects `on` at O1 or above, or when execution-time optimization effort is at least 0.2; otherwise, it selects `off`. `off` disables pipelining, `on` considers all eligible all-gathers, and `explicit` considers only operations marked with `is_pipelineable=true`. `xla_gpu_pipeline_all_reduce` | Enum (default/off/on/explicit) | Control all-reduce pipelining with the same mode semantics and default as `xla_gpu_pipeline_all_gather`. `xla_gpu_pipeline_reduce_scatter` | Enum (default/off/on/explicit) | Control reduce-scatter pipelining with the same mode semantics as `xla_gpu_pipeline_all_gather`. The flag default is `on`. `xla_gpu_enable_pipelined_host_offloading` | Boolean (true/false) | Enable pipelining of host offloading instructions. `xla_gpu_enable_while_loop_double_buffering` | Boolean (true/false) | Enable double-buffering for while loop. `xla_gpu_enable_all_gather_combine_by_dim` | Boolean (true/false) | Combine all-gather ops with the same gather dimension or irrespective of their dimension. `xla_gpu_enable_reduce_scatter_combine_by_dim` | Boolean (true/false) | Combine reduce-scatter ops with the same dimension or irrespective of their dimension. --- ### Gpu Aot (`docs/gpu_aot.md`) # 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](gpu_aot_compatibility.md). ## Workflow and APIs 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`](https://github.com/openxla/xla/blob/main/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`](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_aot_compilation_result.h) (a `CompiledModule`; [`AotCompilationResult`](https://github.com/openxla/xla/blob/main/xla/service/compiler.h) is a deprecated alias). This object encapsulates the serialized [`GpuExecutableProto`](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_executable.proto) and exposes `SerializeAsString()` and `LoadExecutable()`. * **Command-line tooling:** The [`xla_compile`](https://github.com/openxla/xla/blob/main/xla/service/xla_compile_main.cc) tool compiles HLO or StableHLO modules into serialized executables: ```bash 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](https://github.com/openxla/xla/blob/main/xla/backends/gpu/target_config/README.md)). ## What is serialized The top-level container for serialized GPU programs is [`GpuExecutableProto`](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_executable.proto): * `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](gpu_aot_compatibility.md). ## How an artifact is loaded and executed 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 [`CommandBufferThunk`](https://github.com/openxla/xla/blob/main/xla/backends/gpu/runtime/command_buffer_thunk.cc)s 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 Not all executable load failures are caused by serialization regressions. Two common environment mismatches produce load-time errors: * **Compute capability mismatch:** [`GpuExecutable::FromProto`](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_executable.cc) 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`](https://github.com/openxla/xla/blob/main/xla/stream_executor/cuda/cuda_runtime_abi_version.cc) 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`) # 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](gpu_aot.md). The contract applies exclusively to the XLA:GPU backend ([`GpuExecutableProto`](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_executable.proto)). CPU and TPU AOT executables use different serialization mechanisms and are out of scope. ## Compatibility windows {#compatibility} 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} 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](#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 | Change | Verdict | Compatibility impact | Safe procedure | | :--- | :--- | :--- | :--- | | Compiler pass, fusion, or codegen optimization (same schema) | Safe | None | Land directly. | | In-memory runtime refactor with unchanged behavior for existing thunks (caches, buffer ownership) | Safe | None | Land directly with unit tests. | | `ToProto`/`FromProto` refactor with identical serialized bytes | Safe | None | Land directly; include round-trip tests. | | Add new proto field whose unset default preserves legacy behavior | Safe | None | Ensure `FromProto` handles unset state cleanly. | | Add new proto field required for runtime correctness | Breaking | Forward compatibility | Staged rollout (runtime support first). | | Emit new thunk kind, oneof variant, or enum value | Breaking | Forward compatibility | Staged rollout (add runtime thunk first). | | Emit new HLO opcode into serialized module | Breaking | Forward compatibility | Staged rollout; gate emission until runtime window passes. | | Remove thunk kind, proto field, or `FromProto` branch | Breaking | Backward compatibility | Wait 6 months after emission ceased; reserve tags. | | Change existing field number or proto data type | Breaking | Forward & backward compatibility | Never modify existing field tags; allocate a new tag. | | Change semantic meaning of an existing field | Breaking | Forward & backward compatibility | Introduce a new field representing the new behavior. | | Rename or remove registered kernel symbol or FFI handler | Breaking | Backward compatibility | Retain alias registration under legacy name for 6 months. | | Alter load-time thunk pass or toggle runtime-read flag default | Depends | Backward compatibility | Allowed only if correct for every artifact emitted in the last 6 months. | ## Staged rollout for breaking changes {#staged-rollout} 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 Every compatibility fallback in the runtime should document its expiration date: ```cpp // Backward-compatibility fallback for legacy AOT-compiled kernels without // an explicit CollectiveKernelSpec. // Can be removed in (6 months backward compatibility window). ``` Public OpenXLA rollout examples: * [PR #49046](https://github.com/openxla/xla/pull/49046): Introduced dedicated stream assignment in the runtime first, keeping emission disabled by default during the rollout window. * [PR #46865](https://github.com/openxla/xla/pull/46865): Added runtime thunk support for `CollectiveReduce`, with compiler emission guarded behind an experimental flag. ## Authoring runtime code for deserialized executables 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 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`](https://github.com/openxla/xla/blob/main/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc)). * **Golden AOT tests:** End-to-end compatibility tests live under [`xla/tests/aot_compatibility/gpu`](https://github.com/openxla/xla/blob/main/xla/tests/aot_compatibility/gpu) (such as `collective_ops_aot_test.cc`), running against serialized golden executables in `executables//v/`. * **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`) # 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](https://github.com/openxla/stablehlo/tree/main) (“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](https://opensource.googleblog.com/2023/05/pjrt-simplifying-ml-hardware-and-framework-integration.html) 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 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 A number of notable optimization passes happen on HLO, as HLO->HLO rewrites. ### SPMD Partitioner The XLA SPMD partitioner, as described in [GSPMD: General and Scalable Parallelization for MLComputation Graphs](https://arxiv.org/pdf/2105.04663.pdf), 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](https://dl.acm.org/doi/pdf/10.1145/3567955.3567959) 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 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 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 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 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 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 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 The XLA:GPU backend generates high-performance LLVM IR directly for a number of operations (reductions, transposes, etc). ### Triton code generation For more advanced fusions which include matrix multiplication or softmax, XLA:GPU uses [Triton](https://github.com/openai/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 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`) # 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 ### Using Environment Variables 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: ```shell 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 ```shell XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH" ``` 1. HLO protos ```shell XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH" ``` 1. HLO Snapshots ```shell XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH" ``` 1. Graph render with graphviz server (only works well for small graphs) ```shell XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH" ``` 1. Graph render to HTML file (only works well for small graphs) ```shell 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. ```shell XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_syntax_sugar_async_ops=true" ``` ## Dump Specific Intermediate Passes In addition to the standard pre-optimized / final-optimized HLOs, you can also dump the state of HLOs after a particular compiler pass. ```shell 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: ```shell 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: ```shell XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*" ``` ### JAX-specific Options #### 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: ```python jax.jit(f).lower(*args).as_text('hlo') ``` For dumping to files during HLO compilation passes, specify: ```python 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 [`jaxpr`s](https://docs.jax.dev/en/latest/jaxpr.html) are JAX's intermediate representation for program traces. To dump this, set the environment variables: ```shell JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr ``` Learn more in JAX documentation on [Exporting and serializing staged-out computations: Debugging](https://docs.jax.dev/en/latest/export/export.html#debugging). ## Google Colab ### Environment variables 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: ```python 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 Similar to local execution; for live, interactive introspection you can directly print a computation’s pre-optimized HLO: ```python 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: ```python 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: ```shell 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: ```shell --xla_mosaic_dump_to=/tmp/mosaic_dumps ``` Or, set TPU init arguments as an environment variable: ```shell export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps" ``` Check out the [JAX documentation on Pallas and Mosaic](https://docs.jax.dev/en/latest/pallas/index.html) to learn more. ## More with HLO Dumps ### Finding the right computation 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 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](tools.md#hlo-opt-convert-hlo-module-formats). ### Replay 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: ```shell bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu /tmp/xladump/module_4561.before_optimizations.txt ``` GPU backend: ```shell 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`) # 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 You can use HLO isolation components through the standard OpenXLA/TensorFlow build mechanism (Bazel) or via binary distribution. ### Source and Bazel Setup 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: ```bash bazel build -c opt //third_party/tensorflow/compiler/xla/tools/hlo_isolation:hlo_isolation_test ``` ## Command-Line Interface (CLI) 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 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 ```bash ./hlo_isolation_test \ --hlo_file=/path/to/failing_fusion.hlo \ --test_platform=gpu \ --reference_platform=interpreter ``` ## Result and Artifact Dumps 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 On numeric mismatch, the tool writes the following debug artifacts: 1. The interactive HTML mismatch report (`failed-module-.html`). 2. The failed HLO submodule text (`failed-module-.txt`). 3. The expected output literal (`failed--expected.txt`). 4. The actual output literal (`failed--actual.txt`). 5. The mismatching elements summary (`failed--mismatches.txt`). ### Interactive HTML Graph UI When numeric mismatches are detected during isolation testing, an interactive HTML report (`failed-module-.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 - **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--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--expected.txt`). ## C++ Integration and API 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 The core API provides functional interfaces to run modules and fetch structured reports: ```cpp #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> results = xla::hlo_isolation::RunIsolationPipeline( input_hlo_module, &my_test_runner, &my_reference_runner, options); ``` ### Using the Test Mixin 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`): ```cpp #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> {}; TEST_F(MyCustomPassIsolationTest, ChecksMyFusionSanity) { RunAndVerifyIsolationTest(my_failing_module); } ``` ## Sharded Execution (K8s/Slurm) 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 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. ```yaml 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 - **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`) # HLO Passes This document outlines the [HLO](https://openxla.org/xla/terminology) optimizations and transformations passes in the [XLA compiler](https://openxla.org/xla/architecture). ## Introduction 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](https://openxla.org/xla/operation_semantics) of the arrays to make the optimization or transformation easier. For example: * [`AlgebraicSimplifier`:](https://github.com/openxla/xla/blob/c37fc6a383b870f43cef82280418fcefcc90b0f8/xla/hlo/transforms/simplifiers/algebraic_simplifier.h#L417) 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`:](https://github.com/openxla/xla/tree/main/xla/hlo/transforms/simplifiers/hlo_rematerialization.h) A pass that recomputes selected expressions in the computation to reduce memory pressure caused by long live ranges of array-shaped values. ## Developer details The base class for HLO passes can be found in [`xla/hlo/pass/hlo_pass_interface.h`](https://github.com/openxla/xla/blob/main/xla/hlo/pass/hlo_pass_interface.h). HLO pass should not extend this class directly but instead should extend [`HloModulePass`](https://github.com/openxla/xla/blob/main/xla/hlo/pass/hlo_pass_interface.h#L142). See also [XLA HLO Pass Framework](https://github.com/openxla/xla/tree/main/xla/hlo/pass#readme). ### Tooling and Testing 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](https://openxla.org/xla/tools#hlo-opt_hlo_pass_development_and_debugging). For information on writing unit tests for HLO Passes see [Testing HLO Passes](https://openxla.org/xla/test_hlo_passes). ## Hardware-independent HLO Pass Examples 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`](https://github.com/openxla/xla/tree/main/xla/hlo/transforms). ### Rematerialization See also [`HloRematerialization`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/hlo_rematerialization.h). 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 See also [`AlgebraicSimplifier`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/algebraic_simplifier.h). A grab bag of simplifications, optimizations, and canonicalizations. Analogous to [LLVM’s `instcombine` pass](https://llvm.org/docs/Passes.html#instcombine-combine-redundant-instructions). ### Constant Folding See also [`HloConstantFolding`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/hlo_constant_folding.h). Replaces expressions which can be evaluated at compile time with their constant equivalent. ### Dead Code Elimination See also [`HloDCE`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/hlo_dce.h) . Removes operations with unused results (fast implementation). ### Call Graph Flattening See also [`FlattenCallGraph`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/flatten_call_graph.h). 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 See also [`ReshapeMover`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/reshape_mover.h). 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 See also [`ZeroSizedHloElimination`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/zero_sized_hlo_elimination.h). 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 Passes specific to the TPU backend. ### Model parallelism 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`](https://github.com/openxla/xla/blob/main/xla/service/sharding_propagation.h). Pass to support dividing operations across devices along non-batch dimensions. ### Handling of bfloat16 See also [`BFloat16ConversionFolding`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/bfloat16_conversion_folding.h), [`BFloat16MixedPrecisionRemoval`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/float_normalization.h), and [`BFloat16Propagation`](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/bfloat16_propagation.h). 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 See also [`GatherExpander`](https://github.com/openxla/xla/blob/main/xla/service/gather_expander.h), and [`BatchNormExpander`](https://github.com/openxla/xla/blob/main/xla/service/batchnorm_expander.h). 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 Passes specific to the GPU backend are found in [`xla/service/gpu`](https://github.com/openxla/xla/tree/main/xla/service/gpu). These passes can be identified as classes defined in `namespace gpu`. ### cuDNN Rewriter See also [`CudnnFusedConvRewriter`](https://github.com/openxla/xla/blob/main/xla/service/gpu/transforms/cudnn_fused_conv_rewriter.h) and [`CudnnNormRewriter`](https://github.com/openxla/xla/blob/main/xla/service/gpu/transforms/cudnn_norm_rewriter.h). Rewrites fused convolution and norm operations into their respective library calls in cuDNN. ## CPU-specific HLO Pass Examples Passes specific to the CPU backend are found in [`xla/service/cpu`](https://github.com/openxla/xla/tree/main/xla/service/cpu). These passes can be identified as classes defined in `namespace cpu`. ### Convolution Canonicalization See also [`ConvCanonicalization`](https://github.com/openxla/xla/blob/main/xla/service/cpu/conv_canonicalization.h). Canonicalizes convolutions so that they can be lowered to a fast implementation in Eigen. ### Operation Parallelization See also [`ParallelTaskAssigner`](https://github.com/openxla/xla/blob/main/xla/service/cpu/parallel_task_assignment.h). Partitions HLOs into tasks to run on separate threads. ## Analysis passes 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`](https://github.com/openxla/xla/tree/main/xla/hlo/analysis). ### Analysis Pass Examples #### Dataflow Analysis See also [`HloDataflowAnalysis`](https://github.com/openxla/xla/tree/main/xla/hlo/analysis/hlo_dataflow_analysis.h). Identifies all HLO values in the graph and their uses. #### Alias Analysis See also [`HloAliasAnalysis`](https://github.com/openxla/xla/tree/main/xla/hlo/analysis/hlo_alias_analysis.h). Identifies must-alias relationships between values in the program. #### Computation Cost Analysis See also [`HloCostAnalysis`](https://github.com/openxla/xla/tree/main/xla/service/hlo_cost_analysis.h). Computes FLOP count and memory usage for all operations in the program. #### HLO Verification See also [`HloVerifier`](https://github.com/openxla/xla/tree/main/xla/service/hlo_verifier.h). Verifies various invariants of the HLO graph. --- ### Hlo To Thunks (`docs/hlo_to_thunks.md`) # 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". HLO to thunks diagram ## Pre-optimization HLO 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 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 This includes passes like the [Shardy Partitioner](https://openxla.org/shardy/overview) or those for SPMD sharding. ### Optimization passes This can include both legalization passes and simplification passes. ### Collective optimization passes Similar to **Optimization passes**, but focuses on collective ops. ### Layout assignment passes 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 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 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 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 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 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 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 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 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](https://github.com/openxla/xla/tree/main/xla/service/gpu/ir_emitter_unnested.h). `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](https://github.com/openxla/xla/tree/main/xla/service/cpu/thunk_emitter.h) 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 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](https://github.com/openxla/xla/tree/main/xla/stream_executor/gpu/gpu_command_buffer.h). A command buffer is represented in a thunk sequence by [CommandBufferThunk](https://github.com/openxla/xla/tree/main/xla/backends/gpu/runtime/command_buffer_thunk.h). The emitter does not produce this thunk directly from HLO instructions. Instead, this is done by [CommandBufferConversionPass](https://github.com/openxla/xla/tree/main/xla/backends/gpu/runtime/command_buffer_conversion_pass.h) that runs on the ThunkSequence itself. The pass identifies contiguous sub-sequences of compatible thunks (e.g., a series of `KernelThunk`s and `GemmThunk`s). 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 The final product of the XLA compilation pipeline is a self-contained, platform-specific [Executable](https://github.com/openxla/xla/tree/main/xla/service/executable.h). 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](https://github.com/openxla/xla/tree/main/xla/pjrt/pjrt_executable.h)), 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](https://github.com/openxla/xla/tree/main/xla/service/gpu/gpu_compiler.cc) and [CpuCompiler](https://github.com/openxla/xla/tree/main/xla/service/cpu/cpu_compiler.cc) 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`) # Indexing Analysis HLO indexing analysis is a [dataflow analysis](https://en.wikipedia.org/wiki/Data-flow_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>` ```c 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 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`](./symbolic_expressions.md) infrastructure and adds HLO semantics. ### Coalescing 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 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 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](https://github.com/openxla/xla/blob/main/xla/service/gpu/triton_tiling_propagation.h) 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 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](https://openxla.org/xla/operation_semantics#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 To minimize recomputation, we need a framework for symbolic computation. This is implemented as [`SymbolicExpr` and `SymbolicMap`](./symbolic_expr.md). *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. ```c 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 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 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 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 constraints_; }; ``` *Code reference: [indexing_map.h#L114](https://github.com/openxla/xla/blob/18261672d5c2704a6e1017a0b66bd70fabb7a445/xla/hlo/analysis/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](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#concatenate). `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 ` in `, 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 ### Elementwise For elementwise ops the indexing map is an identity. ```c 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) Broadcasting means that some of the dimensions will be removed when we map output to input and added when we map input to output. ```c 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) Iota has no input tensor operand, so there is no input index arguments. ```c 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) DynamicSlice has offsets known only at runtime. ```c 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) ```c 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) Only the simplified gather is supported. See [gather_simplifier.h](https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/gather_simplifier.h). ```c 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) Indexing map for transpose is a permutation of input/output dimensions. ```c 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) Indexing map for reverse changes the reverted dimensions to `upper_bound(d_i) - d_i`: ```c 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) Variadic reduction have several inputs and several initial values, the map from output to input adds the reduced dimensions. ```c 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) 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. ```c 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) Reshapes come in different flavors. #### Collapse shape This is a "linearizing" reshape from N-D to 1D. ```c 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. ```c 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. ```c 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 ```c 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 A bitcast op can be represented as a [sequence of transpose-reshape-transpose](https://github.com/openxla/xla/blob/578b6df240be94c3c84129fd83f34487efc623a5/xla/shape_util.h#L813). Therefore, its indexing maps are just a composition of indexing maps for this sequence. ### [Concatenate](https://openxla.org/xla/operation_semantics#concatenate) 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. ```c 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) Indexing maps for dot are very similar to the ones of reduce. ```c 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: ``` (d0, d1, d2)[s0] -> (d0, d1, s0), domain: d0 in [0, 3], d1 in [0, 127], d2 in [0, 255], s0 in [0, 63] ``` - p1 -> output: ``` (d0, d1, d2)[s0] -> (d0, s0, d1), domain: d0 in [0, 3], d1 in [0, 255], d2 in [0, 63], s0 in [0, 127] ``` ### [Pad](https://openxla.org/xla/operation_semantics#pad) Indexing of PadOp is the inverse of SliceOp indexing. ```c p0 = f32[4, 4] parameter(0) p1 = f32[] parameter(1) pad = f32[12, 16] pad(p0, p1), padding=1_4_1x4_8_0 ``` The padding config `1_4_1x4_8_0` denotes `lowPad_highPad_interiorPad_dim_0 x lowPad_highPad_interiorPad_dim_1`. The output to input maps: - output -> p0: ``` (d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4), domain: d0 in [1, 7], d1 in [4, 7], (d0 - 1) mod 2 in [0, 0] ``` - output -> p1: ``` (d0, d1) -> (), domain: d0 in [0, 11], d1 in [0, 15] ``` ### [ReduceWindow](https://openxla.org/xla/operation_semantics#reducewindow) ReduceWindow in XLA also performs padding. Therefore, the indexing maps can be computed as a composition of ReduceWindow indexing that does not do any padding and PadOp's indexing. ```c c_inf = f32[] constant(-inf) p0 = f32[1024, 514] parameter(0) outpu = f32[1024, 3] reduce-window(p0, c_inf), window={size=1x512 pad=0_0x0_0}, to_apply=max ``` The output to input maps: - `output -> p0`: ``` (d0, d1)[s0] -> (d0, d1 + s0), domain: d0 in [0, 1023], d1 in [0, 2], s0 in [0, 511] ``` - `output -> c_inf`: ``` (d0, d1) -> (), domain: d0 in [0, 1023], d1 in [0, 2] ``` ## Indexing Maps for Fusion Indexing map for fusion op is a composition of indexing maps for every op in the cluster. It can happen that some inputs are read several times with different access patterns. ### One input, several indexing maps Here is an example for `p0 + transpose(p0)`. ```c f { p0 = f32[1000, 1000] parameter(0) transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0} ROOT a0 = f32[1000, 1000] add(p0, transpose_p0) } ``` The output-to-input indexing maps for `p0` will be `(d0, d1) -> (d0, d1)` and `(d0, d1) -> (d1, d0)`. It means that to compute one element of the output we might need to read the input parameter twice. ### One input, deduplicated indexing map There are cases when the indexing maps are actually the same, even though it is not immediately obvious. ```c f { p0 = f32[20, 10, 50] parameter(0) lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2} lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1) lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1} rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0} rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1) rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2} ROOT output = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2) } ``` The output-to-input indexing map for `p0` in this case is just `(d0, d1, d2) -> (d2, d0, d1)`. ### Softmax The output-to-input indexing maps for `parameter 0` for softmax: ``` (d0, d1, d2)[s0] -> (d0, d1, s0), domain: d0 in [0, 1], d1 in [0, 64], d2 in [0, 124], s0 in [0, 124] ``` and ``` (d0, d1, d2) -> (d0, d1, d2), domain: d0 in [0, 1], d1 in [0, 64], d2 in [0, 124] ``` where `s0` refers to the innermost dimension of the input. For more examples see [indexing_analysis_test.cc](https://github.com/openxla/xla/blob/main/xla/hlo/analysis/indexing_analysis_test.cc). ## Indexing Map simplification We can leverage the knowledge about lower and upper bounds of the sub-expressions in the symbolic maps to simplify them even more. Indexing map simplification can rewrite the following expressions. 1. `(d0, d1) -> (d0 + d1 floordiv 16, d1 mod 16)` for **d** in `[0, 6] x [0, 14]` becomes `(d0, d1) -> (d0, d1)` 2. `(d0, d1, d2) -> ((100d0 + 10d1 + d2) floorDiv 100, ((100d0 + 10d1 + d2) mod 100) floordiv 10, d2 mod 10)` for `di in [0, 9]` becomes `(d0, d1, d2) -> (d0, d1, d2)`. 3. `(d0, d1, d2) -> ((16d0 + 4d1 + d2) floordiv 8, (16d0 + 4d1 + d2) mod 8)` for `d_i in [0, 9]` becomes `(d0, d1, d2) -> (2d0 + (4d1 + d2) floordiv 8,(4d1 + d2) mod 8)`. 4. `(d0, d1) -> (-(-11d0 - d1 + 109) floordiv 11 + 9)` for **d** in `[0, 9] x [0, 10]` becomes `(d0, d1) -> (d0)`. It allows us to understand that some of the chained reshapes in HLO cancel each other. For instance: ```c p0 = f32[10, 10, 10] parameter(0) reshape1 = f32[50, 20] reshape(p0) reshape2 = f32[10, 10, 10] reshape(reshape1) ``` After the composition of indexing maps and their simplification we will get the following map: `(d0, d1, d2) -> (d0, d1, d2)`. Indexing map simplification also simplifies the constraints. 1. Constraints of type `lower_bound <= symbolic_expr (floordiv, ceildiv, +, -, *, mod, min, max) constant <= upper_bound` are rewritten as `updated_lower_bound <= symbolic_expr <= updated_upped_bound`. 2. Constraints that are always satisfied, e.g. `d0 + s0 in [0, 20]` for `d0 in [0, 5]` and `s0 in [1, 3]` are eliminated. 3. Symbolic expressions in the constraints are optimized as the indexing symbolic map above. For more examples see [indexing_map_test.cc](https://github.com/openxla/xla/blob/main/xla/hlo/analysis/indexing_map_test.cc). --- ### Lhs Cost Model (`docs/lhs_cost_model.md`) # LHS Cost Model --- ## tldr; This page describes the internals of the cost model used by Latency Hiding Scheduler. If you are interested in tuning the model go straight to the [Tuning section](#tuning). The Latency Hiding Scheduler (LHS) is a compiler pass that schedules a HLO DAG in a way that minimizes wall time. Its decisions are guided by the unified cost model, which uses a mixture of performance tables and analytical models. In particular XLA embeds performance tables for a GEMMs and fast-interconnect collectives, and uses analytical networking and fusion cost model for other cases. The rest of the document describes the inner workings of these on a high level. --- ## Performance tables – ICI collectives Performance table consist of two main components: a collector and an interpolator. ### Collector [The **collector**](https://github.com/openxla/xla/blob/e6a0a911eb79f540d501458f953393ede9e0048c/xla/tools/collective_perf_table_gen_main.cc#L1) is a C++ tool responsible for generating the performance tables for collective operations. It measures the performance of individual HLO ops (e.g., `all-gather`, `all-reduce`) across a statically defined parameter space. #### How It Works The tool performs a sweep over a range of collective ops, transfer sizes, and transfer schemes for a given cluster. It uses the existing multi-host HLO runner infrastructure and `ExecutionProfile` data to run the generated HLO and gather performance metrics. #### Data Collection Parameters Latency tables are collected for a cross-product of the following parameters: * **Collective Type**: * `all-reduce` * `all-gather` * `reduce-scatter` * **Transfer Size**: * Logarithmic scale from 1024B up to 2GiB (e.g., 1024B, 2048B, 4096B, ...) * **Transfer Scheme**: * `rail-aligned` * `non-rail-aligned` This sweep is run for intra-node clusters with **2, 4, and 8 devices**. #### Output The result of a collection run is a latency table in `.pbtxt` format (approximately 116 KB per platform). ### Interpolator [The **interpolator**](https://github.com/openxla/xla/blob/e6a0a911eb79f540d501458f953393ede9e0048c/xla/service/gpu/model/collective_interpolator.h#L40) is the compiler component that consumes the generated performance tables to provide runtime estimates during compilation. #### Internal Data Structure On initialization, the Interpolator processes the performance table into a map. This map uses a tuple of `(collective_type, transfer_scheme)` as its **key**. The **value** associated with each key is a 2D Euclidean plane. This plane indexes the **network throughput** (measured by the Collector) based on two axes: 1. Transfer size. 2. Number of devices involved. #### Lookup and Interpolation When the compiler encounters a collective operation, the Interpolator performs the following steps: 1. It identifies the correct 2D throughput plane using the operation's `(collective_type, transfer_scheme)` as the map key. 2. It then uses a **weighted average retrieval** (based on Euclidean distance) within that 2D plane, using the operation's `(transfer_size, num_devices)` as the query point. 3. The result of this lookup is a single, unique **network throughput** value. ### Rationale: Throughput and Extrapolation The system is designed to store **network throughput** rather than raw latency. This design choice significantly simplifies extrapolating performance for transfer sizes not explicitly present in the table. If the latency tables capture network bandwidth saturation at a collective size `S`, the throughput `T` at that point is considered the maximum. For any new collective of size `S'` > `S`, the runtime can be estimated as: $$\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated}}}$$ This allows the model to estimate performance for collectives of any size, even those larger than the 2GiB maximum measured by the Collector. **Important:** This extrapolation model relies on the assumption that the generated latency tables **capture true network bandwidth saturation**. If the tables do not contain measurements at or beyond the saturation point, the interpolator will: * **Underestimate** the maximum throughput. * Consequently, **overestimate** the runtime for large transfers. In general XLA:GPU teams maintains performance tables, but in cases user decide to provide their own, it is the responsibility of the user generating the tables to ensure they are representative and include measurements in the bandwidth-saturated region for the target hardware. --- ## Performance tables – GEMMs Similar to the system for collectives, GEMM latency tables are supported by two components: a **collector** and an **interpolator**. ### Collector [The **collector**](https://github.com/openxla/xla/blob/e6a0a911eb79f540d501458f953393ede9e0048c/xla/tools/matmul_perf_table_gen_main.cc) is a C++ tool that computes performance tables for General Matrix Multiplications (GEMMs). It measures the performance of matrix multiplications at the HLO `dot` op level. #### How It Works The tool performs a sweep over a static space of GEMM dimensions (batch, two non-contracting, and one contracting dimension) and data types. * **Default Data Types:** `LHS = bf16,f32`, `RHS = bf16,f32`, `OUT = bf16,f32`. * **Infrastructure:** Re-uses the HLO op profiler. #### Collection Parameters Latency tables are collected for a cross-product of the following dimensions: * **batch:** `{1, 2, 4}` * **m (non-contracting):** `{256, 512, ..., 4096}` * **n (non-contracting):** `{256, 512, ..., 4096}` * **k (contracting):** `{256, 512, ..., 4096}` #### Output and Storage A full sweep generates a `.pbtxt` latency table, ready to be consumed by interpolator. ### Interpolator [The **interpolator**](https://github.com/openxla/xla/blob/e6a0a911eb79f540d501458f953393ede9e0048c/xla/service/gpu/model/matmul_interpolator.h#L35) is the compiler component that uses the generated tables to estimate GEMM performance. #### Rationale: FLOPS Saturation The collected latency tables allow the interpolator to reconstruct **FLOPS** for each entry: $$\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime}}$$ A key insight is that FLOPS **saturate** at a certain point; that is, the hardware reaches peak FLOPS beyond a certain matrix shape. This saturation allows the use of the same extrapolation method employed for collectives. #### Lookup and Interpolation The interpolator builds a **4D Euclidean space** from the table data. To provide a performance estimate, it performs a **weighted-average interpolation** within this 4D space. If there's no table for a certain data type, as a heuristic each dimension is normalized to the number of bytes. --- ## Analytical Cost Model - DCN ### S-curve Collective Cost Model The **S-curve** model is a fully analytical networking roofline model. #### Overview The model is designed to estimate the performance of collective operations based on a set of fixed network properties. #### Model Inputs The model requires two categories of inputs: 1. **Fixed Network Properties (User-Defined):** * Collective launch overhead * NIC speed * RTT (round trip time) By default, XLA auto-detects a platform and uses values for the most common architectures. These properties are configurable by the user. See [Tuning section](#tuning) for details. 2. **Per-Collective Inputs:** * Collective type (e.g., `AllGather`, `ReduceScatter`) * Transfer size * Number of nodes involved in the communication #### Integration The S-curve model is integrated into `XLA:GPU` and is being used on Hopper, and Blackwell. --- ## Analytical Cost Model - Fusions For other kernels we rely on the [GPU performance cost model](https://github.com/openxla/xla/blob/e6a0a911eb79f540d501458f953393ede9e0048c/xla/service/gpu/model/gpu_performance_model.h) to estimate the right runtimes. You can read more about it [here](https://github.com/openxla/xla/discussions/10065). --- ## Tuning S-curve model can be tuned by issuing right XLA flags. Default configuration should be good enough in majority of cases, but the model control is exposed in other cases. ``` export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK) export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE" ``` --- ### Oom Debugging (`docs/oom_debugging.md`) # Debug OOM errors with XProf Out Of Memory (OOM) errors occur when the accelerator's (GPU or TPU) High Bandwidth Memory (HBM) capacity is exhausted. Some common causes for OOM issues and debugging techniques are detailed in [E1000 - Compile Time HBM OOM documentation](./errors/error_1000.md) and [JAX documentation on GPU memory allocation](https://docs.jax.dev/en/latest/gpu_memory_allocation.html#common-causes-of-oom-failures). This page describes how to use **XProf's Memory Viewer tool** to visualize your JAX program's memory usage, identify peak usage instances, and debug OOM errors. This involves the following steps: 1. Run your program with [`jax.profiler.trace`](https://docs.jax.dev/en/latest/_autosummary/jax.profiler.trace.html#jax.profiler.trace) to capture the profile. 2. Start XProf in the background, and use the [Memory Viewer tool](https://openxla.org/xprof/memory_viewer) to view memory utilization details. ## Example program The following JAX program leads to an OOM error: ```python import jax from jax import random import jax.numpy as jnp @jax.profiler.trace("/tmp/xprof") @jax.jit def oom(): a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16) return a @ a if __name__ == "__main__": oom() ``` **Note:** Prefer `jax.profiler.trace` instead of `jax.profiler.start_trace`/`jax.profiler.stop_trace` because the `jax.profiler.trace` context manager handles profiling in an exception safe manner. On a TPU-machine, this program fails with: ```shell XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: ``` (On a GPU-machine, the error looks like: `XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.`) ## Run XProf Install `xprof` (`pip install xprof`), and start an XProf instance specifying the directory where the profile is stored: ```shell xprof --logdir=/tmp/xprof/ --port=6006 ``` Go to the instance (on a local machine, at `http://localhost:6006`). In the *Tools* dropdown, select *Memory Viewer*, and in the Memory Viewer tool window, select *HBM* in the *Memory Types* dropdown (usually selected by default). The [XProf: Memory Viewer tool documentation](https://openxla.org/xprof/memory_viewer#memory_viewer_components) describes the components of the tool and the information presented. Focus on the *HLO Ops at Peak Memory Allocation* section that shows three buffer charts at the peak memory usage point. The buffer includes: - **Program Inputs and Outputs:** Training batches, optimizer states, etc. - **TensorCore and SparseCore Temporaries:** Dynamic memory required for intermediate calculations (like activations, gradients, etc.) You can hover on the buffer charts to get more details about the Op like it's size, shape, allocation type, and more. This can help you identify and evaluate Ops that may have high or long-lasting temporaries, any large input/intermediate/output tensors that have inefficient padding, etc., that are contributing to the peak memory and need to be adjusted or optimized. Learn specific debugging techniques in [E1000: Debugging](https://openxla.org/xla/errors/error_1000#debugging). --- ### Operation Semantics (`docs/operation_semantics.md`) # Operation semantics The following describes the semantics of operations defined in the [`XlaBuilder`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h) interface. Typically, these operations map one-to-one to operations defined in the RPC interface in [`xla_data.proto`](https://github.com/openxla/xla/tree/main/xla/xla_data.proto). A note on nomenclature: the generalized data type XLA deals with is an N-dimensional array holding elements of some uniform type (such as 32-bit float). Throughout the documentation, *array* is used to denote an arbitrary-dimensional array. For convenience, special cases have more specific and familiar names; for example a *vector* is a 1-dimensional array and a *matrix* is a 2-dimensional array. Learn more about the structure of an Op in [*Shapes and layout*](shapes.md) and [*Tiled Layout*](tiled_layout.md). ## Abs See also [`XlaBuilder::Abs`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise abs `x -> |x|`. **`Abs(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - abs](https://openxla.org/stablehlo/spec#abs). ## Add See also [`XlaBuilder::Add`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise addition of `lhs` and `rhs`. **`Add(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Add: **`Add(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - add](https://openxla.org/stablehlo/spec#add). ## AddDependency See also [`HloInstruction::AddDependency`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `AddDependency` may appear in HLO dumps, but they are not intended to be constructed manually by end users. > **Note:** `AddDependency` is only found in HLO. It is not found in StableHLO. ## AfterAll See also [`XlaBuilder::AfterAll`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). AfterAll takes a variadic number of tokens and produces a single token. Tokens are primitive types which can be threaded between side-effecting operations to enforce ordering. `AfterAll` can be used as a join of tokens for ordering an operation after a set of operations. **`AfterAll(tokens)`** Arguments | Type | Semantics --------- | ----------------- | ------------------------- `tokens` | vector of `XlaOp` | variadic number of tokens For StableHLO information see [StableHLO - after_all](https://openxla.org/stablehlo/spec#after_all). ## AllGather See also [`XlaBuilder::AllGather`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs concatenation across replicas. **`AllGather(operand, all_gather_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)`** | Arguments | Type | Semantics | | ----------------------- | -------------------- | --------------------------- | | `operand` | `XlaOp` | Array to concatenate across | : : : replicas : | `all_gather_dimension` | `int64` | Concatenation dimension | | `shard_count` | `int64` | The size of each replica | : : : group : | `replica_groups` | vector of vectors of | Groups between which the | : : `int64` : concatenation is performed : | `channel_id` | optional | Optional channel ID for | : : `ChannelHandle` : cross-module communication : | `layout` | optional `Layout` | Creates a layout pattern | : : : that will capture the : : : : matched layout in the : : : : argument : | `use_global_device_ids` | optional `bool` | Returns true if the ids in | : : : the ReplicaGroup config : : : : represent a global id : - `replica_groups` is a list of replica groups between which the concatenation is performed (replica id for the current replica can be retrieved using [`ReplicaId`](#replicaid)). The order of replicas in each group determines the order in which their inputs are located in the result. `replica_groups` must either be empty (in which case all replicas belong to a single group, ordered from `0` to `N - 1`), or contain the same number of elements as the number of replicas. For example, `replica_groups = {0, 2}, {1, 3}` performs concatenation between the replicas `0` and `2`, and `1` and `3`. - `shard_count` is the size of each replica group. We need this in cases where `replica_groups` are empty. - `channel_id` is used for cross-module communication: only `all-gather` operations with the same `channel_id` can communicate to each other. - `use_global_device_ids` Returns true if the ids in the ReplicaGroup config represent a global id of (replica_id * partition_count + partition_id) instead of a replica id. This enables more flexible grouping of devices if this all-reduce is both cross-partition and cross-replica. The output shape is the input shape with the `all_gather_dimension` made `shard_count` times larger. For example, if there are two replicas and the operand has the value `[1.0, 2.5]` and `[3.0, 5.25]` respectively on the two replicas, then the output value from this op where `all_gather_dim` is `0` will be `[1.0, 2.5, 3.0,5.25]` on both replicas. The API of `AllGather` is internally decomposed into 2 HLO instructions (`AllGatherStart` and `AllGatherDone`). See also [`HloInstruction::CreateAllGatherStart`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `AllGatherStart`, `AllGatherDone` serve as primitives in HLO. These ops may appear in HLO dumps, but they are not intended to be constructed manually by end users. > **Note:** `AllGatherStart` and `AllGatherDone` are only found in HLO. They are > not found in StableHLO. For StableHLO information see [StableHLO - all_gather](https://openxla.org/stablehlo/spec#all_gather). ## AllReduce See also [`XlaBuilder::AllReduce`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs a custom computation across replicas. **`AllReduce(operand, computation, replica_groups, channel_id, shape_with_layout, use_global_device_ids)`** | Arguments | Type | Semantics | | ----------------------- | --------------------- | -------------------------- | | `operand` | `XlaOp` | Array or a non-empty tuple | : : : of arrays to reduce across : : : : replicas : | `computation` | `XlaComputation` | Reduction computation | | `replica_groups` | `ReplicaGroup` vector | Groups between which the | : : : reductions are performed : | `channel_id` | optional | Optional channel ID for | : : `ChannelHandle` : cross-module communication : | `shape_with_layout` | optional `Shape` | Defines the layout of the | : : : data transferred : | `use_global_device_ids` | optional `bool` | Returns true if the ids in | : : : the ReplicaGroup config : : : : represent a global id : - When `operand` is a tuple of arrays, the all-reduce is performed on each element of the tuple. - `replica_groups` is a list of replica groups between which the reduction is performed (replica id for the current replica can be retrieved using [`ReplicaId`](#replicaid)). `replica_groups` must either be empty (in which case all replicas belong to a single group), or contain the same number of elements as the number of replicas. For example, `replica_groups = {0, 2}, {1, 3}` performs reduction between the replicas `0` and `2`, and `1` and `3`. - `channel_id` is used for cross-module communication: only `all-reduce` operations with the same `channel_id` can communicate to each other. - `shape_with_layout`: forces the layout of the AllReduce to the given layout. This is used to guarantee the same layout for a group of AllReduce ops compiled separately. - `use_global_device_ids` Returns true if the ids in the ReplicaGroup config represent a global id of (replica_id * partition_count + partition_id) instead of a replica id. This enables more flexible grouping of devices if this all-reduce is both cross-partition and cross-replica. The output shape is the same as the input shape. For example, if there are two replicas and the operand has the value `[1.0, 2.5]` and `[3.0, 5.25]` respectively on the two replicas, then the output value from this op and summation computation will be `[4.0, 7.75]` on both replicas. If the input is a tuple, the output is a tuple as well. Computing the result of `AllReduce` requires having one input from each replica, so if one replica executes an `AllReduce` node more times than another, then the former replica will wait forever. Since the replicas are all running the same program, there are not a lot of ways for that to happen, but it is possible when a while loop's condition depends on data from `infeed` and the data that is `infeed` causes the while loop to iterate more times on one replica than another. The API of `AllReduce` is internally decomposed into 2 HLO instructions (`AllReduceStart` and `AllReduceDone`). See also [`HloInstruction::CreateAllReduceStart`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `AllReduceStart` and `AllReduceDone` serve as primitives in HLO. These ops may appear in HLO dumps, but they are not intended to be constructed manually by end users. > **Note:** `AllGatherStart` and `AllGatherDone` are only found in HLO. They are > not found in StableHLO. For StableHLO information see > [StableHLO - all_reduce](https://openxla.org/stablehlo/spec#all_reduce). ### CrossReplicaSum See also [`XlaBuilder::CrossReplicaSum`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs `AllReduce` with a summation computation. **`CrossReplicaSum(operand, replica_groups)`** | Arguments | Type | Semantics | | :------------------- | :------------------- | :---------------------------- | | **`operand`** | XlaOp | Array or a non-empty tuple of | : : : arrays to reduce across : : : : replicas : | **`replica_groups`** | vector of vectors of | Groups between which the | : : **`int64`** : reductions are performed : Returns the sum of the operand value within each subgroup of replicas. All replicas supply one input to the sum and all replicas receive the resulting sum for each subgroup. ## AllToAll See also [`XlaBuilder::AllToAll`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). AllToAll is a collective operation that sends data from all cores to all cores. It has two phases: 1. The scatter phase. On each core, the operand is split into `split_count` number of blocks along the `split_dimensions`, and the blocks are scattered to all cores, e.g., the ith block is sent to the ith core. 2. The gather phase. Each core concatenates the received blocks along the `concat_dimension`. The participating cores can be configured by: - `replica_groups`: each ReplicaGroup contains a list of replica ids participating in the computation (the replica id for the current replica can be retrieved using [`ReplicaId`](#replicaid)). AllToAll will be applied within subgroups in the specified order. For example, `replica_groups = {{1,2,3}, {4,5,0}}` means that an AllToAll will be applied within replicas `{1, 2, 3}`, and in the gather phase, and the received blocks will be concatenated in the same order of 1, 2, 3. Then, another AllToAll will be applied within replicas 4, 5, 0, and the concatenation order is also 4, 5, 0. If `replica_groups` is empty, all replicas belong to one group, in the concatenation order of their appearance. Prerequisites: - The dimension size of the operand on the `split_dimension` is divisible by `split_count`. - The operand's shape is not tuple. **`AllToAll(operand, split_dimension, concat_dimension, split_count, replica_groups, layout, channel_id)`** | Arguments | Type | Semantics | | ------------------ | ------------------------ | ---------------------------- | | `operand` | `XlaOp` | n dimensional input array | | `split_dimension` | `int64` | A value in the interval | : : : `[0,n)` that names the : : : : dimension along which the : : : : operand is split : | `concat_dimension` | `int64` | A value in the interval | : : : `[0,n)` that names the : : : : dimension along which the : : : : split blocks are : : : : concatenated : | `split_count` | `int64` | The number of cores that | : : : participate in this : : : : operation. If : : : : `replica_groups` is empty, : : : : this should be the number of : : : : replicas; otherwise, this : : : : should be equal to the : : : : number of replicas in each : : : : group. : | `replica_groups` | `ReplicaGroup`vector | Each group contains a list | : : : of replica ids. : | `layout` | optional `Layout` | user-specified memory layout | | `channel_id` | optional `ChannelHandle` | unique identifier for each | : : : send/recv pair : See [xla::shapes for more information on shapes and layouts.](https://openxla.org/xla/shapes). For StableHLO information see [StableHLO - all_to_all](https://openxla.org/stablehlo/spec#all_to_all). ### AllToAll - Example 1. ```cpp XlaBuilder b("alltoall"); auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x"); AllToAll( x, /*split_dimension=*/ 1, /*concat_dimension=*/ 0, /*split_count=*/ 4); ``` In the above example, there are 4 cores participating in the Alltoall. On each core, the operand is split into 4 parts along dimension 1, so each part has shape f32[4,4]. The 4 parts are scattered to all cores. Then each core concatenates the received parts along dimension 0, in the order of core 0-4. So the output on each core has shape f32[16,4]. ### AllToAll - Example 2 - StableHLO In the above example, there are 2 replicas participating in the AllToAll. On each replica, the operand has shape f32[2,4]. The operand is split into 2 parts along dimension 1, so each part has shape f32[2,2]. The 2 parts are then exchanged across the replicas according to their position in the replica group. Each replica collects its corresponding part from both operands and concatenates them along dimension 0. As a result, the output on each replica has shape f32[4,2]. ### RaggedAllToAll See also [`XlaBuilder::RaggedAllToAll`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). RaggedAllToAll performs a collective all-to-all operation, where the input and output are ragged tensors. **`RaggedAllToAll(input, input_offsets, send_sizes, output, output_offsets, recv_sizes, replica_groups, channel_id)`** | Arguments | Type | Semantics | | ---------------- | ------------------------ | ----------------------------- | | `input` | `XlaOp` | N array of type T | | `input_offsets` | `XlaOp` | N array of type T | | `send_sizes` | `XlaOp` | N array of type T | | `output` | `XlaOp` | N array of type T | | `output_offsets` | `XlaOp` | N array of type T | | `recv_sizes` | `XlaOp` | N array of type T | | `replica_groups` | `ReplicaGroup` vector | Each group contains a list of | : : : replica ids. : | `channel_id` | optional `ChannelHandle` | unique identifier for each | : : : send/recv pair : Ragged tensors are defined by a set of three tensors: - `data`: the `data`tensor is “ragged” along its outermost dimension, along which each indexed element has variable size. - `offsets`': the `offsets` tensor indexes the outermost dimension of the `data` tensor, and represents the starting offset of each ragged element of the `data` tensor. - `sizes`: the `sizes` tensor represents the size of each ragged element of the `data` tensor, where the size is specified in units of sub-elements. A sub-element is defined as the suffix of the ‘data’ tensor shape obtained by removing the outermost “ragged” dimension. - The `offsets` and `sizes` tensors must have the same size. An example ragged tensor: ```cpp data: [8,3] = {{a,b,c},{d,e,f},{g,h,i},{j,k,l},{m,n,o},{p,q,r},{s,t,u},{v,w,x}} offsets: [3] = {0, 1, 4} sizes: [3] = {1, 3, 4} // Index 'data' at 'offsets'[0], 'sizes'[0]' // {a,b,c} // Index 'data' at 'offsets'[1], 'sizes'[1]' // {d,e,f},{g,h,i},{j,k,l} // Index 'data' at 'offsets'[2], 'sizes'[2]' // {m,n,o},{p,q,r},{s,t,u},{v,w,x} ``` `output_offsets` must be sharded in a way that each replica has offsets in the target replica output perspective. For i-th output offset, the current replica will send `input[input_offsets[i]:input_offsets[i]+send_sizes[i]]` update to `i`-th replica that will be written to `output_i[output_offsets[i]:output_offsets[i]+send_sizes[i]]` in `i`-th replica `output`. For example, if we have 2 replicas: ```cpp replica 0: input: [1, 2, 2] output:[0, 0, 0, 0] input_offsets: [0, 1] send_sizes: [1, 2] output_offsets: [0, 0] recv_sizes: [1, 1] replica 1: input: [3, 4, 0] output: [0, 0, 0, 0] input_offsets: [0, 1] send_sizes: [1, 1] output_offsets: [1, 2] recv_sizes: [2, 1] // replica 0's result will be: [1, 3, 0, 0] // replica 1's result will be: [2, 2, 4, 0] ``` The ragged all-to-all HLO has the following arguments: - `input`: ragged input data tensor. - `output`: ragged output data tensor. - `input_offsets`: ragged input offsets tensor. - `send_sizes`: ragged send sizes tensor. - `output_offsets`: array of ragged offsets in the target replica output. - `recv_sizes`: ragged recv sizes tensor. The `*_offsets` and `*_sizes` tensors must all have the same shape. Two shapes are supported for the `*_offsets` and `*_sizes` tensors: - `[num_devices]` where ragged-all-to-all may send at most one update to each remote device in the replica group. For example: ```cpp for (remote_device_id : replica_group) { SEND input[input_offsets[remote_device_id]], output[output_offsets[remote_device_id]], send_sizes[remote_device_id] } ``` - `[num_devices, num_updates]` where ragged-all-to-all may send up to `num_updates` updates the same remote device (each at different offsets), for each remote device in the replica group. For example: ```cpp for (remote_device_id : replica_group) { for (update_idx : num_updates) { SEND input[input_offsets[remote_device_id][update_idx]], output[output_offsets[remote_device_id][update_idx]]], send_sizes[remote_device_id][update_idx] } } ``` > **Note:** `RaggedAllToAll` is only found in HLO. It is not found in StableHLO. ## And See also [`XlaBuilder::And`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise AND of two tensors `lhs` and `rhs`. **`And(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for And: **`And(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - and](https://openxla.org/stablehlo/spec#and). ## Async See also [`HloInstruction::CreateAsyncStart`, `HloInstruction::CreateAsyncUpdate`, `HloInstruction::CreateAsyncDone`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `AsyncDone`, `AsyncStart`, and `AsyncUpdate` are internal HLO instructions used for Asynchronous operations and serve as primitives in HLO. These ops may appear in HLO dumps but they are not intended to be constructed manually by end users. > **Note:** `AsyncStart`, `AsyncUpdate`, and `AsyncDone` are only found in HLO. > They are not found in StableHLO. ## Atan2 See also [`XlaBuilder::Atan2`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise atan2 operation on `lhs` and `rhs`. **`Atan2(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Atan2: **`Atan2(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - atan2](https://openxla.org/stablehlo/spec#atan2). ## BatchNormGrad See also [`XlaBuilder::BatchNormGrad`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h) and [the original batch normalization paper](https://arxiv.org/abs/1502.03167) for a detailed description of the algorithm. Calculates gradients of batch norm. **`BatchNormGrad(operand, scale, batch_mean, batch_var, grad_output, epsilon, feature_index)`** Arguments | Type | Semantics --------------- | ------- | ---------------------------------------------------- `operand` | XlaOp | n dimensional array to be normalized (x) `scale` | XlaOp | 1 dimensional array ($\gamma$) `batch_mean` | XlaOp | 1 dimensional array ($\mu$) `batch_var` | XlaOp | 1 dimensional array ($\sigma^2$) `grad_output` | XlaOp | Gradients passed to `BatchNormTraining` ($\nabla y$) `epsilon` | `float` | Epsilon value ($\epsilon$) `feature_index` | `int64` | Index to feature dimension in `operand` For each feature in the feature dimension (`feature_index` is the index for the feature dimension in `operand`), the operation calculates the gradients with respect to `operand`, `offset`, and `scale` across all the other dimensions. The `feature_index` must be a valid index for the feature dimension in `operand`. The three gradients are defined by the following formulas (assuming a 4-dimensional array as `operand` and with feature dimension index `l`, batch size `m` and spatial sizes `w` and `h`): $$ \begin{split} c_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) \\\\ d_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\\\ \nabla x_{ijkl} &= \frac{\gamma_{l}}{\sqrt{\sigma^2_{l}+\epsilon}} \left( \nabla y_{ijkl} - d_l - c_l (x_{ijkl} - \mu_{l}) \right) \\\\ \nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon}} \right) \\\\\ \nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \end{split} $$ The inputs `batch_mean` and `batch_var` represent moments values across batch and spatial dimensions. The output type is a tuple of three handles: | Outputs | Type | Semantics | | -------------- | ----- | --------------------------- | | `grad_operand` | XlaOp | gradient with respect to | : : : input **`operand`** : : : : ($\nabla x$) : | `grad_scale` | XlaOp | gradient with respect to | : : : input **`scale` ** : : : : ($\nabla\gamma$) : | `grad_offset` | XlaOp | gradient with respect to | : : : input : : : : **`offset`**($\nabla\beta$) : For StableHLO information see [StableHLO - batch_norm_grad](https://openxla.org/stablehlo/spec#batch_norm_grad). ## BatchNormInference See also [`XlaBuilder::BatchNormInference`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h) and [the original batch normalization paper](https://arxiv.org/abs/1502.03167) for a detailed description of the algorithm. Normalizes an array across batch and spatial dimensions. **`BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)`** Arguments | Type | Semantics --------------- | ------- | --------------------------------------- `operand` | XlaOp | n dimensional array to be normalized `scale` | XlaOp | 1 dimensional array `offset` | XlaOp | 1 dimensional array `mean` | XlaOp | 1 dimensional array `variance` | XlaOp | 1 dimensional array `epsilon` | `float` | Epsilon value `feature_index` | `int64` | Index to feature dimension in `operand` For each feature in the feature dimension (`feature_index` is the index for the feature dimension in `operand`), the operation calculates the mean and variance across all the other dimensions and uses the mean and variance to normalize each element in `operand`. The `feature_index` must be a valid index for the feature dimension in `operand`. `BatchNormInference` is equivalent to calling `BatchNormTraining` without computing `mean` and `variance` for each batch. It uses the input `mean` and `variance` instead as estimated values. The purpose of this op is to reduce latency in inference, hence the name `BatchNormInference`. The output is an n-dimensional, normalized array with the same shape as input `operand`. For StableHLO information see [StableHLO - batch_norm_inference](https://openxla.org/stablehlo/spec#batch_norm_inference). ## BatchNormTraining See also [`XlaBuilder::BatchNormTraining`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h) and [`the original batch normalization paper`](https://arxiv.org/abs/1502.03167) for a detailed description of the algorithm. Normalizes an array across batch and spatial dimensions. **`BatchNormTraining(operand, scale, offset, epsilon, feature_index)`** Arguments | Type | Semantics --------------- | ------- | ---------------------------------------- `operand` | `XlaOp` | n dimensional array to be normalized (x) `scale` | `XlaOp` | 1 dimensional array ($\gamma$) `offset` | `XlaOp` | 1 dimensional array ($\beta$) `epsilon` | `float` | Epsilon value ($\epsilon$) `feature_index` | `int64` | Index to feature dimension in `operand` For each feature in the feature dimension (`feature_index` is the index for the feature dimension in `operand`), the operation calculates the mean and variance across all the other dimensions and uses the mean and variance to normalize each element in `operand`. The `feature_index` must be a valid index for the feature dimension in `operand`. The algorithm goes as follows for each batch in `operand` $x$ that contains `m` elements with `w` and `h` as the size of spatial dimensions (assuming `operand` is a 4 dimensional array): - Calculates batch mean $\mu_l$ for each feature `l` in feature dimension: $\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}$ - Calculates batch variance $\sigma^2_l$: $\sigma^2_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h (x_{ijkl} - \mu_l)^2$ - Normalizes, scales and shifts: $y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon}}+\beta_l$ The epsilon value, usually a small number, is added to avoid divide-by-zero errors. The output type is a tuple of three `XlaOp`s: | Outputs | Type | Semantics | | ------------ | ------- | ------------------------------------------------ | | `output` | `XlaOp` | n dimensional array with the same shape as input | : : : `operand` (y) : | `batch_mean` | `XlaOp` | 1 dimensional array ($\mu$) | | `batch_var` | `XlaOp` | 1 dimensional array ($\sigma^2$) | The `batch_mean` and `batch_var` are moments calculated across the batch and spatial dimensions using the formulas above. For StableHLO information see [StableHLO - batch_norm_training](https://openxla.org/stablehlo/spec#batch_norm_training). ## Bitcast See also [`HloInstruction::CreateBitcast`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `Bitcast` may appear in HLO dumps, but they are not intended to be constructed manually by end users. > **Note:** `Bitcast` is only found in HLO. It is not found in StableHLO. ## BitcastConvertType See also [`XlaBuilder::BitcastConvertType`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Similar to a `tf.bitcast` in TensorFlow, performs an element-wise bitcast operation from a data shape to a target shape. The input and output size must match: e.g. `s32` elements become `f32` elements via bitcast routine, and one `s32` element will become four `s8` elements. Bitcast is implemented as a low-level cast, so machines with different floating-point representations will give different results. **`BitcastConvertType(operand, new_element_type)`** Arguments | Type | Semantics ------------------ | --------------- | --------------------------- `operand` | `XlaOp` | array of type T with dims D `new_element_type` | `PrimitiveType` | type U The dimensions of the operand and the target shape must match, apart from the last dimension which will change by the ratio of the primitive size before and after the conversion. The source and destination element types must not be tuples. For StableHLO information see [StableHLO - bitcast_convert](https://openxla.org/stablehlo/spec#bitcast_convert). ### Bitcast-converting to primitive type of different width `BitcastConvert` HLO instruction supports the case where the size of the output element type `T'` is not equal to the size of the input element `T`. As the whole operation is conceptually a bitcast and does not change the underlying bytes, the shape of the output element has to change. For `B = sizeof(T), B' = sizeof(T')`, there are two possible cases. First, when `B > B'`, the output shape gets a new minor-most dimension of size `B/B'`. For example: ```cpp f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input) ``` The rule remains the same for effective scalars: ```cpp f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input) ``` Alternatively, for `B' > B` the instruction requires the last logical dimension of the input shape to be equal to `B'/B`, and this dimension is dropped during the conversion: ```cpp f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input) ``` Note that conversions between different bitwidths are not elementwise. ## Broadcast See also [`XlaBuilder::Broadcast`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Adds dimensions to an array by duplicating the data in the array. **`Broadcast(operand, broadcast_sizes)`** Arguments | Type | Semantics ----------------- | ------------------- | ------------------------------- `operand` | `XlaOp` | The array to duplicate `broadcast_sizes` | `ArraySlice` | The sizes of the new dimensions The new dimensions are inserted on the left, i.e. if `broadcast_sizes` has values `{a0, ..., aN}` and the operand shape has dimensions `{b0, ..., bM}` then the shape of the output has dimensions `{a0, ..., aN, b0, ..., bM}`. The new dimensions index into copies of the operand, i.e. ```cpp output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM] ``` For example, if `operand` is a scalar `f32` with value `2.0f`, and `broadcast_sizes` is `{2, 3}`, then the result will be an array with shape `f32[2, 3]` and all the values in the result will be `2.0f`. For StableHLO information see [StableHLO - broadcast](https://openxla.org/stablehlo/spec#broadcast). ### BroadcastInDim See also [`XlaBuilder::BroadcastInDim`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Expands the size and number of dimensions of an array by duplicating the data in the array. **`BroadcastInDim(operand, out_dim_size, broadcast_dimensions)`** | Arguments | Type | Semantics | | ---------------------- | ------------------- | ----------------------------- | | `operand` | `XlaOp` | The array to duplicate | | `out_dim_size` | `ArraySlice` | The sizes of the dimensions | : : : of the target shape : | `broadcast_dimensions` | `ArraySlice` | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : Similar to Broadcast, but allows adding dimensions anywhere and expanding existing dimensions with size 1. The `operand` is broadcast to the shape described by `out_dim_size`. `broadcast_dimensions` maps the dimensions of `operand` to the dimensions of the target shape, i.e. the i'th dimension of the operand is mapped to the broadcast_dimension\[i\]'th dimension of the output shape. The dimensions of `operand` must have size 1 or be the same size as the dimension in the output shape they are mapped to. The remaining dimensions are filled with dimensions of size 1. Degenerate-dimension broadcasting then broadcasts along these degenerate dimensions to reach the output shape. The semantics are described in detail on the [broadcasting page](broadcasting.md). ## Call See also [`XlaBuilder::Call`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Invokes a computation with the given arguments. **`Call(computation, operands...)`** | Arguments | Type | Semantics | | ------------- | ---------------------- | ----------------------------------- | | `computation` | `XlaComputation` | computation of type `T_0, T_1, ..., | : : : T_{N-1} -> S` with N parameters of : : : : arbitrary type : | `operands` | sequence of N `XlaOp`s | N arguments of arbitrary type | The arity and types of the `operands` must match the parameters of the `computation`. It is allowed to have no `operands`. ### CompositeCall See also [`XlaBuilder::CompositeCall`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Encapsulates an operation made up (composed) of other StableHLO operations, taking inputs and composite_attributes and producing results. The semantics of the op are implemented by the decomposition attribute. The composite op can be replaced with its decomposition without changing program semantics. In cases where inlining the decomposition does not provide the same op semantics, prefer using custom_call. The version field (defaults to 0) is used to denote when a composite's semantics change. This op is implemented as a `kCall` with attribute `is_composite=true`. The `decomposition` field is specified by the `computation` attribute. The frontend attributes store the remaining attributes prefixed with `composite.`. Example CompositeCall op: ```cpp f32[] call(f32[] %cst), to_apply=%computation, is_composite=true, frontend_attributes = { composite.name="foo.bar", composite.attributes={n = 1 : i32, tensor = dense<1> : tensor}, composite.version="1" } ``` **`CompositeCall(computation, operands..., name, attributes, version)`** | Arguments | Type | Semantics | | ------------- | ---------------------- | ----------------------------------- | | `computation` | `XlaComputation` | computation of type `T_0, T_1, ..., | : : : T_{N-1} -> S` with N parameters of : : : : arbitrary type : | `operands` | sequence of N `XlaOp`s | variadic number of values | | `name` | `string` | name of the composite | | `attributes` | optional `string` | optional stringified dictionary of | : : : attributes : | `version` | optional `int64` | number to version updates to | : : : semantics of the composite op : An op’s `decomposition` isn’t a field called, but instead appears as a to_apply attribute that points to the function which contains the lower-level implementation, i.e. `to_apply=%funcname` More information on composite and decomposition can be found on [StableHLO Specification](https://openxla.org/stablehlo/spec#composite). ## Cbrt See also [`XlaBuilder::Cbrt`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise cubic root operation `x -> cbrt(x)`. **`Cbrt(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Cbrt also supports the optional `result_accuracy` argument: **`Cbrt(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - cbrt](https://openxla.org/stablehlo/spec#cbrt). ## Ceil See also [`XlaBuilder::Ceil`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise ceil `x -> ⌈x⌉`. **`Ceil(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - ceil](https://openxla.org/stablehlo/spec#ceil). ## Cholesky See also [`XlaBuilder::Cholesky`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Computes the [Cholesky decomposition](https://en.wikipedia.org/wiki/Cholesky_decomposition) of a batch of symmetric (Hermitian) positive definite matrices. **`Cholesky(a, lower)`** | Arguments | Type | Semantics | | --------- | ------- | ----------------------------------------------------- | | `a` | `XlaOp` | an array of a complex or floating-point type with > 2 | : : : dimensions. : | `lower` | `bool` | whether to use the upper or lower triangle of `a`. | If `lower` is `true`, computes lower-triangular matrices `l` such that $a = l . l^T$. If `lower` is `false`, computes upper-triangular matrices `u` such that $a = u^T . u$. Input data is read only from the lower/upper triangle of `a`, depending on the value of `lower`. Values from the other triangle are ignored. Output data is returned in the same triangle; the values in the other triangle are implementation-defined and may be anything. If `a` has greater than 2 dimensions, `a` is treated as a batch of matrices, where all except the minor 2 dimensions are batch dimensions. If `a` is not symmetric (Hermitian) positive definite, the result is implementation-defined. For StableHLO information see [StableHLO - cholesky](https://openxla.org/stablehlo/spec#cholesky). ## Clamp See also [`XlaBuilder::Clamp`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Clamps an operand to within the range between a minimum and maximum value. **`Clamp(min, operand, max)`** Arguments | Type | Semantics --------- | ------- | --------------- `min` | `XlaOp` | array of type T `operand` | `XlaOp` | array of type T `max` | `XlaOp` | array of type T Given an operand and minimum and maximum values, returns the operand if it is in the range between the minimum and maximum, else returns the minimum value if the operand is below this range or the maximum value if the operand is above this range. That is, `clamp(a, x, b) = min(max(a, x), b)`. All three arrays must be the same shape. Alternatively, as a restricted form of [broadcasting](broadcasting.md), `min` and/or `max` can be a scalar of type `T`. Example with scalar `min` and `max`: ```cpp let operand: s32[3] = {-1, 5, 9}; let min: s32 = 0; let max: s32 = 6; ==> Clamp(min, operand, max) = s32[3]{0, 5, 6}; ``` For StableHLO information see [StableHLO - clamp](https://openxla.org/stablehlo/spec#clamp). ## Collapse See also [`XlaBuilder::Collapse`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). and the `tf.reshape` operation. Collapses dimensions of an array into one dimension. **`Collapse(operand, dimensions)`** Arguments | Type | Semantics ------------ | -------------- | ----------------------------------------------- `operand` | `XlaOp` | array of type T `dimensions` | `int64` vector | in-order, consecutive subset of T's dimensions. Collapse replaces the given subset of the operand's dimensions by a single dimension. The input arguments are an arbitrary array of type T and a compile-time-constant vector of dimension indices. The dimension indices must be an in-order (low to high dimension numbers), consecutive subset of T's dimensions. Thus, {0, 1, 2}, {0, 1}, or {1, 2} are all valid dimension sets, but {1, 0} or {0, 2} are not. They are replaced by a single new dimension, in the same position in the dimension sequence as those they replace, with the new dimension size equal to the product of original dimension sizes. The lowest dimension number in `dimensions` is the slowest varying dimension (most major) in the loop nest which collapses these dimensions, and the highest dimension number is fastest varying (most minor). See the `tf.reshape` operator if more general collapse ordering is needed. For example, let v be an array of 24 elements: ```cpp let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}}, {{20, 21, 22}, {25, 26, 27}}, {{30, 31, 32}, {35, 36, 37}}, {{40, 41, 42}, {45, 46, 47}}}; // Collapse to a single dimension, leaving one dimension. let v012 = Collapse(v, {0,1,2}); then v012 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27, 30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47}; // Collapse the two lower dimensions, leaving two dimensions. let v01 = Collapse(v, {0,1}); then v01 == f32[4x6] {{10, 11, 12, 15, 16, 17}, {20, 21, 22, 25, 26, 27}, {30, 31, 32, 35, 36, 37}, {40, 41, 42, 45, 46, 47}}; // Collapse the two higher dimensions, leaving two dimensions. let v12 = Collapse(v, {1,2}); then v12 == f32[8x3] {{10, 11, 12}, {15, 16, 17}, {20, 21, 22}, {25, 26, 27}, {30, 31, 32}, {35, 36, 37}, {40, 41, 42}, {45, 46, 47}}; ``` > **Note:** For further information on `Collapse` see [HLO - Reshape](#reshape) ## Clz See also [`XlaBuilder::Clz`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise count leading zeros. **`Clz(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function > **Note:** `Clz` is not found directly in StableHLO, but is analogous to > [StableHlo - count_leading_zeros](https://openxla.org/stablehlo/spec#count_leading_zeros). ## CollectiveBroadcast See also [`XlaBuilder::CollectiveBroadcast`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Broadcasts data across replicas. Data is sent from the first replica id in each group to the other ids in the same group. If a replica id is not in any replica group, the output on that replica is a tensor consisting of 0(s) in `shape`. **`CollectiveBroadcast(operand, replica_groups, channel_id)`** | Arguments | Type | Semantics | | ---------------- | ------------------------ | ----------------------------- | | `operand` | `XlaOp` | The operand to the function | | `replica_groups` | `ReplicaGroup`vector | Each group contains a list of | : : : replica ids : | `channel_id` | optional `ChannelHandle` | unique identifier for each | : : : send/recv pair : For StableHLO information see [StableHLO - collective_broadcast](https://openxla.org/stablehlo/spec#collective_broadcast). ## CollectivePermute See also [`XlaBuilder::CollectivePermute`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). CollectivePermute is a collective operation that sends and receives data across replicas. **`CollectivePermute(operand, source_target_pairs, channel_id)`** | Arguments | Type | Semantics | | --------------------- | ------------------------ | ------------------------- | | `operand` | `XlaOp` | n dimensional input array | | `source_target_pairs` | `` vector | A list of | : : : (source_replica_id, : : : : target_replica_id) pairs. : : : : For each pair, the : : : : operand is sent from : : : : source replica to target : : : : replica. : | `channel_id` | optional `ChannelHandle` | Optional channel ID for | : : : cross-module : : : : communication : Note that there are the following restrictions on the `source_target_pairs`: - Any two pairs should not have the same target replica id, and they should not have the same source replica id. - If a replica id is not a target in any pair, then the output on that replica is a tensor consisting of 0(s) with the same shape as the input. The API of `CollectivePermute` operation is internally decomposed into 2 HLO instructions (`CollectivePermuteStart` and `CollectivePermuteDone`). See also [`HloInstruction::CreateCollectivePermuteStart`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `CollectivePermuteStart` and `CollectivePermuteDone` serve as primitives in HLO. These ops may appear in HLO dumps, but they are not intended to be constructed manually by end users. For StableHLO information see [StableHLO - collective_permute](https://openxla.org/stablehlo/spec#collective_permute). > **Note:** `CollectivePermuteStart` and `CollectivePermuteDone` are only found > in HLO. They are not found in StableHLO. ## Compare See also [`XlaBuilder::Compare`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise comparison of `lhs` and `rhs` of the following: ### Eq See also [`XlaBuilder::Eq`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise **equal-to** comparison of `lhs` and `rhs`. $lhs = rhs$ **`Eq(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Eq: **`Eq(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). Support a total order over the floating point numbers exists for Eq, by enforcing: $$-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.$$ **`EqTotalOrder(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : For StableHLO information see [StableHLO - compare](https://openxla.org/stablehlo/spec#compare). ### Ne See also [`XlaBuilder::Ne`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise **not equal-to** comparison of `lhs` and `rhs`. $lhs != rhs$ **`Ne(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Ne: **`Ne(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). Support a total order over the floating point numbers exists for Ne, by enforcing: $$-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.$$ **`NeTotalOrder(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : For StableHLO information see [StableHLO - compare](https://openxla.org/stablehlo/spec#compare). ### Ge See also [`XlaBuilder::Ge`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise **greater-or-equal-than** comparison of `lhs` and `rhs`. $lhs >= rhs$ **`Ge(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Ge: **`Ge(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). Support a total order over the floating point numbers exists for Gt, by enforcing: $$-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.$$ **`GtTotalOrder(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : For StableHLO information see [StableHLO - compare](https://openxla.org/stablehlo/spec#compare). ### Gt See also [`XlaBuilder::Gt`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise **greater-than** comparison of `lhs` and `rhs`. $lhs > rhs$ **`Gt(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Gt: **`Gt(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - compare](https://openxla.org/stablehlo/spec#compare). ### Le See also [`XlaBuilder::Le`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise **less-or-equal-than** comparison of `lhs` and `rhs`. $lhs <= rhs$ **`Le(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Le: **`Le(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). Support a total order over the floating point numbers exists for Le, by enforcing: $$-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.$$ **`LeTotalOrder(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : For StableHLO information see [StableHLO - compare](https://openxla.org/stablehlo/spec#compare). ### Lt See also [`XlaBuilder::Lt`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise **less-than** comparison of `lhs` and `rhs`. $lhs < rhs$ **`Lt(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Lt: **`Lt(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). Support a total order over the floating point numbers exists for Lt, by enforcing: $$-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.$$ **`LtTotalOrder(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : For StableHLO information see [StableHLO - compare](https://openxla.org/stablehlo/spec#compare). ## Complex See also [`XlaBuilder::Complex`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise conversion to a complex value from a pair of real and imaginary values, `lhs` and `rhs`. **`Complex(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Complex: **`Complex(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - complex](https://openxla.org/stablehlo/spec#complex). ## ConcatInDim (Concatenate) See also [`XlaBuilder::ConcatInDim`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Concatenate composes an array from multiple array operands. The array has the same number of dimensions as each of the input array operands (which must have the same number of dimensions as each other) and contains the arguments in the order that they were specified. **`Concatenate(operands..., dimension)`** | Arguments | Type | Semantics | | ----------- | --------------------- | -------------------------------------- | | `operands` | sequence of N `XlaOp` | N arrays of type T with dimensions | : : : [L0, L1, ...]. Requires N >= 1. : | `dimension` | `int64` | A value in the interval `[0, N)` that | : : : names the dimension to be concatenated : : : : between the `operands`. : With the exception of `dimension` all dimensions must be the same. This is because XLA does not support "ragged" arrays. Also note that 0-dimensional values cannot be concatenated (as it's impossible to name the dimension along which the concatenation occurs). 1-dimensional example: ```cpp Concat({{2, 3}, {4, 5}, {6, 7}}, 0) //Output: {2, 3, 4, 5, 6, 7} ``` 2-dimensional example: ```cpp let a = {{1, 2}, {3, 4}, {5, 6}}; let b = {{7, 8}}; Concat({a, b}, 0) //Output: {{1, 2}, // {3, 4}, // {5, 6}, // {7, 8}} ``` Diagram: For StableHLO information see [StableHLO - concatenate](https://openxla.org/stablehlo/spec#concatenate). ## Conditional See also [`XlaBuilder::Conditional`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Conditional(predicate, true_operand, true_computation, false_operand, false_computation)`** | Arguments | Type | Semantics | | ------------------- | ---------------- | ---------------------------------- | | `predicate` | `XlaOp` | Scalar of type `PRED` | | `true_operand` | `XlaOp` | Argument of type $T_0$ | | `true_computation` | `XlaComputation` | XlaComputation of type $T_0 \to S$ | | `false_operand` | `XlaOp` | Argument of type $T_1$ | | `false_computation` | `XlaComputation` | XlaComputation of type $T_1 \to S$ | Executes `true_computation` if `predicate` is `true`, `false_computation` if `predicate` is `false`, and returns the result. The `true_computation` must take in a single argument of type $T_0$ and will be invoked with `true_operand` which must be of the same type. The `false_computation` must take in a single argument of type $T_1$ and will be invoked with `false_operand` which must be of the same type. The type of the returned value of `true_computation` and `false_computation` must be the same. Note that only one of `true_computation` and `false_computation` will be executed depending on the value of `predicate`. **`Conditional(branch_index, branch_computations, branch_operands)`** | Arguments | Type | Semantics | | --------------------- | ------------------------------ | ------------------------------------------------------------ | | `branch_index` | `XlaOp` | Scalar of type `S32` | | `branch_computations` | sequence of N `XlaComputation` | XlaComputations of type $T_0 \to S , T_1 \to S , ..., T_{N-1} \to S$ | | `branch_operands` | sequence of N `XlaOp` | Arguments of type $T_0 , T_1 , ..., T_{N-1}$ | Executes `branch_computations[branch_index]`, and returns the result. If `branch_index` is an `S32` which is < 0 or >= N, then `branch_computations[N-1]` is executed as the default branch. Each `branch_computations[b]` must take in a single argument of type $T_b$ and will be invoked with `branch_operands[b]` which must be of the same type. The type of the returned value of each `branch_computations[b]` must be the same. Note that only one of the `branch_computations` will be executed depending on the value of `branch_index`. For StableHLO information see [StableHLO - if](https://openxla.org/stablehlo/spec#if). ## Constant See also [`XlaBuilder::ConstantLiteral`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Produces an `output` from a constant `literal`. **`Constant(literal)`** Arguments | Type | Semantics --------- | -------------- | -------------------------------------- `literal` | `LiteralSlice` | constant view of an existing `Literal` For StableHLO information see [StableHLO - constant](https://openxla.org/stablehlo/spec#constant). ## ConvertElementType See also [`XlaBuilder::ConvertElementType`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Similar to an element-wise `static_cast` in C++, `ConvertElementType` performs an element-wise conversion operation from a data shape to a target shape. The dimensions must match, and the conversion is an element-wise one; e.g. `s32` elements become `f32` elements via an `s32`-to-`f32` conversion routine. **`ConvertElementType(operand, new_element_type)`** Arguments | Type | Semantics ------------------ | --------------- | --------------------------- `operand` | `XlaOp` | array of type T with dims D `new_element_type` | `PrimitiveType` | type U The dimensions of the operand and the target shape must match. The source and destination element types must not be tuples. A conversion such as `T=s32` to `U=f32` will perform a normalizing int-to-float conversion routine such as round-to-nearest-even. > **Note:** The precise float-to-int and visa-versa conversions are currently > unspecified, but may become additional arguments to the convert operation in > the future. Not all possible conversions have been implemented for all > targets. ```cpp let a: s32[3] = {0, 1, 2}; let b: f32[3] = convert(a, f32); then b == f32[3]{0.0, 1.0, 2.0} ``` For StableHLO information see [StableHLO - convert](https://openxla.org/stablehlo/spec#convert). ## Conv (Convolution) See also [`XlaBuilder::Conv`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Computes a convolution of the kind used in neural networks. Here, a convolution can be thought of as a n-dimensional window moving across a n-dimensional base area and a computation is performed for each possible position of the window. `Conv` Enqueues a convolution instruction onto the computation, which uses the default convolution dimension numbers with no dilation. The padding is specified in a short-hand way as either SAME or VALID. SAME padding pads the input (`lhs`) with zeroes so that the output has the same shape as the input when not taking striding into account. VALID padding simply means no padding. **`Conv(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)`** | Arguments | Type | Semantics | | :----------------------- | :------------------ | :-------------------------- | | `lhs` | `XlaOp` | (n+2)-dimensional array of | : : : inputs : | `rhs` | `XlaOp` | (n+2)-dimensional array of | : : : kernel weights : | `window_strides` | `ArraySlice` | n-d array of kernel strides | | `padding` | `Padding` | enum of padding | | `feature_group_count` | int64 | the number of feature | : : : groups : | `batch_group_count` | int64 | the number of batch groups | | `precision_config` | optional | enum for level of precision | : : `PrecisionConfig` : : | `preferred_element_type` | optional | enum of scalar element type | : : `PrimitiveType` : : Increasing levels of controls are available for `Conv`: - [ConvWithGeneralPadding](#ConvWithGeneralPadding) - [ConvWithGeneralDimensions](#ConvWithGeneralDimensions) - [ConvGeneral](#ConvGeneral) - [ConvGeneralDilated](#convgeneraldilated) Let n be the number of spatial dimensions. The `lhs` argument is an (n+2)-dimensional array describing the base area. This is called the input, even though of course the rhs is also an input. In a neural network, these are the input activations. The n+2 dimensions are, in this order: * `batch`: Each coordinate in this dimension represents an independent input for which convolution is carried out. * `z/depth/features`: Each (y,x) position in the base area has a vector associated to it, which goes into this dimension. * `spatial_dims`: Describes the `n` spatial dimensions that define the base area that the window moves across. The `rhs` argument is an (n+2)-dimensional array describing the convolutional filter/kernel/window. The dimensions are, in this order: * `output-z`: The `z` dimension of the output. * `input-z`: The size of this dimension times `feature_group_count` should equal the size of the `z` dimension in lhs. * `spatial_dims`: Describes the `n` spatial dimensions that define the n-d window that moves across the base area. The `window_strides` argument specifies the stride of the convolutional window in the spatial dimensions. For example, if the stride in the first spatial dimension is 3, then the window can only be placed at coordinates where the first spatial index is divisible by 3. The `padding` argument specifies the amount of zero padding to be applied to the base area. The amount of padding can be negative -- the absolute value of negative padding indicates the number of elements to remove from the specified dimension before doing the convolution. `padding[0]` specifies the padding for dimension `y` and `padding[1]` specifies the padding for dimension `x`. Each pair has the low padding as the first element and the high padding as the second element. The low padding is applied in the direction of lower indices while the high padding is applied in the direction of higher indices. For example, if `padding[1]` is `(2,3)` then there will be a padding by 2 zeroes on the left and by 3 zeroes on the right in the second spatial dimension. Using padding is equivalent to inserting those same zero values into the input (`lhs`) before doing the convolution. The `lhs_dilation` and `rhs_dilation` arguments specify the dilation factor to be applied to the lhs and rhs, respectively, in each spatial dimension. If the dilation factor in a spatial dimension is d, then d-1 holes are implicitly placed between each of the entries in that dimension, increasing the size of the array. The holes are filled with a no-op value, which for convolution means zeroes. Dilation of the rhs is also called atrous convolution. For more details, see `tf.nn.atrous_conv2d`. Dilation of the lhs is also called transposed convolution. For more details, see `tf.nn.conv2d_transpose`. The `feature_group_count` argument (default value 1) can be used for grouped convolutions. `feature_group_count` needs to be a divisor of both the input and the output feature dimension. If `feature_group_count` is greater than 1, it means that conceptually the input and output feature dimension and the `rhs` output feature dimension are split evenly into many `feature_group_count` groups, each group consisting of a consecutive subsequence of features. The input feature dimension of `rhs` needs to be equal to the `lhs` input feature dimension divided by `feature_group_count` (so it already has the size of a group of input features). The i-th groups are used together to compute `feature_group_count` for many separate convolutions. The results of these convolutions are concatenated together in the output feature dimension. For depthwise convolution the `feature_group_count` argument would be set to the input feature dimension, and the filter would be reshaped from `[filter_height, filter_width, in_channels, channel_multiplier]` to `[filter_height, filter_width, 1, in_channels * channel_multiplier]`. For more details, see `tf.nn.depthwise_conv2d`. The `batch_group_count` (default value 1) argument can be used for grouped filters during backpropagation. `batch_group_count` needs to be a divisor of the size of the `lhs` (input) batch dimension. If `batch_group_count` is greater than 1, it means that the output batch dimension should be of size `input batch / batch_group_count`. The `batch_group_count` must be a divisor of the output feature size. The output shape has these dimensions, in this order: * `batch`: The size of this dimension times `batch_group_count` should equal the size of the `batch` dimension in lhs. * `z`: Same size as `output-z` on the kernel (`rhs`). * `spatial_dims`: One value for each valid placement of the convolutional window. The figure above shows how the `batch_group_count` field works. Effectively, we slice each lhs batch into `batch_group_count` groups, and do the same for the output features. Then, for each of these groups we do pairwise convolutions and concatenate the output along the output feature dimension. The operational semantics of all the other dimensions (feature and spatial) remain the same. The valid placements of the convolutional window are determined by the strides and the size of the base area after padding. To describe what a convolution does, consider a 2d convolution, and pick some fixed `batch`, `z`, `y`, `x` coordinates in the output. Then `(y,x)` is a position of a corner of the window within the base area (e.g. the upper left corner, depending on how you interpret the spatial dimensions). We now have a 2d window, taken from the base area, where each 2d point is associated to a 1d vector, so we get a 3d box. From the convolutional kernel, since we fixed the output coordinate `z`, we also have a 3d box. The two boxes have the same dimensions, so we can take the sum of the element-wise products between the two boxes (similar to a dot product). That is the output value. Note that if `output-z` is e.g., 5, then each position of the window produces 5 values in the output into the `z` dimension of the output. These values differ in what part of the convolutional kernel is used - there is a separate 3d box of values used for each `output-z` coordinate. So you could think of it as 5 separate convolutions with a different filter for each of them. Here is pseudo-code for a 2d convolution with padding and striding: ```cpp for (b, oz, oy, ox) { // output coordinates value = 0; for (iz, ky, kx) { // kernel coordinates and input z iy = oy*stride_y + ky - pad_low_y; ix = ox*stride_x + kx - pad_low_x; if ((iy, ix) inside the base area considered without padding) { value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx); } } output(b, oz, oy, ox) = value; } ``` `precision_config` is used to indicate the precision configuration. The level dictates whether hardware should attempt to generate more machine code instructions to provide more accurate dtype emulation when needed (i.e. emulating f32 on a TPU that only supports bf16 matmuls). Values may be `DEFAULT`, `HIGH`, `HIGHEST`. Additional details [in the MXU sections](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus). `preferred_element_type` is a scalar element of higher/lower precision output types used for accumulation. `preferred_element_type` recommends the accumulation type for the given operation, however it is not guaranteed. This allows for some hardware backends to instead accumulate in a different type and convert to the preferred output type. For StableHLO information see [StableHLO - convolution](https://openxla.org/stablehlo/spec#convolution). ### ConvWithGeneralPadding See also [`XlaBuilder::ConvWithGeneralPadding`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`ConvWithGeneralPadding(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)`** Same as [`Conv`](#conv-convolution) where padding configuration is explicit. | Arguments | Type | Semantics | | :----------------------- | :------------------ | :-------------------------- | | `lhs` | `XlaOp` | (n+2)-dimensional array of | : : : inputs : | `rhs` | `XlaOp` | (n+2)-dimensional array of | : : : kernel weights : | `window_strides` | `ArraySlice` | n-d array of kernel strides | | `padding` | `ArraySlice< | n-d array of (low, high) | : : pair>` : padding : | `feature_group_count` | int64 | the number of feature | : : : groups : | `batch_group_count` | int64 | the number of batch groups | | `precision_config` | optional | enum for level of precision | : : `PrecisionConfig` : : | `preferred_element_type` | optional | enum of scalar element type | : : `PrimitiveType` : : ### ConvWithGeneralDimensions See also [`XlaBuilder::ConvWithGeneralDimensions`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`ConvWithGeneralDimensions(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)`** Same as [`Conv`](#conv-convolution) where dimension numbers are explicit. | Arguments | Type | Semantics | | :----------------------- | :---------------------------- | :---------------- | | `lhs` | `XlaOp` | (n+2)-dimensional | : : : array of inputs : | `rhs` | `XlaOp` | (n+2)-dimensional | : : : array of kernel : : : : weights : | `window_strides` | `ArraySlice` | n-d array of | : : : kernel strides : | `padding` | `Padding` | enum of padding | | `dimension_numbers` | `ConvolutionDimensionNumbers` | the number of | : : : dimensions : | `feature_group_count` | int64 | the number of | : : : feature groups : | `batch_group_count` | int64 | the number of | : : : batch groups : | `precision_config` | optional `PrecisionConfig` | enum for level of | : : : precision : | `preferred_element_type` | optional `PrimitiveType` | enum of scalar | : : : element type : ### ConvGeneral See also [`XlaBuilder::ConvGeneral`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`ConvGeneral(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)`** Same as [`Conv`](#conv-convolution) where dimension numbers and padding configuration is explicit | Arguments | Type | Semantics | | :----------------------- | :---------------------------- | :---------------- | | `lhs` | `XlaOp` | (n+2)-dimensional | : : : array of inputs : | `rhs` | `XlaOp` | (n+2)-dimensional | : : : array of kernel : : : : weights : | `window_strides` | `ArraySlice` | n-d array of | : : : kernel strides : | `padding` | `ArraySlice< | n-d array of | : : pair>` : (low, high) : : : : padding : | `dimension_numbers` | `ConvolutionDimensionNumbers` | the number of | : : : dimensions : | `feature_group_count` | int64 | the number of | : : : feature groups : | `batch_group_count` | int64 | the number of | : : : batch groups : | `precision_config` | optional `PrecisionConfig` | enum for level of | : : : precision : | `preferred_element_type` | optional `PrimitiveType` | enum of scalar | : : : element type : ### ConvGeneralDilated See also [`XlaBuilder::ConvGeneralDilated`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`ConvGeneralDilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type, window_reversal)`** Same as [`Conv`](#conv-convolution) where padding configuration, dilation factors, and dimension numbers are explicit. | Arguments | Type | Semantics | | :----------------------- | :---------------------------- | :---------------- | | `lhs` | `XlaOp` | (n+2)-dimensional | : : : array of inputs : | `rhs` | `XlaOp` | (n+2)-dimensional | : : : array of kernel : : : : weights : | `window_strides` | `ArraySlice` | n-d array of | : : : kernel strides : | `padding` | `ArraySlice< | n-d array of | : : pair>` : (low, high) : : : : padding : | `lhs_dilation` | `ArraySlice` | n-d lhs dilation | : : : factor array : | `rhs_dilation` | `ArraySlice` | n-d rhs dilation | : : : factor array : | `dimension_numbers` | `ConvolutionDimensionNumbers` | the number of | : : : dimensions : | `feature_group_count` | int64 | the number of | : : : feature groups : | `batch_group_count` | int64 | the number of | : : : batch groups : | `precision_config` | optional `PrecisionConfig` | enum for level of | : : : precision : | `preferred_element_type` | optional `PrimitiveType` | enum of scalar | : : : element type : | `window_reversal` | optional `vector` | flag used to | : : : logically reverse : : : : dimension before : : : : applying the : : : : convolution : ## Copy See also [`HloInstruction::CreateCopyStart`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `Copy` is internally decomposed into 2 HLO instructions `CopyStart` and `CopyDone`. `Copy` along with `CopyStart` and `CopyDone` serve as primitives in HLO. These ops may appear in HLO dumps, but they are not intended to be constructed manually by end users. > **Note:** `Copy`, `CopyStart`, `CopyDone` are only found in HLO. They are not > found in StableHLO. ## Cos See also[`XlaBuilder::Cos`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise cosine `x -> cos(x)`. **`Cos(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Cos also supports the optional `result_accuracy` argument: **`Cos(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - cosine](https://openxla.org/stablehlo/spec#cosine). ## Cosh See also [`XlaBuilder::Cosh`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise hyperbolic cosine `x -> cosh(x)`. **`Cosh(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Cosh also supports the optional `result_accuracy` argument: **`Cosh(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). > **Note:** `Cosh` is only found in HLO and not found in StableHLO. CHLO `Cosh` > in Frameworks will lower to HLO `Cosh` see > [StableHLO - chlo.cosh](https://openxla.org/stablehlo/generated/chlo?hl=en#chlocosh_chlocoshop) ## CustomCall See also [`XlaBuilder::CustomCall`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Call a user-provided function within a computation. CustomCall documentation is provided in [Developer details - XLA Custom Calls](https://openxla.org/xla/custom_call) For StableHLO information see [StableHLO - custom_call](https://openxla.org/stablehlo/spec#custom_call). ## Div See also [`XlaBuilder::Div`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise division of dividend `lhs` and divisor `rhs`. **`Div(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T Integer division overflow (signed/unsigned division/remainder by zero or signed division/remainder of `INT_SMIN` with `-1`) produces an implementation defined value. The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Div: **`Div(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - divide](https://openxla.org/stablehlo/spec#divide). ## Domain See also [`HloInstruction::CreateDomain`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `Domain` may appear in HLO dumps, but it is not intended to be constructed manually by end users. > **Note:** `Domain` is only found in HLO. It is not found in StableHLO. ## Dot See also [`XlaBuilder::Dot`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Dot(lhs, rhs, precision_config, preferred_element_type)`** | Arguments | Type | Semantics | | :----------------------- | :---------------- | :-------------------------- | | `lhs` | `XlaOp` | array of type T | | `rhs` | `XlaOp` | array of type T | | `precision_config` | optional | enum for level of precision | : : `PrecisionConfig` : : | `preferred_element_type` | optional | enum of scalar element type | : : `PrimitiveType` : : The exact semantics of this operation depend on the ranks of the operands: | Input | Output | Semantics | | :-------------------------- | :------------- | :--------------------------- | | vector [n] `dot` vector [n] | scalar | vector dot product | | matrix [m x k] `dot` vector | vector [m] | matrix-vector multiplication | : [k] : : : | matrix [m x k] `dot` matrix | matrix [m x n] | matrix-matrix multiplication | : [k x n] : : : The operation performs sum of products over the second dimension of `lhs` (or the first if it has 1 dimension) and the first dimension of `rhs`. These are the "contracted" dimensions. The contracted dimensions of `lhs` and `rhs` must be of the same size. In practice, it can be used to perform dot products between vectors, vector/matrix multiplications or matrix/matrix multiplications. `precision_config` is used to indicate the precision configuration. The level dictates whether hardware should attempt to generate more machine code instructions to provide more accurate dtype emulation when needed (i.e. emulating f32 on a TPU that only supports bf16 matmuls). Values may be `DEFAULT`, `HIGH`, `HIGHEST`. Additional details [in the MXU sections](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus). `preferred_element_type` is a scalar element of higher/lower precision output types used for accumulation. `preferred_element_type` recommends the accumulation type for the given operation, however it is not guaranteed. This allows for some hardware backends to instead accumulate in a different type and convert to the preferred output type. For StableHLO information see [StableHLO - dot](https://openxla.org/stablehlo/spec#dot). ### DotGeneral See also [`XlaBuilder::DotGeneral`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`DotGeneral(lhs, rhs, dimension_numbers, precision_config, preferred_element_type)`** | Arguments | Type | Semantics | | :----------------------- | :-------------------- | :--------------------- | | `lhs` | `XlaOp` | array of type T | | `rhs` | `XlaOp` | array of type T | | `dimension_numbers` | `DotDimensionNumbers` | contracting and batch | : : : dimension numbers : | `precision_config` | optional | enum for level of | : : `PrecisionConfig` : precision : | `preferred_element_type` | optional | enum of scalar element | : : `PrimitiveType` : type : Similar to Dot, but allows contracting and batch dimension numbers to be specified for both the `lhs` and `rhs`. | DotDimensionNumbers Fields | Type | Semantics | | :--------------------------- | :------------- | :-------------------------- | | `lhs_contracting_dimensions` | repeated int64 | `lhs` contracting dimension | : : : numbers : | `rhs_contracting_dimensions` | repeated int64 | `rhs` contracting dimension | : : : numbers : | `lhs_batch_dimensions` | repeated int64 | `lhs` batch dimension | : : : numbers : | `rhs_batch_dimensions` | repeated int64 | `rhs` batch dimension | : : : numbers : DotGeneral performs the sum of products over contracting dimensions specified in `dimension_numbers`. Associated contracting dimension numbers from the `lhs` and `rhs` do not need to be the same but must have the same dimension sizes. Example with contracting dimension numbers: ```cpp lhs = { {1.0, 2.0, 3.0}, {4.0, 5.0, 6.0} } rhs = { {1.0, 1.0, 1.0}, {2.0, 2.0, 2.0} } DotDimensionNumbers dnums; dnums.add_lhs_contracting_dimensions(1); dnums.add_rhs_contracting_dimensions(1); DotGeneral(lhs, rhs, dnums) -> { { 6.0, 12.0}, {15.0, 30.0} } ``` Associated batch dimension numbers from the `lhs` and `rhs` must have the same dimension sizes. Example with batch dimension numbers (batch size 2, 2x2 matrices): ```cpp lhs = { { {1.0, 2.0}, {3.0, 4.0} }, { {5.0, 6.0}, {7.0, 8.0} } } rhs = { { {1.0, 0.0}, {0.0, 1.0} }, { {1.0, 0.0}, {0.0, 1.0} } } DotDimensionNumbers dnums; dnums.add_lhs_contracting_dimensions(2); dnums.add_rhs_contracting_dimensions(1); dnums.add_lhs_batch_dimensions(0); dnums.add_rhs_batch_dimensions(0); DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0}, {3.0, 4.0} }, { {5.0, 6.0}, {7.0, 8.0} } } ``` Input | Output | Semantics ----------------------------------- | -------------- | ------------ [b0, m, k] `dot` [b0, k, n] | [b0, m, n] | batch matmul [b0, b1, m, k] `dot` [b0, b1, k, n] | [b0, b1, m, n] | batch matmul It follows that the resulting dimension number starts with the batch dimension, then the `lhs` non-contracting/non-batch dimension, and finally the `rhs` non-contracting/non-batch dimension. `precision_config` is used to indicate the precision configuration. The level dictates whether hardware should attempt to generate more machine code instructions to provide more accurate dtype emulation when needed (i.e. emulating f32 on a TPU that only supports bf16 matmuls). Values may be `DEFAULT`, `HIGH`, `HIGHEST`. Additional details [can be found in the MXU sections](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus). `preferred_element_type` is a scalar element of higher/lower precision output types used for accumulation. `preferred_element_type` recommends the accumulation type for the given operation, however it is not guaranteed. This allows for some hardware backends to instead accumulate in a different type and convert to the preferred output type. For StableHLO information see [StableHLO - dot_general](https://openxla.org/stablehlo/spec#dot_general). ### ScaledDot See also [`XlaBuilder::ScaledDot`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`ScaledDot(lhs, lhs_scale, rhs, rhs_scale, dimension_number, precision_config,preferred_element_type)`** | Arguments | Type | Semantics | | ------------------------ | ------------------------- | --------------------- | | `lhs` | `XlaOp` | array of type T | | `rhs` | `XlaOp` | array of type T | | `lhs_scale` | `XlaOp` | array of type T | | `rhs_scale` | `XlaOp` | array of type T | | `dimension_number` | `ScatterDimensionNumbers` | Dimension numbers for | : : : scatter operation : | `precision_config` | `PrecisionConfig` | enum for level of | : : : precision : | `preferred_element_type` | optional `PrimitiveType` | enum of scalar | : : : element type : Similar to [DotGeneral](#dotgeneral). Creates a scaled dot op with operands 'lhs', 'lhs_scale', 'rhs', and 'rhs_scale', with contracting and batch dimensions specified in 'dimension_numbers'. > **Note:** `ScaledDot` is only found in HLO. It is not found in StableHLO. ### RaggedDot See also [`XlaBuilder::RaggedDot`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). For a breakdown of `RaggedDot` computation see [StableHLO - chlo.ragged_dot](https://openxla.org/stablehlo/generated/chlo?hl=en#chloragged_dot_chloraggeddotop) > **Note:** `RaggedDot` is only found in HLO. It is not found in StableHLO. ## DynamicReshape See also [`XlaBuilder::DynamicReshape`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). This operation is functionally identical to [reshape](#reshape), but the result shape is specified dynamically via output_shape. **`DynamicReshape(operand, dim_sizes, new_size_bounds, dims_are_dynamic)`** Arguments | Type | Semantics ------------------ | ----------------- | ------------------------------ `operand` | `XlaOp` | N dimensional array of type T `dim_sizes` | vector of `XlaOP` | N dimensional vector sizes `new_size_bounds` | vector of `int63` | N dimensional vector of bounds `dims_are_dynamic` | vector of `bool` | N dimensional dynamic dim For StableHLO information see [StableHLO - dynamic_reshape](https://openxla.org/stablehlo/spec#dynamic_reshape). ## DynamicSlice See also [`XlaBuilder::DynamicSlice`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). DynamicSlice extracts a sub-array from the input array at dynamic `start_indices`. The size of the slice in each dimension is passed in `size_indices`, which specify the end point of exclusive slice intervals in each dimension: [start, start + size). The shape of `start_indices` must be 1-dimensional, with dimension size equal to the number of dimensions of `operand`. **`DynamicSlice(operand, start_indices, slice_sizes)`** | Arguments | Type | Semantics | | --------------- | --------------------- | ---------------------------------- | | `operand` | `XlaOp` | N dimensional array of type T | | `start_indices` | sequence of N `XlaOp` | List of N scalar integers | : : : containing the starting indices of : : : : the slice for each dimension. : : : : Value must be greater than or : : : : equal to zero. : | `size_indices` | `ArraySlice` | List of N integers containing the | : : : slice size for each dimension. : : : : Each value must be strictly : : : : greater than zero, and start + : : : : size must be less than or equal to : : : : the size of the dimension to avoid : : : : wrapping modulo dimension size. : The effective slice indices are computed by applying the following transformation for each index `i` in `[1, N)` before performing the slice: ```cpp start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - slice_sizes[i]) ``` This ensures that the extracted slice is always in-bounds with respect to the operand array. If the slice is in-bounds before the transformation is applied, the transformation has no effect. 1-dimensional example: ```cpp let a = {0.0, 1.0, 2.0, 3.0, 4.0}; let s = {2}; DynamicSlice(a, s, {2}); // Result: {2.0, 3.0} ``` 2-dimensional example: ```cpp let b = { {0.0, 1.0, 2.0}, {3.0, 4.0, 5.0}, {6.0, 7.0, 8.0}, {9.0, 10.0, 11.0} } let s = {2, 1} DynamicSlice(b, s, {2, 2}); //Result: // { { 7.0, 8.0}, // {10.0, 11.0} } ``` For StableHLO information see [StableHLO - dynamic_slice](https://openxla.org/stablehlo/spec#dynamic_slice). ## DynamicUpdateSlice See also [`XlaBuilder::DynamicUpdateSlice`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). DynamicUpdateSlice generates a result which is the value of the input array `operand`, with a slice `update` overwritten at `start_indices`. The shape of `update` determines the shape of the sub-array of the result which is updated. The shape of `start_indices` must be 1-dimensional, with dimension size equal to the number of dimensions of `operand`. **`DynamicUpdateSlice(operand, update, start_indices)`** | Arguments | Type | Semantics | | --------------- | --------------------- | ---------------------------------- | | `operand` | `XlaOp` | N dimensional array of type T | | `update` | `XlaOp` | N dimensional array of type T | : : : containing the slice update. Each : : : : dimension of update shape must be : : : : strictly greater than zero, and : : : : start + update must be less than : : : : or equal to the operand size for : : : : each dimension to avoid generating : : : : out-of-bounds update indices. : | `start_indices` | sequence of N `XlaOp` | List of N scalar integers | : : : containing the starting indices of : : : : the slice for each dimension. : : : : Value must be greater than or : : : : equal to zero. : The effective slice indices are computed by applying the following transformation for each index `i` in `[1, N)` before performing the slice: ```cpp start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i]) ``` This ensures that the updated slice is always in-bounds with respect to the operand array. If the slice is in-bounds before the transformation is applied, the transformation has no effect. 1-dimensional example: ```cpp let a = {0.0, 1.0, 2.0, 3.0, 4.0} let u = {5.0, 6.0} let s = {2} DynamicUpdateSlice(a, u, s) // Result: {0.0, 1.0, 5.0, 6.0, 4.0} ``` 2-dimensional example: ```cpp let b = { {0.0, 1.0, 2.0}, {3.0, 4.0, 5.0}, {6.0, 7.0, 8.0}, {9.0, 10.0, 11.0} } let u = { {12.0, 13.0}, {14.0, 15.0}, {16.0, 17.0} } let s = {1, 1} DynamicUpdateSlice(b, u, s) // Result: // { {0.0, 1.0, 2.0}, // {3.0, 12.0, 13.0}, // {6.0, 14.0, 15.0}, // {9.0, 16.0, 17.0} } ``` For StableHLO information see [StableHLO - dynamic_update_slice](https://openxla.org/stablehlo/spec#dynamic_update_slice). ## Erf See also [`XlaBuilder::Erf`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise error function `x -> erf(x)` where: $$\text{erf}(x) = \frac{2}{\sqrt{\pi}}\int_0^x e^{-t^2} \, dt$$. **`Erf(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Erf also supports the optional `result_accuracy` argument: **`Erf(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). > **Note:** `Erf` is only found in HLO and not found in StableHLO. CHLO `Erf` in > Frameworks will lower to HLO `Erf` see > [StableHLO - chlo.erf](https://openxla.org/stablehlo/generated/chlo?hl=en#chloerf_chloerfop) ## Exp See also [`XlaBuilder::Exp`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise natural exponential `x -> e^x`. **`Exp(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Exp also supports the optional `result_accuracy` argument: **`Exp(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - exponential](https://openxla.org/stablehlo/spec#exponential). ## Expm1 See also [`XlaBuilder::Expm1`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise natural exponential minus one `x -> e^x - 1`. **`Expm1(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Expm1 also supports the optional `result_accuracy` argument: **`Expm1(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - exponential_minus_one](https://openxla.org/stablehlo/spec#exponential_minus_one). ## Fft See also [`XlaBuilder::Fft`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). The XLA FFT operation implements the forward and inverse Fourier Transforms for real and complex inputs/outputs. Multidimensional FFTs on up to 3 axes are supported. **`Fft(operand, ftt_type, fft_length)`** | Arguments | Type | Semantics | | ------------ | ------------------- | ------------------------ | | `operand` | `XlaOp` | The array we are Fourier | : : : transforming. : | `fft_type` | `FftType` | See the table below. | | `fft_length` | `ArraySlice` | The time-domain lengths | : : : of the axes being : : : : transformed. This is : : : : needed in particular for : : : : IRFFT to right-size the : : : : innermost axis, since : : : : `RFFT(fft_length=[16])` : : : : has the same output : : : : shape as : : : : `RFFT(fft_length=[17])`. : | `FftType` | Semantics | | --------- | ---------------------------------------------------------------- | | `FFT` | Forward complex-to-complex FFT. Shape is unchanged. | | `IFFT` | Inverse complex-to-complex FFT. Shape is unchanged. | | `RFFT` | Forward real-to-complex FFT. Shape of the innermost axis is | : : reduced to `fft_length[-1] // 2 + 1` if `fft_length[-1]` is a : : : non-zero value, omitting the reversed conjugate part of the : : : transformed signal beyond the Nyquist frequency. : | `IRFFT` | Inverse real-to-complex FFT (i.e. takes complex, returns real). | : : Shape of the innermost axis is expanded to `fft_length[-1]` if : : : `fft_length[-1]` is a non-zero value, inferring the part of the : : : transformed signal beyond the Nyquist frequency from the reverse : : : conjugate of the `1` to `fft_length[-1] // 2 + 1` entries. : For StableHLO information see [StableHLO - fft](https://openxla.org/stablehlo/spec#fft). #### Multidimensional FFT When more than 1 `fft_length` is provided, this is equivalent to applying a cascade of FFT operations to each of the innermost axes. Note that for the real->complex and complex->real cases, the innermost axis transform is (effectively) performed first (RFFT; last for IRFFT), which is why the innermost axis is the one which changes size. Other axis transforms will then be complex->complex. #### Implementation details CPU FFT is backed by Eigen's TensorFFT. GPU FFT uses cuFFT. ## Floor See also [`XlaBuilder::Floor`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise floor `x -> ⌊x⌋`. **`Floor(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - floor](https://openxla.org/stablehlo/spec#floor). ## Fusion See also [`HloInstruction::CreateFusion`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). `Fusion` operation represents HLO instructions and serves as a primitive in HLO. This op may appear in HLO dumps but is not intended to be constructed manually by end users. > **Note:** `Fusion` is only found in HLO. It is not found in StableHLO. ## Gather The XLA gather operation stitches together several slices (each slice at a potentially different runtime offset) of an input array. For StableHLO information see [StableHLO - gather](https://openxla.org/stablehlo/spec#gather). ### General Semantics See also [`XlaBuilder::Gather`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). For a more intuitive description, see the "Informal Description" section below. **`gather(operand, start_indices, dimension_numbers, slice_sizes, indices_are_sorted)`** | Arguments | Type | Semantics | | -------------------- | ------------------------ | -------------------------- | | `operand` | `XlaOp` | The array we’re gathering | : : : from. : | `start_indices` | `XlaOp` | Array containing the | : : : starting indices of the : : : : slices we gather. : | `dimension_numbers` | `GatherDimensionNumbers` | The dimension in | : : : `start_indices` that : : : : "contains" the starting : : : : indices. See below for a : : : : detailed description. : | `slice_sizes` | `ArraySlice` | `slice_sizes[i]` is the | : : : bounds for the slice on : : : : dimension `i`. : | `indices_are_sorted` | `bool` | Whether the indices are | : : : guaranteed to be sorted by : : : : the caller. : For convenience, we label dimensions in the output array not in `offset_dims` as `batch_dims`. The output is an array with `batch_dims.size` + `offset_dims.size` dimensions. The `operand.rank` must equal the sum of `offset_dims.size` and `collapsed_slice_dims.size`. Also, `slice_sizes.size` has to be equal to `operand.rank`. If `index_vector_dim` is equal to `start_indices.rank` we implicitly consider `start_indices` to have a trailing `1` dimension (i.e. if `start_indices` was of shape `[6,7]` and `index_vector_dim` is `2` then we implicitly consider the shape of `start_indices` to be `[6,7,1]`). The bounds for the output array along dimension `i` is computed as follows: 1. If `i` is present in `batch_dims` (i.e. is equal to `batch_dims[k]` for some `k`) then we pick the corresponding dimension bounds out of `start_indices.shape`, skipping `index_vector_dim` (i.e. pick `start_indices.shape.dims`[`k`] if `k` < `index_vector_dim` and `start_indices.shape.dims`[`k`+`1`] otherwise). 2. If `i` is present in `offset_dims` (i.e. equal to `offset_dims`[`k`] for some `k`) then we pick the corresponding bound out of `slice_sizes` after accounting for `collapsed_slice_dims` (i.e. we pick `adjusted_slice_sizes`[`k`] where `adjusted_slice_sizes` is `slice_sizes` with the bounds at indices `collapsed_slice_dims` removed). Formally, the operand index `In` corresponding to a given output index `Out` is calculated as follows: 1. Let `G` = { `Out`[`k`] for `k` in `batch_dims` }. Use `G` to slice out a vector `S` such that `S`[`i`] = `start_indices`[Combine(`G`, `i`)] where Combine(A, b) inserts b at position `index_vector_dim` into A. Note that this is well defined even if `G` is empty: If `G` is empty then `S` = `start_indices`. 2. Create a starting index, `S``in`, into `operand` using `S` by scattering `S` using `start_index_map`. More precisely: 1. `S``in`[`start_index_map`[`k`]] = `S`[`k`] if `k` < `start_index_map.size`. 2. `S``in`[`_`] = `0` otherwise. 3. Create an index `O``in` into `operand` by scattering the indices at the offset dimensions in `Out` according to the `collapsed_slice_dims` set. More precisely: 1. `O``in`[`remapped_offset_dims`(`k`)] = `Out`[`offset_dims`[`k`]] if `k` < `offset_dims.size` (`remapped_offset_dims` is defined below). 2. `O``in`[`_`] = `0` otherwise. 4. `In` is `O``in` + `S``in` where + is element-wise addition. `remapped_offset_dims` is a monotonic function with domain [`0`, `offset_dims.size`) and range [`0`, `operand.rank`) \ `collapsed_slice_dims`. So if, e.g., `offset_dims.size` is `4`, `operand.rank` is `6` and `collapsed_slice_dims` is {`0`, `2`} then `remapped_offset_dims` is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}. If `indices_are_sorted` is set to true then XLA can assume that `start_indices` are sorted (in ascending order, _after_ scattering its values according to `start_index_map`) by the user. If they are not then the semantics are implementation defined. ### Informal Description and Examples Informally, every index `Out` in the output array corresponds to an element `E` in the operand array, computed as follows: - We use the batch dimensions in `Out` to look up a starting index from `start_indices`. - We use `start_index_map` to map the starting index (whose size may be less than operand.rank) to a "full" starting index into the `operand`. - We dynamic-slice out a slice with size `slice_sizes` using the full starting index. - We reshape the slice by collapsing the `collapsed_slice_dims` dimensions. Since all collapsed slice dimensions must have a bound of 1, this reshape is always legal. - We use the offset dimensions in `Out` to index into this slice to get the input element, `E`, corresponding to output index `Out`. `index_vector_dim` is set to `start_indices.rank` - `1` in all of the examples that follow. More interesting values for `index_vector_dim` do not change the operation fundamentally, but make the visual representation more cumbersome. To get an intuition on how all of the above fits together, let's look at an example that gathers 5 slices of shape `[8,6]` from a `[16,11]` array. The position of a slice into the `[16,11]` array can be represented as an index vector of shape `S64[2]`, so the set of 5 positions can be represented as a `S64[5,2]` array. The behavior of the gather operation can then be depicted as an index transformation that takes [`G`,`O``0`,`O``1`], an index in the output shape, and maps it to an element in the input array in the following way: We first select an (`X`,`Y`) vector from the gather indices array using `G`. The element in the output array at index [`G`,`O``0`,`O``1`] is then the element in the input array at index [`X`+`O``0`,`Y`+`O``1`]. `slice_sizes` is `[8,6]`, which decides the range of O`0` and O`1`, and this in turn decides the bounds of the slice. This gather operation acts as a batch dynamic slice with `G` as the batch dimension. The gather indices may be multidimensional. For instance, a more general version of the example above using a "gather indices" array of shape `[4,5,2]` would translate indices like this: Again, this acts as a batch dynamic slice `G``0` and `G``1` as the batch dimensions. The slice size is still `[8,6]`. The gather operation in XLA generalizes the informal semantics outlined above in the following ways: 1. We can configure which dimensions in the output shape are the offset dimensions (dimensions containing `O``0`, `O``1` in the last example). The output batch dimensions (dimensions containing `G``0`, `G``1` in the last example) are defined to be the output dimensions that are not offset dimensions. 2. The number of output offset dimensions explicitly present in the output shape may be smaller than the input number of dimensions. These "missing" dimensions, which are listed explicitly as `collapsed_slice_dims`, must have a slice size of `1`. Since they have a slice size of `1` the only valid index for them is `0` and eliding them does not introduce ambiguity. 3. The slice extracted from the "Gather Indices" array ((`X`, `Y`) in the last example) may have fewer elements than the input array's number of dimensions, and an explicit mapping dictates how the index should be expanded to have the same number of dimensions as the input. As a final example, we use (2) and (3) to implement `tf.gather_nd`: `G``0` and `G``1` are used to slice out a starting index from the gather indices array as usual, except the starting index has only one element, `X`. Similarly, there is only one output offset index with the value `O``0`. However, before being used as indices into the input array, these are expanded in accordance to "Gather Index Mapping" (`start_index_map` in the formal description) and "Offset Mapping" (`remapped_offset_dims` in the formal description) into [`X`,`0`] and [`0`,`O``0`] respectively, adding up to [`X`,`O``0`]. In other words, the output index [`G``0`,`G``1`,`O``0`] maps to the input index [`GatherIndices`[`G``0`,`G``1`,`0`],`O``0`] which gives us the semantics for `tf.gather_nd`. `slice_sizes` for this case is `[1,11]`. Intuitively this means that every index `X` in the gather indices array picks an entire row and the result is the concatenation of all these rows. ## GetDimensionSize See also [`XlaBuilder::GetDimensionSize`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Returns the size of the given dimension of the operand. The operand must be array shaped. **`GetDimensionSize(operand, dimension)`** | Arguments | Type | Semantics | | ----------- | ------- | --------------------------------------------------- | | `operand` | `XlaOp` | n dimensional input array | | `dimension` | `int64` | A value in the interval `[0, n)` that specifies the | : : : dimension : For StableHLO information see [StableHLO - get_dimension_size](https://openxla.org/stablehlo/spec#get_dimension_size). ## GetTupleElement See also [`XlaBuilder::GetTupleElement`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Indexes into a tuple with a compile-time-constant value. The value must be a compile-time-constant so that shape inference can determine the type of the resulting value. This is analogous to `std::get(t)` in C++. Conceptually: ```cpp let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; let s: s32 = 5; let t: (f32[10], s32) = tuple(v, s); let element_1: s32 = gettupleelement(t, 1); // Inferred shape matches s32. ``` See also `tf.tuple`. **`GetTupleElement(tuple_data, index)`** Argument | Type | Semantics ------------ | ------- | -------------------- `tuple_data` | `XlaOP` | The tuple `index` | `int64` | Index of tuple shape For StableHLO information see [StableHLO - get_tuple_element](https://openxla.org/stablehlo/spec#get_tuple_element). ## Imag See also [`XlaBuilder::Imag`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise imaginary part of a complex (or real) shape. `x -> imag(x)`. If the operand is a floating point type, returns 0. **`Imag(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - imag](https://openxla.org/stablehlo/spec#imag). ## Infeed See also [`XlaBuilder::Infeed`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Infeed(shape, config)`** | Argument | Type | Semantics | | -------- | ----------------- | --------------------------------------------- | | `shape` | `Shape` | Shape of the data read from the Infeed | : : : interface. The layout field of the shape must : : : : be set to match the layout of the data sent : : : : to the device; otherwise its behavior is : : : : undefined. : | `config` | optional `string` | Configuration of the op. | Reads a single data item from the implicit Infeed streaming interface of the device, interpreting the data as the given shape and its layout, and returns a `XlaOp` of the data. Multiple Infeed operations are allowed in a computation, but there must be a total order among the Infeed operations. For example, two `Infeed`'s in the code below have a total order since there is a dependency between the while loops. ```cpp result1 = while (condition, init = init_value) { Infeed(shape) } result2 = while (condition, init = result1) { Infeed(shape) } ``` Nested tuple shapes are not supported. For an empty tuple shape, the Infeed operation is effectively a no-op and proceeds without reading any data from the Infeed of the device. > **Note:** We plan to allow multiple Infeed operations without a total order, > in which case the compiler will provide information about how the Infeed > operations are serialized in the compiled program. For StableHLO information see [StableHLO - infeed](https://openxla.org/stablehlo/spec#infeed). ## Iota See also [`XlaBuilder::Iota`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Iota(shape, iota_dimension)`** Builds a constant literal on device rather than a potentially large host transfer. Creates an array that has specified shape and holds values starting at zero and incrementing by one along the specified dimension. For floating-point types, the produced array is equivalent to `ConvertElementType(Iota(...))` where the `Iota` is of integral type and the conversion is to the floating-point type. Arguments | Type | Semantics ---------------- | ------- | -------------------------------------- `shape` | `Shape` | Shape of the array created by `Iota()` `iota_dimension` | `int64` | The dimension to increment along. For example, `Iota(s32[4, 8], 0)` returns ```cpp [[0, 0, 0, 0, 0, 0, 0, 0 ], [1, 1, 1, 1, 1, 1, 1, 1 ], [2, 2, 2, 2, 2, 2, 2, 2 ], [3, 3, 3, 3, 3, 3, 3, 3 ]] ``` `Iota(s32[4, 8], 1)` returns ```cpp [[0, 1, 2, 3, 4, 5, 6, 7 ], [0, 1, 2, 3, 4, 5, 6, 7 ], [0, 1, 2, 3, 4, 5, 6, 7 ], [0, 1, 2, 3, 4, 5, 6, 7 ]] ``` For StableHLO information see [StableHLO - iota](https://openxla.org/stablehlo/spec#iota). ## IsFinite See also [`XlaBuilder::IsFinite`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Tests whether each element of `operand` is finite, i.e., is not positive or negative infinity, and is not `NaN`. Returns an array of `PRED` values with the same shape as the input, where each element is `true` if and only if the corresponding input element is finite. **`IsFinite(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - is_finite](https://openxla.org/stablehlo/spec#is_finite). ## Log See also [`XlaBuilder::Log`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise natural logarithm `x -> ln(x)`. **`Log(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Log also supports the optional `result_accuracy` argument: **`Log(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - log](https://openxla.org/stablehlo/spec#log). ## Log1p See also [`XlaBuilder::Log1p`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise shifted natural logarithm `x -> ln(1+x)`. **`Log1p(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Log1p also supports the optional `result_accuracy` argument: **`Log1p(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - log_plus_one](https://openxla.org/stablehlo/spec#log_plus_one). ## Logistic See also [`XlaBuilder::Logistic`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise logistic function computation `x -> logistic(x)`. **`Logistic(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Logistic also supports the optional `result_accuracy` argument: **`Logistic(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - logistic](https://openxla.org/stablehlo/spec#logistic). ## Map See also [`XlaBuilder::Map`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Map(operands..., computation, dimensions)`** | Arguments | Type | Semantics | | ----------------- | ---------------------- | ------------------------------ | | `operands` | sequence of N `XlaOp`s | N arrays of types T_0..T_{N-1} | | `computation` | `XlaComputation` | Computation of type `T_0, T_1, | : : : .., T_{N + M -1} -> S` with N : : : : parameters of type T and M of : : : : arbitrary type. : | `dimensions` | `int64` array | Array of map dimensions | | `static_operands` | sequence of N `XlaOp`s | Static ops for the map | : : : operation : Applies a scalar function over the given `operands` arrays, producing an array of the same dimensions where each element is the result of the mapped function applied to the corresponding elements in the input arrays. The mapped function is an arbitrary computation with the restriction that it has N inputs of scalar type `T` and a single output with type `S`. The output has the same dimensions as the operands except that the element type T is replaced with S. For example: `Map(op1, op2, op3, computation, par1)` maps `elem_out <- computation(elem1, elem2, elem3, par1)` at each (multi-dimensional) index in the input arrays to produce the output array. For StableHLO information see [StableHLO - map](https://openxla.org/stablehlo/spec#map). ## Max See also [`XlaBuilder::Max`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise max operation on tensors `lhs` and `rhs`. **`Max(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Max: **`Max(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - maximum](https://openxla.org/stablehlo/spec#maximum). ## Min See also [`XlaBuilder::Min`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise min operation on `lhs` and `rhs`. **`Min(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Min: **`Min(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - minimum](https://openxla.org/stablehlo/spec#minimum). ## Mul See also [`XlaBuilder::Mul`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise product of `lhs` and `rhs`. **`Mul(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Mul: **`Mul(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - multiply](https://openxla.org/stablehlo/spec#multiply). ## Mulhi See also [`XlaBuilder::Mulhi`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise multiplication of two N-bit integer arrays `lhs` and `rhs`, returning an N-bit integer array containing the most significant N bits of the upcasted (N+N-bit) product: $$ \text{mulhi}(x, y) = \text{downcast}((\text{upcast}(x) * \text{upcast}(y)) >> N) $$ **`Mulhi(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | ------------------------------------------------- lhs | XlaOp | Left-hand-side operand: array of integer type T rhs | XlaOp | Right-hand-side operand: array of integer type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Mulhi: **`Mulhi(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | --------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : integer type T : | rhs | XlaOp | Right-hand-side operand: array of | : : : integer type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). > **Note:** `Mulhi` is only found in HLO and not found in StableHLO. CHLO > `Mulhi` in Frameworks will lower to HLO `Mulhi` see > [StableHLO - chlo.mulhi](https://openxla.org/stablehlo/generated/chlo?hl=en#chlomulhi_chlomulhiop) ## Neg See also [`XlaBuilder::Neg`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise negation `x -> -x`. **`Neg(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - negate](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#negate) ## Not See also [`XlaBuilder::Not`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise logical not `x -> !(x)`. **`Not(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - not](https://openxla.org/stablehlo/spec#not). ## OptimizationBarrier See also [`XlaBuilder::OptimizationBarrier`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Blocks any optimization pass from moving computations across the barrier. **`OptimizationBarrier(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Ensures that all inputs are evaluated before any operators that depend on the barrier's outputs. For StableHLO information see [StableHLO - optimization_barrier](https://openxla.org/stablehlo/spec#optimization_barrier). ## Or See also [`XlaBuilder::Or`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise OR of `lhs` and `rhs` . **`Or(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Or: **`Or(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - or](https://openxla.org/stablehlo/spec#or). ## Outfeed See also [`XlaBuilder::Outfeed`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Writes inputs to the outfeed. **`Outfeed(operand, shape_with_layout, outfeed_config)`** Arguments | Type | Semantics ------------------- | -------- | ---------------------------------------------- `operand` | `XlaOp` | array of type `T` `shape_with_layout` | `Shape` | Defines the layout of the data transferred `outfeed_config` | `string` | Constant of config for the Outfeed instruction `shape_with_layout` communicates the laid out shape that we want to outfeed. For StableHLO information see [StableHLO - outfeed](https://openxla.org/stablehlo/spec#outfeed). ## Pad See also [`XlaBuilder::Pad`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Pad(operand, padding_value, padding_config)`** | Arguments | Type | Semantics | | ---------------- | --------------- | --------------------------------------- | | `operand` | `XlaOp` | array of type `T` | | `padding_value` | `XlaOp` | scalar of type `T` to fill in the added | : : : padding : | `padding_config` | `PaddingConfig` | padding amount on both edges (low, | : : : high) and between the elements of each : : : : dimension : Expands the given `operand` array by padding around the array as well as between the elements of the array with the given `padding_value`. `padding_config` specifies the amount of edge padding and the interior padding for each dimension. `PaddingConfig` is a repeated field of `PaddingConfigDimension`, which contains three fields for each dimension: `edge_padding_low`, `edge_padding_high`, and `interior_padding`. `edge_padding_low` and `edge_padding_high` specify the amount of padding added at the low-end (next to index 0) and the high-end (next to the highest index) of each dimension respectively. The amount of edge padding can be negative -- the absolute value of negative padding indicates the number of elements to remove from the specified dimension. `interior_padding` specifies the amount of padding added between any two elements in each dimension; it may not be negative. Interior padding occurs logically before edge padding, so in the case of negative edge padding, elements are removed from the interior-padded operand. This operation is a no-op if the edge padding pairs are all (0, 0) and the interior padding values are all 0. The figure below shows examples of different `edge_padding` and `interior_padding` values for a two-dimensional array. For StableHLO information see [StableHLO - pad](https://openxla.org/stablehlo/spec#pad). ## Parameter See also [`XlaBuilder::Parameter`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). `Parameter` represents an argument input to a computation. > **Note:** `Parameter` is only found in HLO. It is not found in StableHLO. ## PartitionID See also [`XlaBuilder::BuildPartitionId`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Produces `partition_id` of the current process. **`PartitionID(shape)`** Arguments | Type | Semantics --------- | ------- | ----------------- `shape` | `Shape` | Shape of the data `PartitionID` may appear in HLO dumps but it is not intended to be constructed manually by end users. For StableHLO information see [StableHLO - partition_id](https://openxla.org/stablehlo/spec#partition_id). ## PopulationCount See also [`XlaBuilder::PopulationCount`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Computes the number of bits set in each element of `operand`. **`PopulationCount(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - popcnt](https://openxla.org/stablehlo/spec#popcnt). ## Pow See also [`XlaBuilder::Pow`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise exponentiation of `lhs` by `rhs`. **`Pow(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Pow: **`Pow(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - power](https://openxla.org/stablehlo/spec#power). ## Real See also [`XlaBuilder::Real`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise real part of a complex (or real) shape. `x -> real(x)`. If the operand is a floating point type, `Real` returns the same value. **`Real(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - real](https://openxla.org/stablehlo/spec#real). ## Recv See also [`XlaBuilder::Recv`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). `Recv`, `RecvWithTokens`, and `RecvToHost` are operations that serve as communication primitives in HLO. These ops typically appear in HLO dumps as part of low-level input/output or cross-device transfer, but they are not intended to be constructed manually by end users. **`Recv(shape, handle)`** Arguments | Type | Semantics --------- | --------------- | ----------------------------------------- `shape` | `Shape` | shape of the data to receive `handle` | `ChannelHandle` | unique identifier for each send/recv pair Receives data of the given shape from a `Send` instruction in another computation that shares the same channel handle. Returns a XlaOp for the received data. For StableHLO information see [StableHLO - recv](https://openxla.org/stablehlo/spec#recv). ### RecvDone See also [`HloInstruction::CreateRecv` and `HloInstruction::CreateRecvDone`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). Similar to [`Send`](#send), the client API of `Recv` operation represents synchronous communication. However, the instruction is internally decomposed into 2 HLO instructions (`Recv` and `RecvDone`) to enable asynchronous data transfers. **`Recv(const Shape& shape, int64 channel_id)`** Allocates resources required to receive data from a [`Send`](#send) instruction with the same channel_id. Returns a context for the allocated resources, which is used by a following `RecvDone` instruction to wait for the completion of the data transfer. The context is a tuple of {receive buffer (shape), request identifier (U32)} and it can only be used by a `RecvDone` instruction. Given a context created by a `Recv` instruction, waits for the data transfer to complete and return the received data. > **Note:** `RecvDone` is only found in HLO. It is not found in StableHLO. ## Reduce See also [`XlaBuilder::Reduce`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Applies a reduction function to one or more arrays in parallel. **`Reduce(operands..., init_values..., computation, dimensions_to_reduce)`** | Arguments | Type | Semantics | | :--------------------- | :-------------------- | :-------------------------- | | `operands` | Sequence of N `XlaOp` | N arrays of types `T_0,..., | : : : T_{N-1}`. : | `init_values` | Sequence of N `XlaOp` | N scalars of types | : : : `T_0,..., T_{N-1}`. : | `computation` | `XlaComputation` | computation of type | : : : `T_0,..., T_{N-1}, T_0, : : : : ...,T_{N-1} ->` : : : : `Collate(T_0,..., : : : : T_{N-1})`. : | `dimensions_to_reduce` | `int64` array | unordered array of | : : : dimensions to reduce. : Where: * N is required to be greater or equal to 1. * The computation has to be "roughly" associative (see below). * All input arrays must have the same dimensions. * All initial values have to form an identity under `computation`. * If `N = 1`, `Collate(T)` is `T`. * If `N > 1`, `Collate(T_0, ..., T_{N-1})` is a tuple of `N` elements of type `T`. This operation reduces one or more dimensions of each input array into scalars. The number of dimensions of each returned array is `number_of_dimensions(operand) - len(dimensions)`. The output of the op is `Collate(Q_0, ..., Q_N)` where `Q_i` is an array of type `T_i`, the dimensions of which are described below. Different backends are allowed to reassociate the reduction computation. This can lead to numerical differences, as some reduction functions like addition are not associative for floats. However, if the range of the data is limited, floating-point addition is close enough to be associative for most practical uses. For StableHLO information see [StableHLO - reduce](https://openxla.org/stablehlo/spec#reduce). ### Examples When reducing across one dimension in a single 1D array with values `[10, 11, 12, 13]`, with reduction function `f` (this is `computation`) then that could be computed as `f(10, f(11, f(12, f(init_value, 13)))` but there are also many other possibilities, e.g. `f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))` The following is a rough pseudo-code example of how reduction could be implemented, using summation as the reduction computation with an initial value of 0. ```python result_shape <- remove all dims in dimensions from operand_shape # Iterate over all elements in result_shape. The number of r's here is equal # to the number of dimensions of the result. for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...: # Initialize this result element result[r0, r1...] <- 0 # Iterate over all the reduction dimensions for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...: # Increment the result element with the value of the operand's element. # The index of the operand's element is constructed from all ri's and di's # in the right order (by construction ri's and di's together index over the # whole operand shape). result[r0, r1...] += operand[ri... di] ``` Here's an example of reducing a 2D array (matrix). The shape has 2 dimensions, dimension 0 of size 2 and dimension 1 of size 3: Results of reducing dimensions 0 or 1 with an "add" function: Note that both reduction results are 1D arrays. The diagram shows one as column and another as row just for visual convenience. For a more complex example, here is a 3D array. Its number of dimensions is 3, dimension 0 of size 4, dimension 1 of size 2 and dimension 2 of size 3. For simplicity, the values 1 to 6 are replicated across dimension 0. Similarly to the 2D example, we can reduce just one dimension. If we reduce dimension 0, for example, we get a 2-dimensional array where all values across dimension 0 were folded into a scalar: ```text | 4 8 12 | | 16 20 24 | ``` If we reduce dimension 2, we also get a 2-dimensional array where all values across dimension 2 were folded into a scalar: ```text | 6 15 | | 6 15 | | 6 15 | | 6 15 | ``` Note that the relative order between the remaining dimensions in the input is preserved in the output, but some dimensions may get assigned new numbers (since the number of dimensions changes). We can also reduce multiple dimensions. Add-reducing dimensions 0 and 1 produces the 1D array `[20, 28, 36]`. Reducing the 3D array over all its dimensions produces the scalar `84`. ### Variadic Reduce When `N > 1`, reduce function application is slightly more complex, as it is applied simultaneously to all inputs. The operands are supplied to the computation in the following order: * Running reduced value for the first operand * ... * Running reduced value for the N'th operand * Input value for the first operand * ... * Input value for the N'th operand For example, consider the following reduction function, which can be used to compute the max and the argmax of a 1-D array in parallel: ```python f: (Float, Int, Float, Int) -> Float, Int f(max, argmax, value, index): if value >= max: return (value, index) else: return (max, argmax) ``` For 1-D Input arrays `V = Float[N], K = Int[N]`, and init values `I_V = Float, I_K = Int`, the result `f_(N-1)` of reducing across the only input dimension is equivalent to the following recursive application: ```python f_0 = f(I_V, I_K, V_0, K_0) f_1 = f(f_0.first, f_0.second, V_1, K_1) ... f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1)) ``` Applying this reduction to an array of values, and an array of sequential indices (i.e. iota), will co-iterate over the arrays, and return a tuple containing the maximal value and the matching index. ## ReducePrecision See also [`XlaBuilder::ReducePrecision`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Models the effect of converting floating-point values to a lower-precision format (such as IEEE-FP16) and back to the original format. The number of exponent and mantissa bits in the lower-precision format can be specified arbitrarily, although all bit sizes may not be supported on all hardware implementations. **`ReducePrecision(operand, exponent_bits, mantissa_bits)`** Arguments | Type | Semantics --------------- | ------- | ------------------------------------------------- `operand` | `XlaOp` | array of floating-point type `T`. `exponent_bits` | `int32` | number of exponent bits in lower-precision format `mantissa_bits` | `int32` | number of mantissa bits in lower-precision format The result is an array of type `T`. The input values are rounded to the nearest value representable with the given number of mantissa bits (using "ties to even" semantics), and any values that exceed the range specified by the number of exponent bits are clamped to positive or negative infinity. `NaN` values are retained, although they may be converted to canonical `NaN` values. The lower-precision format must have at least one exponent bit (in order to distinguish a zero value from an infinity, since both have a zero mantissa), and must have a non-negative number of mantissa bits. The number of exponent or mantissa bits may exceed the corresponding value for type `T`; the corresponding portion of the conversion is then simply a no-op. For StableHLO information see [StableHLO - reduce_precision](https://openxla.org/stablehlo/spec#reduce_precision). ## ReduceScatter See also [`XlaBuilder::ReduceScatter`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). ReduceScatter is a collective operation that effectively does an AllReduce and then scatters the result by splitting it into `shard_count` blocks along the `scatter_dimension` and replica `i` in the replica group receives the `ith` shard. **`ReduceScatter(operand, computation, scatter_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)`** | Arguments | Type | Semantics | | ----------------------- | --------------------- | -------------------------- | | `operand` | `XlaOp` | Array or a non-empty tuple | : : : of arrays to reduce across : : : : replicas. : | `computation` | `XlaComputation` | Reduction computation | | `scatter_dimension` | `int64` | Dimension to scatter. | | `shard_count` | `int64` | Number of blocks to split | : : : `scatter_dimension` : | `replica_groups` | `ReplicaGroup` vector | Groups between which the | : : : reductions are performed : | `channel_id` | optional | Optional channel ID for | : : `ChannelHandle` : cross-module communication : | `layout` | optional `Layout` | user-specified memory | : : : layout : | `use_global_device_ids` | optional `bool` | user-specified flag | - When `operand` is a tuple of arrays, the reduce-scatter is performed on each element of the tuple. - `replica_groups` is a list of replica groups between which the reduction is performed (replica id for the current replica can be retrieved using [`ReplicaId`](#replicaid)). The order of replicas in each group determines the order in which the all-reduce result will be scattered. `replica_groups` must either be empty (in which case all replicas belong to a single group), or contain the same number of elements as the number of replicas. When there are more than one replica groups, they all must be of the same size. For example, `replica_groups = {0, 2}, {1, 3}` performs reduction between the replicas `0` and `2`, and `1` and `3` and then scatters the result. - `shard_count` is the size of each replica group. We need this in cases where `replica_groups` are empty. If `replica_groups` is not empty, `shard_count` must be equal to the size of each replica group. - `channel_id` is used for cross-module communication: only `reduce-scatter` operations with the same `channel_id` can communicate with each other. - `layout` See [xla::shapes for more information on layouts.](https://openxla.org/xla/shapes) - `use_global_device_ids` is a user-specified flag. When `false`(default) the numbers in `replica_groups` are [`ReplicaId`](#replicaid) when `true` the `replica_groups` represent a global id of (`ReplicaID`*`partition_count` + `partition_id`). For example: - With 2 replicas and 4 partitions, - replica_groups={{0,1,4,5},{2,3,6,7}} and use_global_device_ids=true - group[0] = (0,0), (0,1), (1,0), (1,1) - group[1] = (0,2), (0,3), (1,2), (1,3) - where each pair is (replica_id, partition_id). The output shape is the input shape with the `scatter_dimension` made `shard_count` times smaller. For example, if there are two replicas and the operand has the value `[1.0, 2.25]` and `[3.0, 5.25]` respectively on the two replicas, then the output value from this op where `scatter_dim` is `0` will be `[4.0]` for the first replica and `[7.5]` for the second replica. For StableHLO information see [StableHLO - reduce_scatter](https://openxla.org/stablehlo/spec#reduce_scatter). ### ReduceScatter - Example 1 - StableHLO In the above example, there are 2 replicas participating in the ReduceScatter. On each replica, the operand has shape f32[2,4]. An all-reduce (sum) is performed across the replicas, producing a reduced value of shape f32[2,4] on each replica. This reduced value is then split into 2 parts along dimension 1, so each part has shape f32[2,2]. Each replica within the process group receives the part corresponding to its position in the group. As a result, the output on each replica has shape f32[2,2]. ## ReduceWindow See also [`XlaBuilder::ReduceWindow`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Applies a reduction function to all elements in each window of a sequence of N multi-dimensional arrays, producing a single or a tuple of N multi-dimensional arrays as output. Each output array has the same number of elements as the number of valid positions of the window. A pooling layer can be expressed as a `ReduceWindow`. Similar to [`Reduce`](#reduce), the applied `computation` is always passed the `init_values` on the left-hand side. **`ReduceWindow(operands..., init_values..., computation, window_dimensions, window_strides, padding)`** | Arguments | Type | Semantics | | ------------------- | ------------------- | -------------------------------- | | `operands` | `N XlaOps` | A sequence of N | : : : multi-dimensional arrays of : : : : types `T_0,..., T_{N-1}`, each : : : : representing the base area on : : : : which the window is placed. : | `init_values` | `N XlaOps` | A sequence of N scalars of types | : : : `T_0,..., T_{N-1}`, representing : : : : the starting values for the : : : : reduction, one for each of the N : : : : operands. : | `computation` | `XlaComputation` | Reduction function of type `T_0, | : : : ..., T_{N-1}, T_0, ..., T_{N-1} : : : : -> Collate(T_0, ..., T_{N-1})`, : : : : to apply to elements in each : : : : window of all the input : : : : operands. : | `window_dimensions` | `ArraySlice` | array of integers for window | : : : dimension values : | `window_strides` | `ArraySlice` | array of integers for window | : : : stride values : | `base_dilations` | `ArraySlice` | array of integers for base | : : : dilation values : | `window_dilations` | `ArraySlice` | array of integers for window | : : : dilation values : | `padding` | `Padding` | padding type for window | : : : (Padding\:\:kSame, which pads so : : : : as to have the same output shape : : : : as input if the stride is 1, or : : : : Padding\:\:kValid, which uses no : : : : padding and "stops" the window : : : : once it no longer fits) : Where: * N is required to be greater or equal to 1. * All input arrays must have the same dimensions. * All initial values have to form an identity under `computation`. * If `N = 1`, `Collate(T)` is `T`. * If `N > 1`, `Collate(T_0, ..., T_{N-1})` is a tuple of `N` elements of type `(T0,...T{N-1})`. For StableHLO information see [StableHLO - reduce_window](https://openxla.org/stablehlo/spec#reduce_window). ### ReduceWindow - Example 1 Input is a matrix of size [4x6] and both window_dimensions and window_stride_dimensions are [2x3]. ```cpp // Create a computation for the reduction (maximum). XlaComputation max; { XlaBuilder builder(client_, "max"); auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y"); auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x"); builder.Max(y, x); max = builder.Build().value(); } // Create a ReduceWindow computation with the max reduction computation. XlaBuilder builder(client_, "reduce_window_2x3"); auto shape = ShapeUtil::MakeShape(F32, {4, 6}); auto input = builder.Parameter(0, shape, "input"); builder.ReduceWindow( input, /*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)), *max, /*window_dimensions=*/{2, 3}, /*window_stride_dimensions=*/{2, 3}, Padding::kValid); ``` Stride of 1 in a dimension specifies that the position of a window in the dimension is 1 element away from its adjacent window. In order to specify that no windows overlap with each other, window_stride_dimensions should be equal to window_dimensions. The figure below illustrates the use of two different stride values. Padding is applied to each dimension of the input and the calculations are the same as though the input came in with the dimensions it has after padding. For a non-trivial padding example, consider computing reduce-window minimum (initial value is `MAX_FLOAT`) with dimension `3` and stride `2` over the input array `[10000, 1000, 100, 10, 1]`. Padding `kValid` computes minimums over two _valid_ windows: `[10000, 1000, 100]` and `[100, 10, 1]`, resulting in the output `[100, 1]`. Padding `kSame` first pads the array so that the shape after the reduce-window would be the _same_ as input for stride one by adding initial elements on both sides, getting `[MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE]`. Running reduce-window over the padded array operates on three windows `[MAX_VALUE, 10000, 1000]`, `[1000, 100, 10]`, `[10, 1, MAX_VALUE]`, and yields `[1000, 10, 1]`. The evaluation order of the reduction function is arbitrary and may be non-deterministic. Therefore, the reduction function should not be overly sensitive to reassociation. See the discussion about associativity in the context of [`Reduce`](#reduce) for more details. ### ReduceWindow - Example 2 - StableHLO In the above example: Input) The operand has an input shape of S32[3,2]. With a values of `[[1,2],[3,4],[5,6]]` Step 1) Base dilation with factor 2 along the row dimension inserts holes between each row of the operand. Padding of 2 rows at the top and 1 row at the bottom is applied after dilation. As a result, the tensor becomes taller. Step 2) A window of shape [2,1] is defined, with window dilation [3,1]. This means each window selects two elements from the same column, but the second element is taken three rows below the first rather than directly beneath it. Step 3) The windows are then slid across the operand with stride [4,1]. This causes the window to move down four rows at a time, while shifting one column at a time horizontally. Padding cells are filled with the `init_value` (in this case `init_value = 0`). Values 'falling into' dilation cells are ignored. Because of the stride and padding, some windows overlap only zeros and holes, while others overlap real input values. Step 4) Within each window, the elements are combined using the reduction function (a, b) → a + b, starting from an initial value of 0. The top two windows see only padding and holes, so their results are 0. The bottom windows capture the values 3 and 4 from the input and return those as results. Results) The final output has shape S32[2,2], with values: `[[0,0],[3,4]]` ## Rem See also [`XlaBuilder::Rem`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise remainder of dividend `lhs` and divisor `rhs`. The sign of the result is taken from the dividend, and the absolute value of the result is always less than the divisor's absolute value. **`Rem(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Rem: **`Rem(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - remainder](https://openxla.org/stablehlo/spec#remainder). ## ReplicaId See also [`XlaBuilder::ReplicaId`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Returns the unique ID (U32 scalar) of the replica. **`ReplicaId()`** The unique ID of each replica is an unsigned integer in the interval `[0, N)`, where `N` is the number of replicas. Since all the replicas are running the same program, a `ReplicaId()` call in the program will return a different value on each replica. For StableHLO information see [StableHLO - replica_id](https://openxla.org/stablehlo/spec#replica_id). ## Reshape See also [`XlaBuilder::Reshape`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). and the [`Collapse`](#collapse) operation. Reshapes the dimensions of an array into a new configuration. **`Reshape(operand, dimensions)`** Arguments | Type | Semantics ------------ | -------------- | --------------------------------- `operand` | `XlaOp` | array of type T `dimensions` | `int64` vector | vector of sizes of new dimensions Conceptually, reshape first flattens an array into a one-dimensional vector of data values, and then refines this vector into a new shape. The input arguments are an arbitrary array of type T, a compile-time-constant vector of dimension indices, and a compile-time-constant vector of dimension sizes for the result. The `dimensions` vector determines the size of the output array. The value at index 0 in `dimensions` is the size of dimension 0, the value at index 1 is the size of dimension 1, and so on. The product of the `dimensions` dimensions must equal the product of the operand's dimension sizes. When refining the collapsed array into the multidimensional array defined by `dimensions`, the dimensions in `dimensions` are ordered from slowest varying (most major) and to fastest varying (most minor). For example, let v be an array of 24 elements: ```cpp let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}}, {{20, 21, 22}, {25, 26, 27}}, {{30, 31, 32}, {35, 36, 37}}, {{40, 41, 42}, {45, 46, 47}}}; let v012_24 = Reshape(v, {24}); then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27, 30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47}; let v012_83 = Reshape(v, {8,3}); then v012_83 == f32[8x3] {{10, 11, 12}, {15, 16, 17}, {20, 21, 22}, {25, 26, 27}, {30, 31, 32}, {35, 36, 37}, {40, 41, 42}, {45, 46, 47}}; ``` As a special case, reshape can transform a single-element array to a scalar and vice versa. For example, ```cpp Reshape(f32[1x1] {{5}}, {}) == 5; Reshape(5, {1,1}) == f32[1x1] {{5}}; ``` For StableHLO information see [StableHLO - reshape](https://openxla.org/stablehlo/spec#reshape). ### Reshape (explicit) See also [`XlaBuilder::Reshape`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Reshape(shape, operand)`** Reshape op that uses an explicit target shape. Arguments | Type | Semantics --------- | ------- | ---------------------- `shape` | `Shape` | Output shape of type T `operand` | `XlaOp` | array of type T ## Rev (reverse) See also [`XlaBuilder::Rev`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Rev(operand, dimensions)`** Arguments | Type | Semantics ------------ | ------------------- | --------------------- `operand` | `XlaOp` | array of type T `dimensions` | `ArraySlice` | dimensions to reverse Reverses the order of elements in the `operand` array along the specified `dimensions`, generating an output array of the same shape. Each element of the operand array at a multidimensional index is stored into the output array at a transformed index. The multidimensional index is transformed by reversing the index in each dimension to be reversed (i.e., if a dimension of size N is one of the reversing dimensions, its index i is transformed into N - 1 - i). One use for the `Rev` operation is to reverse the convolution weight array along the two window dimensions during the gradient computation in neural networks. For StableHLO information see [StableHLO - reverse](https://openxla.org/stablehlo/spec#reverse). ## RngNormal See also [`XlaBuilder::RngNormal`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Constructs an output of a given shape with random numbers generated following the $N(\mu, \sigma)$ normal distribution. The parameters $\mu$ and $\sigma$, and output shape have to have a floating point elemental type. The parameters furthermore have to be scalar valued. **`RngNormal(mu, sigma, shape)`** | Arguments | Type | Semantics | | --------- | ------- | ----------------------------------------------------- | | `mu` | `XlaOp` | Scalar of type T specifying mean of generated numbers | | `sigma` | `XlaOp` | Scalar of type T specifying standard deviation of | : : : generated : | `shape` | `Shape` | Output shape of type T | For StableHLO information see [StableHLO - rng](https://openxla.org/stablehlo/spec#rng). ## RngUniform See also [`XlaBuilder::RngUniform`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Constructs an output of a given shape with random numbers generated following the uniform distribution over the interval $[a,b)$. The parameters and output element type have to be a boolean type, an integral type or a floating point types, and the types have to be consistent. The CPU and GPU backends currently only support F64, F32, F16, BF16, S64, U64, S32 and U32. Furthermore, the parameters need to be scalar valued. If $b <= a$ the result is implementation-defined. **`RngUniform(a, b, shape)`** Arguments | Type | Semantics --------- | ------- | --------------------------------------------------- `a` | `XlaOp` | Scalar of type T specifying lower limit of interval `b` | `XlaOp` | Scalar of type T specifying upper limit of interval `shape` | `Shape` | Output shape of type T For StableHLO information see [StableHLO - rng](https://openxla.org/stablehlo/spec#rng). ## RngBitGenerator See also [`XlaBuilder::RngBitGenerator`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Generates an output with a given shape filled with uniform random bits using the specified algorithm (or backend default) and returns an updated state (with the same shape as initial state) and the generated random data. Initial state is the initial state of the current random number generation. It and the required shape and valid values are dependent on the algorithm used. The output is guaranteed to be a deterministic function of the initial state but it is *not* guaranteed to be deterministic between backends and different compiler versions. **`RngBitGenerator(algorithm, initial_state, shape)`** Arguments | Type | Semantics --------------- | ----------------- | ------------------------------------- `algorithm` | `RandomAlgorithm` | PRNG algorithm to be used. `initial_state` | `XlaOp` | Initial state for the PRNG algorithm. `shape` | `Shape` | Output shape for generated data. Available values for `algorithm`: - `rng_default`: Backend specific algorithm with backend specific shape requirements. - `rng_three_fry`: ThreeFry counter-based PRNG algorithm. The `initial_state` shape is `u64[2]` with arbitrary values. [Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3.](http://www.thesalmons.org/john/random123/papers/random123sc11.pdf) - `rng_philox`: Philox algorithm to generate random numbers in parallel. The `initial_state` shape is `u64[3]` with arbitrary values. [Salmon et al. SC 2011. Parallel random numbers: as easy as 1, 2, 3.](http://www.thesalmons.org/john/random123/papers/random123sc11.pdf) For StableHLO information see [StableHLO - rng_bit_generator](https://openxla.org/stablehlo/spec#rng_bit_generator). ## RngGetAndUpdateState See also [`HloInstruction::CreateRngGetAndUpdateState`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). The API of the various `Rng` operations are internally decomposed into HLO instructions including `RngGetAndUpdateState`. `RngGetAndUpdateState` serves as a primitive in HLO. This op may appear in HLO dumps, but it is not intended to be constructed manually by end users. > **Note:** `RngGetAndUpdateState` is only found in HLO. It is not found in > StableHLO. ## Round See also [`XlaBuilder::Round`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise rounding, ties away from zero. **`Round(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function ### RoundNearestAfz See also [`XlaBuilder::RoundNearestAfz`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise rounding towards the nearest integer, breaking ties away from zero. **`RoundNearestAfz(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - round_nearest_afz](https://openxla.org/stablehlo/spec#round_nearest_afz). ### RoundNearestEven See also [`XlaBuilder::RoundNearestEven`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise rounding, ties to the nearest even. **`RoundNearestEven(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - round_nearest_even](https://openxla.org/stablehlo/spec#round_nearest_even). ## Rsqrt See also [`XlaBuilder::Rsqrt`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise reciprocal of square root operation `x -> 1.0 / sqrt(x)`. **`Rsqrt(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Rsqrt also supports the optional `result_accuracy` argument: **`Rsqrt(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - rsqrt](https://openxla.org/stablehlo/spec#rsqrt). ## Scan See also [`XlaBuilder::Scan`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Applies a reduction function to an array across a given dimension, producing a final state and an array of intermediate values. **`Scan(inputs..., inits..., to_apply, scan_dimension, is_reverse, is_associative)`** | Arguments | Type | Semantics | | :--- | :--- | :--- | | `inputs` | Sequence of `m` `XlaOp`s | The arrays to scan. | | `inits` | Sequence of `k` `XlaOp`s | Initial carries. | | `to_apply` | `XlaComputation` | Computation of type `i_0, ..., i_{m-1}, c_0, ..., c_{k-1} -> (o_0, ..., o_{n-1}, c'_0, ..., c'_{k-1})`. | | `scan_dimension` | `int64` | The dimension to scan across. | | `is_reverse` | `bool` | If true, scan in reverse order. | | `is_associative` | `bool` (tri-state) | If true, the operation is associative. | The function `to_apply` is applied sequentially to elements in `inputs` along `scan_dimension`. If `is_reverse` is false, elements are processed in order `0` to `N-1`, where `N` is the size of `scan_dimension`. If `is_reverse` is true, elements are processed from `N-1` down to `0`. The `to_apply` function takes `m + k` operands: 1. `m` current elements from `inputs`. 2. `k` carry values from the previous step (or `inits` for the first element). The `to_apply` function returns a tuple of `n + k` values: 1. `n` elements of `outputs`. 2. `k` new carry values. The Scan operation produces a tuple of `n + k` values: 1. The `n` output arrays, containing the output values for each step. 2. The final `k` carry values after processing all elements. The types of the `m` inputs must match the types of the first `m` parameters of `to_apply` with an extra scan dimension. The types of the `n` outputs must match the types of the first `n` return values of `to_apply` with an extra scan dimension. The extra scan dimension among all inputs and outputs must have the same size `N`. The types of the last `k` parameters and return values of `to_apply` as well as the `k` inits must match. For example (`m, n, k == 1, N == 3`), for initial carry `i`, input `[a, b, c]`, function `f(x, c) -> (y, c')`, and `scan_dimension=0`, `is_reverse=false`: - Step 0: `f(a, i) -> (y0, c0)` - Step 1: `f(b, c0) -> (y1, c1)` - Step 2: `f(c, c1) -> (y2, c2)` The output of `Scan` is `([y0, y1, y2], c2)`. > **Note:** `Scan` is only found in HLO. It is not found in StableHLO. ## Scatter See also [`XlaBuilder::Scatter`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). The XLA scatter operation generates a sequence of results which are the values of the input array `operands`, with several slices (at indices specified by `scatter_indices`) updated with the sequence of values in `updates` using `update_computation`. **`Scatter(operands..., scatter_indices, updates..., update_computation, dimension_numbers, indices_are_sorted, unique_indices)`** Arguments | Type | Semantics ------------------------------ | ------------------------- | --------- `operands` | Sequence of N `XlaOp` | N arrays of types `T_0, ..., T_N` to be scattered into. `scatter_indices` | `XlaOp` | Array containing the starting indices of the slices that must be scattered to. `updates` | Sequence of N `XlaOp` | N arrays of types `T_0, ..., T_N`. `updates[i]` contains the values that must be used for scattering `operands[i]`. `update_computation` | `XlaComputation` | Computation to be used for combining the existing values in the input array and the updates during scatter. This computation should be of type `T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N)`. `index_vector_dim` | `int64` | The dimension in `scatter_indices` that contains the starting indices. `update_window_dims` | `ArraySlice` | The set of dimensions in `updates` shape that are *window dimensions*. `inserted_window_dims` | `ArraySlice` | The set of *window dimensions* that must be inserted into `updates` shape. `scatter_dims_to_operand_dims` | `ArraySlice` | A dimensions map from the scatter indices to the operand index space. This array is interpreted as mapping `i` to `scatter_dims_to_operand_dims[i]` . It has to be one-to-one and total. `dimension_number` | `ScatterDimensionNumbers` | Dimension numbers for scatter operation `indices_are_sorted` | `bool` | Whether the indices are guaranteed to be sorted by the caller. `unique_indices` | `bool` | Whether the indices are guaranteed to be unique by the caller. Where: * N is required to be greater or equal to 1. * `operands`[`0`], ..., `operands`[`N-1`] must all have the same dimensions. * `updates`[`0`], ..., `updates`[`N-1`] must all have the same dimensions. * If `N = 1`, `Collate(T)` is `T`. * If `N > 1`, `Collate(T_0, ..., T_N)` is a tuple of `N` elements of type `T`. If `index_vector_dim` is equal to `scatter_indices.rank` we implicitly consider `scatter_indices` to have a trailing `1` dimension. We define `update_scatter_dims` of type `ArraySlice` as the set of dimensions in `updates` shape that are not in `update_window_dims`, in ascending order. The arguments of scatter should follow these constraints: - Each `updates` array must have `update_window_dims.size + scatter_indices.rank - 1` dimensions. - Bounds of dimension `i` in each `updates` array must conform to the following: - If `i` is present in `update_window_dims` (i.e. equal to `update_window_dims`[`k`] for some `k`), then the bound of dimension `i` in `updates` must not exceed the corresponding bound of `operand` after accounting for the `inserted_window_dims` (i.e. `adjusted_window_bounds`[`k`], where `adjusted_window_bounds` contains the bounds of `operand` with the bounds at indices `inserted_window_dims` removed). - If `i` is present in `update_scatter_dims` (i.e. equal to `update_scatter_dims`[`k`] for some `k`), then the bound of dimension `i` in `updates` must be equal to the corresponding bound of `scatter_indices`, skipping `index_vector_dim` (i.e. `scatter_indices.shape.dims`[`k`], if `k` < `index_vector_dim` and `scatter_indices.shape.dims`[`k+1`] otherwise). - `update_window_dims` must be in ascending order, not have any repeating dimension numbers, and be in the range `[0, updates.rank)`. - `inserted_window_dims` must be in ascending order, not have any repeating dimension numbers, and be in the range `[0, operand.rank)`. - `operand.rank` must equal the sum of `update_window_dims.size` and `inserted_window_dims.size`. - `scatter_dims_to_operand_dims.size` must be equal to `scatter_indices.shape.dims`[`index_vector_dim`], and its values must be in the range `[0, operand.rank)`. For a given index `U` in each `updates` array, the corresponding index `I` in the corresponding `operands` array into which this update has to be applied is computed as follows: 1. Let `G` = { `U`[`k`] for `k` in `update_scatter_dims` }. Use `G` to look up an index vector `S` in the `scatter_indices` array such that `S`[`i`] = `scatter_indices`[Combine(`G`, `i`)] where Combine(A, b) inserts b at positions `index_vector_dim` into A. 2. Create an index `S``in` into `operand` using `S` by scattering `S` using the `scatter_dims_to_operand_dims` map. More formally: 1. `S``in`[`scatter_dims_to_operand_dims`[`k`]] = `S`[`k`] if `k` < `scatter_dims_to_operand_dims.size`. 2. `S``in`[`_`] = `0` otherwise. 3. Create an index `W``in` into each `operands` array by scattering the indices at `update_window_dims` in `U` according to `inserted_window_dims`. More formally: 1. `W``in`[`window_dims_to_operand_dims`(`k`)] = `U`[`k`] if `k` is in `update_window_dims`, where `window_dims_to_operand_dims` is the monotonic function with domain [`0`, `update_window_dims.size`) and range [`0`, `operand.rank`) \\ `inserted_window_dims`. (For example, if `update_window_dims.size` is `4`, `operand.rank` is `6`, and `inserted_window_dims` is {`0`, `2`} then `window_dims_to_operand_dims` is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}). 2. `W``in`[`_`] = `0` otherwise. 4. `I` is `W``in` + `S``in` where + is element-wise addition. In summary, the scatter operation can be defined as follows. - Initialize `output` with `operands`, i.e. for all indices `J`, for all indices `O` in the `operands`[`J`] array: \ `output`[`J`][`O`] = `operands`[`J`][`O`] - For every index `U` in the `updates`[`J`] array and the corresponding index `O` in the `operand`[`J`] array, if `O` is a valid index for `output`: \ `(output`[`0`][`O`], ..., `output`[`N-1`][`O`]) =`update_computation`(`output`[`0`][`O`], ..., ,`output`[`N-1`][`O`],`updates`[`0`][`U`], ...,`updates`[`N-1`][`U`]) The order in which updates are applied is non-deterministic. So, when multiple indices in `updates` refer to the same index in `operands`, the corresponding value in `output` will be non-deterministic. Note that the first parameter that is passed into the `update_computation` will always be the current value from the `output` array and the second parameter will always be the value from the `updates` array. This is important specifically for cases when the `update_computation` is _not commutative_. If `indices_are_sorted` is set to true then XLA can assume that `scatter_indices` are sorted (in ascending order, _after_ scattering its values according to `scatter_dims_to_operand_dims`) by the user. If they are not then the semantics are implementation defined. If `unique_indices` is set to true then XLA can assume that all elements scattered to are unique. So XLA could use non-atomic operations. If `unique_indices` is set to true and the indices being scattered to are not unique then the semantics is implementation defined. Informally, the scatter op can be viewed as an _inverse_ of the gather op, i.e. the scatter op updates the elements in the input that are extracted by the corresponding gather op. For a detailed informal description and examples, refer to the "Informal Description" section under `Gather`. For StableHLO information see [StableHLO - scatter](https://openxla.org/stablehlo/spec#scatter). ### Scatter - Example 1 - StableHLO In the above image, each row of the table is an example of one update index example. Let's review stepwise from left(Update Index) to right(Result Index): Input) `input` has shape S32[2,3,4,2]. `scatter_indices` have shape S64[2,2,3,2]. `updates` have shape S32[2,2,3,1,2]. Update Index) As part of the input we are given `update_window_dims:[3,4]`. This tell us that `updates`'s dim 3 and dim 4 are window dimensions, highlighted in yellow. This allows us to derive that `update_scatter_dims` = [0,1,2]. Update Scatter Index) Shows us the extracted `updated_scatter_dims` for each. (The non-yellow of column Update Index) Start Index) Looking at the `scatter_indices` tensor image we can see that our values from the previous step (Update scatter Index), give us the location of the start index. From `index_vector_dim` we are also told the dimension of the `starting_indices` that contains the starting indices, which for `scatter_indices` is dim 3 with a size 2. Full Start Index) `scatter_dims_to_operand_dims` = [2,1] tells us the first element of the index vector goes to operand dim 2. The second element of the index vector goes to operand dim 1. The remaining operand dimensions are filled with 0. Full Batching Index) We can see the purple highlighted area is shown in this column(full batching index), the update scatter index column, and update index column. Full Window Index) Computed from the `update_window_dimensions` [3,4]. Result Index) The addition of Full Start Index, Full Batching Index, and Full Window Index in the `operand` tensor. Notice the green highlighted regions correspond to the `operand` figure as well. The last row is skipped because it falls outside of `operand` tensor. ## Select See also [`XlaBuilder::Select`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Constructs an output array from elements of two input arrays, based on the values of a predicate array. **`Select(pred, on_true, on_false)`** Arguments | Type | Semantics ---------- | ------- | ------------------ `pred` | `XlaOp` | array of type PRED `on_true` | `XlaOp` | array of type T `on_false` | `XlaOp` | array of type T The arrays `on_true` and `on_false` must have the same shape. This is also the shape of the output array. The array `pred` must have the same dimensionality as `on_true` and `on_false`, with the `PRED` element type. For each element `P` of `pred`, the corresponding element of the output array is taken from `on_true` if the value of `P` is `true`, and from `on_false` if the value of `P` is `false`. As a restricted form of [broadcasting](broadcasting.md), `pred` can be a scalar of type `PRED`. In this case, the output array is taken wholly from `on_true` if `pred` is `true`, and from `on_false` if `pred` is `false`. Example with non-scalar `pred`: ```cpp let pred: PRED[4] = {true, false, false, true}; let v1: s32[4] = {1, 2, 3, 4}; let v2: s32[4] = {100, 200, 300, 400}; ==> Select(pred, v1, v2) = s32[4]{1, 200, 300, 4}; ``` Example with scalar `pred`: ```cpp let pred: PRED = true; let v1: s32[4] = {1, 2, 3, 4}; let v2: s32[4] = {100, 200, 300, 400}; ==> Select(pred, v1, v2) = s32[4]{1, 2, 3, 4}; ``` Selections between tuples are supported. Tuples are considered to be scalar types for this purpose. If `on_true` and `on_false` are tuples (which must have the same shape!) then `pred` has to be a scalar of type `PRED`. For StableHLO information see [StableHLO - select](https://openxla.org/stablehlo/spec#select) ## SelectAndScatter See also [`XlaBuilder::SelectAndScatter`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). This operation can be considered as a composite operation that first computes `ReduceWindow` on the `operand` array to select an element from each window, and then scatters the `source` array to the indices of the selected elements to construct an output array with the same shape as the operand array. The binary `select` function is used to select an element from each window by applying it across each window, and it is called with the property that the first parameter's index vector is lexicographically less than the second parameter's index vector. The `select` function returns `true` if the first parameter is selected and returns `false` if the second parameter is selected, and the function must hold transitivity (i.e., if `select(a, b)` and `select(b, c)` are `true`, then `select(a, c)` is also `true`) so that the selected element does not depend on the order of the elements traversed for a given window. The function `scatter` is applied at each selected index in the output array. It takes two scalar parameters: 1. Current value at the selected index in the output array 2. The scatter value from `source` that applies to the selected index It combines the two parameters and returns a scalar value that's used to update the value at the selected index in the output array. Initially, all indices of the output array are set to `init_value`. The output array has the same shape as the `operand` array and the `source` array must have the same shape as the result of applying a `ReduceWindow` operation on the `operand` array. `SelectAndScatter` can be used to backpropagate the gradient values for a pooling layer in a neural network. **`SelectAndScatter(operand, select, window_dimensions, window_strides, padding, source, init_value, scatter)`** | Arguments | Type | Semantics | | ------------------- | ------------------- | -------------------------------- | | `operand` | `XlaOp` | array of type T over which the | : : : windows slide : | `select` | `XlaComputation` | binary computation of type `T, T | : : : -> PRED`, to apply to all : : : : elements in each window; returns : : : : `true` if the first parameter is : : : : selected and returns `false` if : : : : the second parameter is selected : | `window_dimensions` | `ArraySlice` | array of integers for window | : : : dimension values : | `window_strides` | `ArraySlice` | array of integers for window | : : : stride values : | `padding` | `Padding` | padding type for window | : : : (Padding\:\:kSame or : : : : Padding\:\:kValid) : | `source` | `XlaOp` | array of type T with the values | : : : to scatter : | `init_value` | `XlaOp` | scalar value of type T for the | : : : initial value of the output : : : : array : | `scatter` | `XlaComputation` | binary computation of type `T, T | : : : -> T`, to apply each scatter : : : : source element with its : : : : destination element : The figure below shows examples of using `SelectAndScatter`, with the `select` function computing the maximal value among its parameters. Note that when the windows overlap, as in the figure (2) below, an index of the `operand` array may be selected multiple times by different windows. In the figure, the element of value 9 is selected by both of the top windows (blue and red) and the binary addition `scatter` function produces the output element of value 8 (2 + 6). The evaluation order of the `scatter` function is arbitrary and may be non-deterministic. Therefore, the `scatter` function should not be overly sensitive to reassociation. See the discussion about associativity in the context of [`Reduce`](#reduce) for more details. For StableHLO information see [StableHLO - select_and_scatter](https://openxla.org/stablehlo/spec#select_and_scatter). ## Send See also [`XlaBuilder::Send`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). `Send`, `SendWithTokens`, and `SendToHost` are operations that serve as communication primitives in HLO. These ops typically appear in HLO dumps as part of low-level input/output or cross-device transfer, but they are not intended to be constructed manually by end users. **`Send(operand, handle)`** Arguments | Type | Semantics --------- | --------------- | ----------------------------------------- `operand` | `XlaOp` | data to send (array of type T) `handle` | `ChannelHandle` | unique identifier for each send/recv pair Sends the given operand data to a [`Recv`](#recv) instruction in another computation that shares the same channel handle. Does not return any data. Similar to the [`Recv`](#recv) operation, the client API of `Send` operation represents synchronous communication, and is internally decomposed into 2 HLO instructions (`Send` and `SendDone`) to enable asynchronous data transfers. See also [`HloInstruction::CreateSend` and `HloInstruction::CreateSendDone`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). **`Send(HloInstruction operand, int64 channel_id)`** Initiates an asynchronous transfer of the operand to the resources allocated by the `Recv` instruction with the same channel id. Returns a context, which is used by a following `SendDone` instruction to wait for the completion of the data transfer. The context is a tuple of {operand (shape), request identifier (U32)} and it can only be used by a `SendDone` instruction. For StableHLO information see [StableHLO - send](https://openxla.org/stablehlo/spec#send). ### SendDone See also [`HloInstruction::CreateSendDone`](https://github.com/openxla/xla/tree/main/xla/hlo/ir/hlo_instruction.h). **`SendDone(HloInstruction context)`** Given a context created by a `Send` instruction, waits for the data transfer to complete. The instruction does not return any data. > **Note:** `SendDone` is only found in HLO. It is not found in StableHLO. ### Scheduling of channel instructions The execution order of the 4 instructions for each channel (`Recv`, `RecvDone`, `Send`, `SendDone`) is as below. * `Recv` happens before `Send` * `Send` happens before `RecvDone` * `Recv` happens before `RecvDone` * `Send` happens before `SendDone` When the backend compilers generate a linear schedule for each computation that communicates via channel instructions, there must not be cycles across the computations. For example, below schedules lead to deadlocks. ## SetDimensionSize See also [`XlaBuilder::SetDimensionSize`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Sets the dynamic size of XlaOp's given dimension. The operand must be array shaped. **`SetDimensionSize(operand, val, dimension)`** | Arguments | Type | Semantics | | ----------- | ------- | --------------------------------------------------- | | `operand` | `XlaOp` | n dimensional input array. | | `val` | `XlaOp` | int32 representing the runtime dynamic size. | | `dimension` | `int64` | A value in the interval `[0, n)` that specifies the | : : : dimension. : Pass through the operand as result, with dynamic dimension tracked by the compiler. Padded values will be ignored by downstream reduction ops. ```cpp let v: f32[10] = f32[10]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; let five: s32 = 5; let six: s32 = 6; // Setting dynamic dimension size doesn't change the upper bound of the static // shape. let padded_v_five: f32[10] = set_dimension_size(v, five, /*dimension=*/0); let padded_v_six: f32[10] = set_dimension_size(v, six, /*dimension=*/0); // sum == 1 + 2 + 3 + 4 + 5 let sum:f32[] = reduce_sum(padded_v_five); // product == 1 * 2 * 3 * 4 * 5 let product:f32[] = reduce_product(padded_v_five); // Changing padding size will yield different result. // sum == 1 + 2 + 3 + 4 + 5 + 6 let sum:f32[] = reduce_sum(padded_v_six); ``` ## ShiftLeft See also [`XlaBuilder::ShiftLeft`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise left-shift operation on `lhs` by `rhs` number of bits. **`ShiftLeft(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for ShiftLeft: **`ShiftLeft(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - shift_left](https://openxla.org/stablehlo/spec#shift_left). ## ShiftRightArithmetic See also [`XlaBuilder::ShiftRightArithmetic`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise arithmetic right-shift operation on `lhs` by `rhs` number of bits. **`ShiftRightArithmetic(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for ShiftRightArithmetic: **`ShiftRightArithmetic(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - shift_right_arithmetic](https://openxla.org/stablehlo/spec#shift_right_arithmetic). ## ShiftRightLogical See also [`XlaBuilder::ShiftRightLogical`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise logical right-shift operation on `lhs` by `rhs` number of bits. **`ShiftRightLogical(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for ShiftRightLogical: **`ShiftRightLogical(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - shift_right_logical](https://openxla.org/stablehlo/spec#shift_right_logical). ## Shuffle See also [`XlaBuilder::Shuffle`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Shuffle(operand, dimensions, mode)`** Arguments | Type | Semantics ------------- | ------------------- | ------------------------------------- `operand` | `XlaOp` | array of type T `dimensions` | `ArraySlice` | dimensions to shuffle `mode` | `ShuffleMode` | the shuffle pattern to apply, together with the attributes of that pattern Shuffles the elements of the `operand` array along the specified `dimensions`, generating an output array of the same shape. Because a shuffle only moves elements around, it never changes the shape of the operand. The `mode` selects which pattern is applied, and carries the attributes of that pattern: Mode | Attributes | Semantics -------- | ---------- | --------------------------------------------------- `rotate` | `shifts` | Rotates the elements to the left by `shifts` ### `rotate` Rotates the elements of the `operand` array along the specified `dimensions` to the left by the corresponding `shifts`, generating an output array of the same shape. Shifts outside of `[0, size(dimension))` are supported, and wrap around cyclically. Each element of the output array at a multidimensional index is retrieved from the operand array at a transformed index. The multidimensional index is transformed by rotating the index along each specified dimension (i.e., for a dimension of size N with a corresponding shift S, the element at output index i is taken from operand index (i + S) % N). 1-dimensional example: ```cpp let a = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; Shuffle(a, {0}, shuffle::Rotate({4})) // Result: {4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3} ``` ## Sign See also [`XlaBuilder::Sign`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Sign(operand)`** Element-wise sign operation `x -> sgn(x)` where $$\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}$$ using the comparison operator of the element type of `operand`. **`Sign(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function For StableHLO information see [StableHLO - sign](https://openxla.org/stablehlo/spec#sign). ## Sin **`Sin(operand)`** Element-wise sine `x -> sin(x)`. See also [`XlaBuilder::Sin`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Sin(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Sin also supports the optional `result_accuracy` argument: **`Sin(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - sine](https://openxla.org/stablehlo/spec#sine). ## Slice See also [`XlaBuilder::Slice`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Slicing extracts a sub-array from the input array. The sub-array has the same number of dimensions as the input and contains the values inside a bounding box within the input array where the dimensions and indices of the bounding box are given as arguments to the slice operation. **`Slice(operand, start_indices, limit_indices, strides)`** | Arguments | Type | Semantics | | --------------- | ------------------- | ------------------------------------ | | `operand` | `XlaOp` | N dimensional array of type T | | `start_indices` | `ArraySlice` | List of N integers containing the | : : : starting indices of the slice for : : : : each dimension. Values must be : : : : greater than or equal to zero. : | `limit_indices` | `ArraySlice` | List of N integers containing the | : : : ending indices (exclusive) for the : : : : slice for each dimension. Each value : : : : must be greater than or equal to the : : : : respective `start_indices` value for : : : : the dimension and less than or equal : : : : to the size of the dimension. : | `strides` | `ArraySlice` | List of N integers that decides the | : : : input stride of the slice. The slice : : : : picks every `strides[d]` element in : : : : dimension `d`. : 1-dimensional example: ```cpp let a = {0.0, 1.0, 2.0, 3.0, 4.0} Slice(a, {2}, {4}) // Result: {2.0, 3.0} ``` 2-dimensional example: ```cpp let b = { {0.0, 1.0, 2.0}, {3.0, 4.0, 5.0}, {6.0, 7.0, 8.0}, {9.0, 10.0, 11.0} } Slice(b, {2, 1}, {4, 3}) // Result: // { { 7.0, 8.0}, // {10.0, 11.0} } ``` For StableHLO information see [StableHLO - slice](https://openxla.org/stablehlo/spec#slice). ## Sort See also [`XlaBuilder::Sort`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`Sort(operands, comparator, dimension, is_stable)`** Arguments | Type | Semantics ------------ | ------------------- | -------------------------------------- `operands` | `ArraySlice` | The operands to sort. `comparator` | `XlaComputation` | The comparator computation to use. `dimension` | `int64` | The dimension along which to sort. `is_stable` | `bool` | Whether stable sorting should be used. If only one operand is provided: * If the operand is a 1-dimensional tensor (an array), the result is a sorted array. If you want to sort the array into ascending order, the comparator should perform a less-than comparison. Formally, after the array is sorted, it holds for all index positions `i, j` with `i < j` that either `comparator(value[i], value[j]) = comparator(value[j], value[i]) = false` or `comparator(value[i], value[j]) = true`. * If the operand has higher number of dimensions, the operand is sorted along the provided dimension. For example, for a 2-dimensional tensor (a matrix), a dimension value of `0` will independently sort every column, and a dimension value of `1` will independently sort each row. If no dimension number is provided, then the last dimension is chosen by default. For the dimension which is sorted, the same sorting order applies as in the 1-dimensional case. If `n > 1` operands are provided: * All `n` operands must be tensors with the same dimensions. The element types of the tensors may be different. * All operands are sorted together, not individually. Conceptually the operands are treated as a tuple. When checking whether the elements of each operand at index positions `i` and `j` need to be swapped, the comparator is called with `2 * n` scalar parameters, where parameter `2 * k` corresponds to the value at position `i` from the `k-th` operand, and parameter `2 * k + 1` corresponds to the value at position `j` from the `k-th` operand. Usually, the comparator would thus compare parameters `2 * k` and `2 * k + 1` with each other and possibly use other parameter pairs as tie breakers. * The result is a tuple that consists of the operands in sorted order (along the provided dimension, as above). The `i-th` operand of the tuple corresponds to the `i-th` operand of Sort. For example, if there are three operands `operand0 = [3, 1]`, `operand1 = [42, 50]`, `operand2 = [-3.0, 1.1]`, and the comparator compares only the values of `operand0` with less-than, then the output of the sort is the tuple `([1, 3], [50, 42], [1.1, -3.0])`. If `is_stable` is set to true, the sort is guaranteed to be stable, that is, if there are elements which are considered to be equal by the comparator, the relative order of the equal values is preserved. Two elements `e1` and `e2` are equal if and only if `comparator(e1, e2) = comparator(e2, e1) = false`. By default, `is_stable` is set to false. For StableHLO information see [StableHLO - sort](https://openxla.org/stablehlo/spec#sort). ## Sqrt See also [`XlaBuilder::Sqrt`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise square root operation `x -> sqrt(x)`. **`Sqrt(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Sqrt also supports the optional `result_accuracy` argument: **`Sqrt(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - sqrt](https://openxla.org/stablehlo/spec#sqrt). ## Sub See also [`XlaBuilder::Sub`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise subtraction of `lhs` and `rhs`. **`Sub(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Sub: **`Sub(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - subtract](https://openxla.org/stablehlo/spec#subtract). ## Tan See also [`XlaBuilder::Tan`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise tangent `x -> tan(x)`. **`Tan(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Tan also supports the optional `result_accuracy` argument: **`Tan(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - tan](https://openxla.org/stablehlo/spec#tan). ## Tanh See also [`XlaBuilder::Tanh`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Element-wise hyperbolic tangent `x -> tanh(x)`. **`Tanh(operand)`** Arguments | Type | Semantics --------- | ------- | --------------------------- `operand` | `XlaOp` | The operand to the function Tanh also supports the optional `result_accuracy` argument: **`Tanh(operand, result_accuracy)`** | Arguments | Type | Semantics | | ----------------- | ------------------------- | --------------------------- | | `operand` | `XlaOp` | The operand to the function | | `result_accuracy` | optional `ResultAccuracy` | The types of accuracy the | : : : user can request for unary : : : : ops with multiple : : : : implementations : For more information on `result_accuracy` see [Result Accuracy](https://github.com/openxla/stablehlo/blob/main/rfcs/20241015-result-accuracy.md). For StableHLO information see [StableHLO - tanh](https://openxla.org/stablehlo/spec#tanh). ## TopK See also [`XlaBuilder::TopK`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). `TopK` finds the values and indices of the `k` largest or smallest elements for the last dimension of the given tensor. **`TopK(operand, k, largest)`** | Arguments | Type | Semantics | | --------- | ------- | ------------------------------------------------------ | | `operand` | `XlaOp` | The tensor from which to extract the top `k` elements. | : : : The tensor must have greater or equal to one : : : : dimensions. The size of the last dimension of the : : : : tensor must be greater or equal to `k`. : | `k` | `int64` | The number of elements to extract. | | `largest` | `bool` | Whether to extract the largest or smallest `k` | : : : elements. : For a 1-dimensional input tensor (an array), finds the `k` largest or smallest entries in the array and outputs a tuple of two arrays `(values, indices)`. Thus `values[j]` is the `j`-th largest/smallest entry in `operand`, and its index is `indices[j]`. For an input tensor with more than 1 dimension, computes the top `k` entries along the last dimension, preserving all other dimensions (rows) in the output. Thus, for an operand of shape `[A, B, ..., P, Q]` where `Q >= k` the output is a tuple `(values, indices)` where: ``` values.shape = indices.shape = [A, B, ..., P, k] ``` If two elements within a row are equal, the lower-index element appears first. > **Note:** `TopK` is only found in HLO and is not found in StableHLO. CHLO > `TopK` in Frameworks will lower to HLO `TopK` see > [StableHLO - chlo.top_k](https://openxla.org/stablehlo/generated/chlo?hl=en#chlotop_k_chlotopkop) ## Transpose See also the `tf.reshape` operation. **`Transpose(operand, permutation)`** Arguments | Type | Semantics ------------- | ------------------- | ------------------------------ `operand` | `XlaOp` | The operand to transpose. `permutation` | `ArraySlice` | How to permute the dimensions. Permutes the operand dimensions with the given permutation, so `∀ i . 0 ≤ i < number of dimensions ⇒ input_dimensions[permutation[i]] = output_dimensions[i]`. This is the same as Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)). For StableHLO information see [StableHLO - transpose](https://openxla.org/stablehlo/spec#transpose). ## TriangularSolve See also [`XlaBuilder::TriangularSolve`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Solves systems of linear equations with lower or upper triangular coefficient matrices by forward- or back-substitution. Broadcasting along leading dimensions, this routine solves one of the matrix systems `op(a) * x = b`, or `x * op(a) = b`, for the variable `x`, given `a` and `b`, where `op(a)` is either `op(a) = a`, or `op(a) = Transpose(a)`, or `op(a) = Conj(Transpose(a))`. **`TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)`** | Arguments | Type | Semantics | | --------------- | ----------- | ------------------------------------------- | | `a` | `XlaOp` | a > 2 dimensional array of a complex or | : : : floating-point type with shape `[..., M, : : : : M]`. : | `b` | `XlaOp` | a > 2 dimensional array of the same type | : : : with shape `[..., M, K]` if `left_side` is : : : : true, `[..., K, M]` otherwise. : | `left_side` | `bool` | indicates whether to solve a system of the | : : : form `op(a) * x = b` (`true`) or `x * : : : : op(a) = b` (`false`). : | `lower` | `bool` | whether to use the upper or lower triangle | : : : of `a`. : | `unit_diagonal` | `bool` | if `true`, the diagonal elements of `a` are | : : : assumed to be `1` and not accessed. : | `transpose_a` | `Transpose` | whether to use `a` as is, transpose it or | : : : take its conjugate transpose. : Input data is read only from the lower/upper triangle of `a`, depending on the value of `lower`. Values from the other triangle are ignored. Output data is returned in the same triangle; the values in the other triangle are implementation-defined and may be anything. If the number of dimensions of `a` and `b` are greater than 2, they are treated as batches of matrices, where all except the minor 2 dimensions are batch dimensions. `a` and `b` must have equal batch dimensions. For StableHLO information see [StableHLO - triangular_solve](https://openxla.org/stablehlo/spec#triangular_solve). ## Tuple See also [`XlaBuilder::Tuple`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). A tuple containing a variable number of data handles, each of which has its own shape. **`Tuple(elements)`** Arguments | Type | Semantics ---------- | ----------------- | ----------------- `elements` | vector of `XlaOp` | N array of type T This is analogous to `std::tuple` in C++. Conceptually: ```cpp let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; let s: s32 = 5; let t: (f32[10], s32) = tuple(v, s); ``` Tuples can be deconstructed (accessed) via the [`GetTupleElement`] (#gettupleelement) operation. For StableHLO information see [StableHLO - tuple](https://openxla.org/stablehlo/spec#tuple). > **Note:** In HLO, tuples are needed for most ops that return >1 result. While > in StableHLO/MLIR, variadic results can be expressed and tuples are not used, > except in custom_calls/get_tuple_element. ## While See also [`XlaBuilder::While`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). **`While(condition, body, init)`** | Arguments | Type | Semantics | | :---------- | :--------------- | :--------------------------------------- | | `condition` | `XlaComputation` | XlaComputation of type `T -> PRED` which | : : : defines the termination condition of the : : : : loop. : | `body` | `XlaComputation` | XlaComputation of type `T -> T` which | : : : defines the body of the loop. : | `init` | `T` | Initial value for the parameter of | : : : `condition` and `body`. : Sequentially executes the `body` until the `condition` fails. This is similar to a typical while loop in many other languages except for the differences and restrictions listed below. * A `While` node returns a value of type `T`, which is the result from the last execution of the `body`. * The shape of the type `T` is statically determined and must be the same across all iterations. The T parameters of the computations are initialized with the `init` value in the first iteration and are automatically updated to the new result from `body` in each subsequent iteration. One main use case of the `While` node is to implement the repeated execution of training in neural networks. Simplified pseudocode is shown below with a graph that represents the computation. The code can be found in [`while_test.cc`](https://github.com/openxla/xla/tree/main/xla/tests/while_test.cc). The type `T` in this example is a `Tuple` consisting of an `int32` for the iteration count and a `vector[10]` for the accumulator. For 1000 iterations, the loop keeps adding a constant vector to the accumulator. ```cpp // Pseudocode for the computation. init = {0, zero_vector[10]} // Tuple of int32 and float[10]. result = init; while (result(0) < 1000) { iteration = result(0) + 1; new_vector = result(1) + constant_vector[10]; result = {iteration, new_vector}; } ``` For StableHLO information see [StableHLO - while](https://openxla.org/stablehlo/spec#while). ## Xor See also [`XlaBuilder::Xor`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). Performs element-wise XOR of `lhs` and `rhs`. **`Xor(lhs, rhs)`** Arguments | Type | Semantics --------- | ----- | --------------------------------------- lhs | XlaOp | Left-hand-side operand: array of type T rhs | XlaOp | Left-hand-side operand: array of type T The arguments' shapes have to be either similar or compatible. See the [broadcasting](broadcasting.md) documentation about what it means for shapes to be compatible. The result of an operation has a shape which is the result of broadcasting the two input arrays. In this variant, operations between arrays of different ranks are *not* supported, unless one of the operands is a scalar. An alternative variant with different-dimensional broadcasting support exists for Xor: **`Xor(lhs,rhs, broadcast_dimensions)`** | Arguments | Type | Semantics | | ------------------- | ----------------- | -------------------------------- | | lhs | XlaOp | Left-hand-side operand: array of | : : : type T : | rhs | XlaOp | Left-hand-side operand: array of | : : : type T : | broadcast_dimension | ArraySlice | Which dimension in the target | : : : shape each dimension of the : : : : operand shape corresponds to : This variant of the operation should be used for arithmetic operations between arrays of different ranks (such as adding a matrix to a vector). The additional broadcast_dimensions operand is a slice of integers specifying the dimensions to use for broadcasting the operands. The semantics are described in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - xor](https://openxla.org/stablehlo/spec#xor). --- ### Persisted Autotuning (`docs/persisted_autotuning.md`) # Persisted autotuning (GPU only) We use OpenAI Triton for generating some of the GPU kernels. Triton allows generating fast GPU kernels for certain fusions, but we have to tune some parameters for each such fusion. This can take a long time if there are many fusions, so we provide a way to load those autotuning results, while still running the other compilation steps normally. Autotuning caches are still useful if we make a few changes: the fusions that are present in the cache will use the cache, and the other ones will be autotuned normally. ## Recommended: Cache directory ``` --xla_gpu_per_fusion_autotune_cache_dir=your/directory ``` Use and maintain a per-fusion autotune cache in the given directory. There will be one file per distinct fusion. The main advantage of this approach is that you can use the same cache directory for multiple XLA runs (of different models) and your cache will grow with each new fusion encountered - speeding up subsequent runs. There is also basic support for running multiple XLA instances with the same cache directory concurrently. XLA will read existing results when they are needed and write new results after they are determined. - The directory must exist before running XLA and it must be writable. - Cache invalidation has to be handled by the user: - Please use an empty directory if you want to start with an empty cache. - XLA version checks must be done by the user: - If you want to use separate caches for different versions of XLA, please use different directories. The cache is turned off by default (when you don't provide the parameter). Limitation: This is not guaranteed to work well in combination with the other caching method described below. ## Alternative: Loading or dumping all results from a given HLO to one file The autotuning results can be dumped/loaded using these parameters: ``` --xla_gpu_dump_autotune_results_to= --xla_gpu_load_autotune_results_from= ``` If we specify a .txt or .textproto file, then the cache will be dumped in textproto format, otherwise in binary protobuf format. ## In tests Persisted autotuning can also be used in tests. It is recommended to use it if the tests are very big, especially if the performance of the test environment is limited. It only works well if the autotune cache contains results generated on the same type of GPU where the tests are being run. ### Making a test use persisted autotuning For now let's assume that the test in question always uses the same GPU type. 1. We have to export the autotune results from the test, for example by specifying these parameters to the test command: ``` --test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabled ``` Sharding must be disabled to correctly get a single autotune cache for all tests. 2. Then we have to upload that cache to our code repository. 3. Then we have to add the cache to the data dependencies of our test target, and load it using an environment variable. ``` data = ["test_autotune_cache.textproto"], env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" + "$(execpath test_autotune_cache.textproto)"}, ``` (It is OK to use sharding in tests that load autotune results.) Please also see the example tests in [xla/backends/gpu/tests/BUILD](https://github.com/openxla/xla/blob/main/xla/backends/gpu/tests/BUILD): - load_autotune_results_using_execpath_test - load_autotune_results_from_test_workspace_test - dump_autotune_results_to_test_outputs_test ### Cache obsolescence If many changes are made to a model, it is possible that the cache will no longer contain all fusions, so the test will become slower. In this case we would have to regenerate the autotuning cache. If we start using a new type of GPU for running the tests, the same applies. The cache may also become obsolete if the XLA compiler evolves and generates different fusions. --- ### Rocm Ci (`docs/rocm_ci.md`) # AMD ROCm Continuous Integration This document describes the AMD ROCm continuous integration (CI) implementation. It covers the GitHub Actions workflow, runner pool, Docker containers, Bazel configurations, tag filters, GPU locking, and the upstream-fork relationship with [`ROCm/xla`](https://github.com/ROCm/xla). The setup comprises two distinct paths: 1. **ROCm CI workflow** (`.github/workflows/rocm_ci.yml`): Runs physical XLA and JAX tests on a self-hosted AMD Instinct GPU runner. 2. **Compile-only ROCm CI** (`.github/workflows/ci.yml`): Builds XLA against a hermetic ROCm toolchain on a generic Linux x86 runner without a GPU. These paths ensure both runtime verification on hardware and compilation reliability on every commit. ## Trigger and concurrency `rocm_ci.yml` runs on: - `pull_request` against `main` (the `jax` job is additionally gated to `github.base_ref == 'main'`). - `workflow_dispatch` (manual). The compile-only ROCm path inside `ci.yml` runs on the same `pull_request` events as the rest of the CI matrix and at post-submit. Both workflows share a concurrency group keyed by workflow + head ref, so a new push to a PR cancels any in-flight run on that PR but does not cancel runs on `main`: ```yaml concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} cancel-in-progress: ${{ github.ref != 'main' }} ``` Permissions are restricted to `contents: read`. No write scopes or artifact uploads are enabled. ## Runner and container configuration ### Runner pool The ROCm jobs use a self-hosted GitHub Actions runner labelled `linux-x86-64-1gpu-amd`. This is a Linux x86_64 host with a single AMD Instinct GPU attached, currently targeting `gfx950`. The label links to the AMD-owned RBE runner fleet, where the build and tests run. The status of this cluster can be monitored on [this dashboard](https://wardite.cluster.engflow.com/). ### Docker image The jobs run in a Docker container. The Docker image is provided by AMD and derived from the official Tensorflow build images. **Key points:** - `--device=/dev/kfd` and `--device=/dev/dri` expose the AMD GPU kernel fusion driver and DRI render nodes to the container. Without these the container cannot see the GPU. - `--group-add video` puts the container's user in the `video` group, which is the convention on Ubuntu for users that may access GPU device nodes. - `--ipc=host` and `--shm-size=64G` are needed because HIP/ROCm collectives (RCCL) and several stream-executor pathways use shared memory for inter-process communication on a single node. - `--cap-add=SYS_PTRACE` and `--security-opt=seccomp=unconfined` enable sanitizer-instrumented test binaries to attach debuggers and bypass seccomp filters that would block sanitizer runtime calls. - `--tmpfs /root/.cache/bazel:rw,exec,size=80g` puts the Bazel disk cache in RAM. `exec` is mandatory because Bazel runs `bazel-out` binaries directly out of the cache. - The two `ci-cert.*` files mounted in from `/data` on the host are the client TLS certificate and private key used to authenticate to the EngFlow remote build cluster. ## Job topology `rocm_ci.yml` defines three jobs: | Job | Purpose | Depends on | | ------------- | ------------------------------------------ | ------------- | | `rocm-config` | Publish the pinned Docker image digest. | — | | `jax` | Build and test JAX against changes to XLA, | `rocm-config` | : : on a single AMD GPU. : : | `xla` | Build and test XLA itself, on a single AMD | `rocm-config` | : : GPU, plus a CPU-only XLA suite. : : `jax` and `xla` run in parallel once `rocm-config` finishes. The `xla` job is *not* gated on `github.base_ref == 'main'`; the `jax` job is. This means the JAX suite only runs for PRs targeted at `main`, while the XLA suite runs for any PR plus `workflow_dispatch`. `timeout-minutes` is set at the job level, with a tighter `timeout-minutes: 60` (XLA single-GPU, JAX) or 80 (XLA CPU suite) on each test step. ## XLA CI job ### Step sequence 1. Check out `openxla/xla` at the PR head. 2. Download the AMD-maintained driver script (`execute_ci_build_upstream.sh`) from the `ROCm/xla` fork. 3. Print CPU (`lscpu`) and GPU (`rocminfo`) information for debuggability. 4. Run the test phases (Single-GPU and CPU) using the downloaded script. ### What `--config=ci_single_gpu` and `--config=ci_rocm_cpu` mean These are defined in `build_tools/rocm/rocm_xla.bazelrc`. The single-GPU suite uses the parallel GPU helper, retries flakes three times, and runs tests in a process-isolated way. The CPU suite uses 200-wide local parallelism and the sanitizer wrapper: ``` build:ci_single_gpu --run_under=//build_tools/rocm:parallel_gpu_execute build:ci_single_gpu --flaky_test_attempts=3 build:ci_rocm_cpu --run_under=//build_tools/rocm:sanitizer_wrapper build:ci_rocm_cpu --local_test_jobs=200 build:ci_rocm_cpu --strategy=TestRunner=local ``` A third config, `ci_multi_gpu`, exists in the bazelrc but is not currently triggered from `rocm_ci.yml` — it is used by the standalone helper `run_xla_multi_gpu.sh` on hosts with ≥4 GPUs. ### Tag filtering `execute_ci_build_upstream.sh` (in the `ROCm/xla` fork) composes a tag filter list using the helper at `build_tools/rocm/rocm_tag_filters.sh`. The legacy in-tree wrapper `build_tools/rocm/run_xla_ci_build.sh` shows the same pattern: ```bash TAG_FILTERS=$($SCRIPT_DIR/rocm_tag_filters.sh) for arg in "$@"; do if [[ "$arg" == "--config=ci_multi_gpu" ]]; then TAG_FILTERS="${TAG_FILTERS},requires-gpu-rocm,requires-gpu-amd,multi_gpu" fi if [[ "$arg" == "--config=ci_single_gpu" ]]; then TAG_FILTERS="${TAG_FILTERS},requires-gpu-rocm,requires-gpu-amd,-multi_gpu" fi if [[ "$arg" == "--config=ci_rocm_cpu" ]]; then TAG_FILTERS="${TAG_FILTERS},gpu,-requires-gpu-rocm,-requires-gpu-amd" fi done ``` The base list from `rocm_tag_filters.sh` is the "things ROCm CI should *never* attempt": ``` -no_gpu -requires-gpu-intel -requires-gpu-nvidia -requires-gpu-cuda -cuda-only -oneapi-only -requires-gpu-sm60 -requires-gpu-sm60-only -requires-gpu-sm70 -requires-gpu-sm70-only -requires-gpu-sm80 -requires-gpu-sm80-only -requires-gpu-sm86 -requires-gpu-sm86-only -requires-gpu-sm89 -requires-gpu-sm89-only -requires-gpu-sm90 -requires-gpu-sm90-only -skip_rocprofiler_sdk -no_oss -oss_excluded -oss_serial ``` The per-config additions then narrow the set: Config | Adds --------------- | ----------------------------------------------------- `ci_single_gpu` | `requires-gpu-rocm`, `requires-gpu-amd`, `-multi_gpu` `ci_multi_gpu` | `requires-gpu-rocm`, `requires-gpu-amd`, `multi_gpu` `ci_rocm_cpu` | `gpu`, `-requires-gpu-rocm`, `-requires-gpu-amd` Note: The `gpu` tag includes tests that touch GPU code paths but can execute on CPU. The CPU portion excludes tests requiring a physical AMD/ROCm GPU, allowing it to exercise GPU codegen and the autotuner harness against the host CPU. ### Tests excluded by name Collectives, sharding propagation, and distributed PJRT require multiple AMD GPUs and are filtered from single-GPU runs. These must be executed on a dedicated multi-GPU host. ## JAX CI job The JAX job runs in the same container but with a different driver script, provided by JAX itself: ```bash ./ci/run_bazel_test_rocm_rbe.sh \ --override_repository=xla="${GITHUB_WORKSPACE}" \ --override_module=xla="${GITHUB_WORKSPACE}" \ --config=single_gpu \ --//jax:build_jaxlib=wheel \ --//jax:build_jax=true \ --local_test_jobs=1 \ --action_env=JAX_ENABLE_X64=0 \ --repo_env=HERMETIC_PYTHON_VERSION=3.14 \ --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" ``` **Flag details:** - `--override_repository=xla=...` and `--override_module=xla=...` redirect JAX's `xla` dependency at the in-tree checkout of the PR. Without these the JAX job would compile against the pinned XLA commit recorded in JAX's WORKSPACE/MODULE.bazel, which would defeat the point of running JAX from an XLA PR. - `--//jax:build_jaxlib=wheel` and `--//jax:build_jax=true` build jaxlib as a wheel artifact and rebuild jax itself from source — this is how JAX currently bootstraps the C++ extension for tests. - `HERMETIC_PYTHON_VERSION=3.14` selects the Python toolchain hermetically (the build is reproducible regardless of the host Python). - `--local_test_jobs=1` keeps the GPU contention low; with only one AMD GPU on the runner there is no benefit to running JAX tests in parallel. - `JAX_ENABLE_X64=0` matches the JAX team's default test posture for fast CI. - The same `TF_ROCM_RBE_DOCKER_IMAGE` is passed down so that any remote worker the build leases from EngFlow runs the same ROCm image. ## Remote build execution (EngFlow) ROCm CI uses an EngFlow remote build cluster (`grpcs://wardite.cluster.engflow.com`) for both remote caching and remote execution. The configuration lives in `build_tools/rocm/rocm_xla.bazelrc`. Authentication to this cluster is handled via mounted certificates (`/data/ci-cert.crt`, `/data/ci-cert.key`). Key configurations: - `REMOTE_GPU_TESTING=1`: Enables remote-execution optimizations (e.g., restricted filesystem access). - `@local_config_rocm//rocm:linux_x64`: Specifies the platform, assuming system-level ROCm installation on workers to resolve the `@local_config_rocm` repository. There is also a `rocm_rbe_dynamic` config that builds locally but tests under Bazel's dynamic scheduler (running each test action concurrently locally and remotely, taking whichever finishes first): ``` build:rocm_rbe_dynamic --config=rocm_rbe build:rocm_rbe_dynamic --spawn_strategy=local test:rocm_rbe_dynamic --experimental_spawn_scheduler test:rocm_rbe_dynamic --strategy=TestRunner=dynamic test:rocm_rbe_dynamic --dynamic_mode=default test:rocm_rbe_dynamic --dynamic_local_strategy=worker,standalone,local test:rocm_rbe_dynamic --dynamic_remote_strategy=remote test:rocm_rbe_dynamic --experimental_local_execution_delay=1000 test:rocm_rbe_dynamic --local_resources=cpu=HOST_CPUS*0.5 ``` The XLA single-GPU step in `rocm_ci.yml` uses essentially this approach by passing `--internal_spawn_scheduler --strategy=TestRunner=dynamic` directly on the command line on top of `--config=rocm_rbe`. The third "umbrella" config, `rocm_ci`, is loaded via: ``` # rocm_xla_ci.bazelrc try-import /usertools/rocm.bazelrc try-import %workspace%/build_tools/rocm/rocm_xla.bazelrc ``` `/usertools/rocm.bazelrc` is provided by the AMD-built Docker image (it encodes the host's ROCm install paths); the in-tree bazelrc is layered on top. `--config=rocm_ci` therefore means "use the in-image ROCm toolchain plus the in-tree XLA configs". ## GPU locking When a GPU test is run on the AMD host, Bazel invokes it via `--run_under=//build_tools/rocm:parallel_gpu_execute`. The script (`build_tools/rocm/parallel_gpu_execute.sh`) does three things: 1. Figure out how many GPUs the runner actually has by calling `rocminfo` and counting `Name: *gfx*` lines. ```bash ROCMINFO=$(find -L "${TEST_SRCDIR:-.}" -name "rocminfo" -path "*/bin/rocminfo" | head -n 1) TF_GPU_COUNT=$($ROCMINFO | grep "Name: *gfx*" | wc -l) ``` `rocminfo` is brought into the action's runfiles via the Bazel target's `data = ["//tensorflow/third_party/rocm/google:rocminfo"]` in `build_tools/rocm/BUILD`, so it is discoverable inside the remote-execution sandbox. 2. If `TF_GPU_COUNT == 0` (e.g. on the RBE default pool, which has no GPU) the script just `exec "$@"` — the test still runs, but without any per-GPU isolation. Tests that genuinely need a GPU are expected to self-fail; this branch exists so that GPU-tagged tests that *can* actually run on CPU still execute on cheap workers. 3. Otherwise, acquire a `flock`-based slot, one per `(gpu, slot_within_gpu)` pair, and export `CUDA_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` to that GPU index for the duration of the test: ```bash for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do for i in `seq 0 $((TF_GPU_COUNT-1))`; do exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1 if flock -n "$lock_fd"; then ( export CUDA_VISIBLE_DEVICES=$i export HIP_VISIBLE_DEVICES=$i "$TEST_BINARY" "$@" ) … fi done done ``` Slots are filled across GPUs sequentially (e.g., Slot 0 on all GPUs first, then Slot 1) to minimize contention before oversubscribing memory. There is also a near-identical copy at `build_tools/ci/parallel_gpu_execute.sh`, which is the CUDA-oriented sibling used by the NVIDIA jobs. They differ mainly in that the ROCm version actually consults `rocminfo` at runtime, whereas the CUDA version trusts `TF_GPU_COUNT=4` by default. A second helper, `//build_tools/rocm:sanitizer_wrapper`, is used by `ci_rocm_cpu` and `ci_multi_gpu`. It is a tiny generated shell script (`echo '#!/bin/bash' > $@; echo 'exec "$$@"' >> $@`) whose only job is to declare the sanitizer ignore lists as runfiles, so that changes to those ignore lists force the affected test actions to re-run. ## Standalone helpers Three scripts in `build_tools/rocm/` are designed to be run by hand against a local AMD machine and follow the same conventions as CI: - **`run_xla.sh`** — single-GPU XLA test sweep. Detects GPU count via `rocm-smi`, computes `N_TEST_JOBS = TF_GPU_COUNT * TF_TESTS_PER_GPU`, pulls `gfx*` ID from `rocminfo` to set `TF_ROCM_AMDGPU_TARGETS`, and invokes `bazel test --config=rocm_ci --config=xla_sgpu …` with a similar tag filter recipe as CI (though they differ slightly in implementation details). Uses `--run_under=//build_tools/ci:parallel_gpu_execute` (the CUDA-style wrapper). - **`run_xla_multi_gpu.sh`** — multi-GPU XLA test sweep. Requires `TF_GPU_COUNT ≥ 4`, otherwise it exits silently. Uses `--config=xla_mgpu`, sets `NCCL_MAX_NCHANNELS=1`, and notably does *not* pass `--run_under=//build_tools/ci:parallel_gpu_execute` — collective tests need to see all the GPUs at once, so per-test GPU pinning would break them. - **`run_xla_ci_build.sh`** — the in-tree legacy version of what the CI now fetches from `ROCm/xla`. Useful as documentation for how `--config=…` arguments map to tag filters. These three exist for AMD developers reproducing CI behaviour locally. ## Compile-only ROCm CI Independently of the GPU runner, `.github/workflows/ci.yml` runs an `XLA Linux x86 GPU ROCm` job on a generic Linux x86 runner with no GPU and no ROCm system install. This ensures compilation reliability on all PRs independently of GPU runner availability. The runner pool and container for this job are: ```yaml { pool: "linux-x86-n4-16", container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest", name: "XLA Linux x86 GPU ROCm", repo: "openxla/xla", } ``` A single matrix entry; no `--device=/dev/kfd`; no AMD certificate volumes. The job dispatches to `build_tools/ci/build.py`, which is the configuration-as-data driver shared with every other CI matrix entry: ```yaml - run: | "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py \ --build="${{ matrix.job_info.name }}_github_actions" ``` The `XLA_LINUX_X86_GPU_ROCM_GITHUB_ACTIONS` build is defined in `build_tools/ci/build.py` as: ```python Build( type_=BuildType.XLA_LINUX_X86_GPU_ROCM_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "rbe_linux_cpu", "rocm_clang_hermetic"), target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, build_tag_filters=rocm_tag_filter, test_tag_filters=rocm_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, subcommand="build", ) ``` **Configuration details:** - `subcommand="build"` — no tests are run. - `configs=("warnings", "rbe_linux_cpu", "rocm_clang_hermetic")` — the build uses the CPU RBE pool (no GPU workers needed) and the `rocm_clang_hermetic` config, which pulls in a hermetic ROCm + Clang toolchain rather than depending on `/opt/rocm` being present. - `rocm_tag_filter` mirrors the runtime filter set used by `rocm_ci.yml` *plus* an explicit `"gpu"` include, since `build.py` does not have the `rocm_tag_filters.sh` script in scope. `build.py` is just a config-as-data factory: it emits three Bazel commands per build — dry-run (`--nobuild`) retry, real build, and `bazel analyze-profile profile.json.gz`: ```python def commands(self) -> List[List[str]]: cmds = [] cmds.extend(self.extra_setup_commands) if not (macos_build or windows_build): cmds.append(retry(self.bazel_command(subcommand="build", extra_options=("--nobuild",)))) cmds.append(self.bazel_command(subcommand=self.subcommand)) cmds.append(["bazel", "analyze-profile", "profile.json.gz"]) return cmds ``` The dry-run-then-build pattern allows retrying transient package-fetch failures without rebuilding. This phase ensures compilation regressions are caught even if GPU runners are unavailable. ## Sanitizer support `build_tools/rocm/rocm_xla.bazelrc` also declares ASan and TSan configs: ``` build:tsan --strip=never build:tsan --copt -fsanitize=thread build:tsan --copt -g build:tsan --copt -fno-omit-frame-pointer build:tsan --linkopt -fsanitize=thread build:tsan --linkopt -g build:tsan --//build_tools/rocm:sanitizer=tsan build:tsan --test_env=TSAN_OPTIONS=suppressions=build_tools/rocm/tsan_ignore_list.txt:… build:tsan --run_under=//build_tools/rocm:sanitizer_wrapper build:asan --test_env=ASAN_OPTIONS=suppressions=build_tools/rocm/asan_ignore_list.txt:use_sigaltstack=0 build:asan --test_env=LSAN_OPTIONS=suppressions=build_tools/rocm/lsan_ignore_list.txt:use_sigaltstack=0 build:asan --//build_tools/rocm:sanitizer=asan build:asan --run_under=//build_tools/rocm:sanitizer_wrapper ``` The `//build_tools/rocm:sanitizer` string flag is consumed by `select({"asan": [...], "tsan": [...]})` in the sanitizer_wrapper filegroup to bring the appropriate suppression file into the test action's runfiles. These configs are not engaged by `rocm_ci.yml` today; they are available for ad-hoc invocations and for downstream pipelines. ## Execution flow For a typical PR landing on `openxla/xla`: ``` PR commit │ ├─► .github/workflows/ci.yml ───────────────► matrix of CPU/GPU/ROCm builds │ │ │ └─► XLA Linux x86 GPU ROCm │ └─► build_tools/ci/build.py │ └─► bazel build (nobuild → real) │ configs: warnings + rbe_linux_cpu + │ rocm_clang_hermetic │ tag filter: rocm_tag_filter (compile-only) │ └─► .github/workflows/rocm_ci.yml ──────────► AMD GPU runner │ ├─► rocm-config: pin Docker image digest │ ├─► jax (1 GPU) │ └─► jax/ci/run_bazel_test_rocm_rbe.sh │ └─► bazel test --config=single_gpu … │ override_module=xla= │ EngFlow RBE │ └─► xla (1 GPU) ├─► wget execute_ci_build_upstream.sh from ROCm/xla ├─► bazel test --config=rocm_ci --config=rocm_rbe │ --config=ci_single_gpu │ (dynamic scheduling, parallel_gpu_execute lock, │ flock per CUDA_VISIBLE_DEVICES, 3 retries) └─► bazel test --config=rocm_ci --config=rocm_rbe --config=ci_rocm_cpu (200-way local, sanitizer_wrapper, no GPU required but exercises GPU-tagged tests on CPU) ``` If anything fails: - An EngFlow link (`https://wardite.cluster.engflow.com/invocation/`) is printed by Bazel; it is the canonical place to look at action logs, cache hits, and per-test stdout. - Profile data goes to `/tf/pkg/profile.json.gz` (this path is hardcoded in the AMD scripts; the container `--tmpfs` plus the `/tf/pkg` directory is where Bazel writes timing data for `bazel analyze-profile`). - `lscpu` and `rocminfo` output are at the top of each job log, which is often enough to triage "is the GPU actually visible to the container". ## Maintenance and updates Common edits and where to make them: | Goal | File | | ------------------------ | ------------------------------------------------- | | Bump ROCm version | `.github/workflows/rocm_ci.yml`, `rocm-config` | : : step — update the sha256 digest : | Add/remove a base tag | `build_tools/rocm/rocm_tag_filters.sh` | : filter : : | Re-tune the compile-only | `build_tools/ci/build.py`, | : ROCm phase : `XLA_LINUX_X86_GPU_ROCM_GITHUB_ACTIONS` block + : : : `rocm_tag_filter` tuple : | Add/remove tests from | `build_tools/rocm/rocm_xla.bazelrc`, | : single-GPU sweep : `test\:xla_sgpu` : | Promote a test to | `build_tools/rocm/rocm_xla.bazelrc`, | : multi-GPU : `test\:xla_mgpu`, and ensure it has the : : : `multi_gpu` tag : | Change EngFlow endpoint | `build_tools/rocm/rocm_xla.bazelrc`, | : or auth path : `build\:rocm_rbe` lines + the `volumes\:` in : : : `rocm_ci.yml` : | Tighten per-test | `build_tools/rocm/parallel_gpu_execute.sh` | : isolation : : | Adjust JAX-side build | The `jax` step in `rocm_ci.yml`, calling JAX's | : flags : `ci/run_bazel_test_rocm_rbe.sh` : | Adjust upstream driver | This lives in `ROCm/xla` on the `rocm-dev-infra` | : script : branch : : : (`build_tools/rocm/execute_ci_build_upstream.sh`) : : : — not in this repo : A change to `build.py` should be paired with `build_tools/ci/golden_commands.txt` regeneration (the readme explains how: `PYTHONDONTWRITEBYTECODE=1 python3 build.py --dump_commands > golden_commands.txt`). Goldens are documentation, not enforcement — CI does not diff against them — but reviewers do read them to see the materialised command lines. ## References - Workflows: `.github/workflows/rocm_ci.yml`, `.github/workflows/ci.yml` - ROCm-specific scripts and bazelrc: `build_tools/rocm/` - Shared CI driver: `build_tools/ci/build.py` - Tag-filter helper: `build_tools/rocm/rocm_tag_filters.sh` - GPU lock helpers: `build_tools/rocm/parallel_gpu_execute.sh`, `build_tools/ci/parallel_gpu_execute.sh` - Sanitizer ignore lists: `build_tools/rocm/{asan,lsan,tsan}_ignore_list.txt` - Upstream-fork driver script (external): `ROCm/xla` on branch `rocm-dev-infra`, file `build_tools/rocm/execute_ci_build_upstream.sh` - EngFlow cluster (RBE + BES): `grpcs://wardite.cluster.engflow.com` --- ### Shapes (`docs/shapes.md`) # Shapes and layout ## Structure of an XLA Op Consider an example HLO: ``` add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)} add(exponential.183, broadcast.3115) ``` This consists of the following components: * Op Name: `add.936` * This is the unique name for the operation. * Shape: `bf16[8,1,1280,16384]` * This is the output shape of the Op. Here the dtype is [bf16](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format) and the shape is `[8,1,1280,16384]`. * Layout (with Tiling): `3,2,0,1:T(8,128)(2,1)` * This describes how the array is stored in memory. `3,2,0,1` denotes the order of the axes in memory (i.e., column major, row major, etc.) and `T(8,128)(2,1)` denotes the tiling & padding used. * Layout is optional. If not specified, there is no tiling and the dimensions are assumed to be ordered from most-major to most-minor. * Operation: `add` * The operation being performed. Here, it is [Add](operation_semantics.md#add), which is also mentioned in the Op name. * Arguments: `exponential.183`, `broadcast.3115` * This operation takes two arguments, specified with their unique names. Let's consider another example, a fusion Op: ``` %fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)} fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32), kind=kCustom, calls=%all-reduce-scatter.3 ``` In addition to the previously described components, this consists of: * Attributes: `kind` and `calls` * These provide more information about the operation being performed, in this case: fusion. * Memory location (memory space identifier): `S(1)` * This denotes the memory space/location where the array is stored. `S(1)` here denotes this array lives in VMEM (on a TPU). * Shape and layout details for the input argument `%fusion.32` The following sections describe Shapes, [Layout](#layout), and [Memory Space Identifiers](#memory-space-identifiers). You can learn more about Tiling in [Tiled Layout](tiled_layout.md). ## Shapes The XLA `ShapeProto` proto ([xla_data.proto](https://github.com/openxla/xla/tree/main/xla/xla_data.proto)) describes the number of dimensions, size, and data type of an N-dimensional array (*array* in short). ### Terminology, notation, and conventions NOTE: in the past, XLA has used the term "rank" to mean the number of dimensions of an array. We have stopped this usage as it's inconsistent with the matrix rank concept in linear algebra. However, you may still see the name `rank` used in legacy documentation and some of the code. * The *true number of dimensions* of an array is the number of dimensions which have a size greater than 1. * Dimensions are numbered from `0` up to `N-1` for an `N` dimensional array. The size of a dimension is a non-negative integer. In particular, size 0 is valid. The dimension numbers are arbitrary labels for convenience. The order of these dimension numbers does not imply a particular minor/major ordering in the layout of the shape. The layout is determined by the `LayoutProto` proto. * By convention, dimensions are listed in increasing order of dimension number. For example, for a 3-dimensional array of size `[A x B x C]`, dimension 0 has size `A`, dimension 1 has size `B`, and dimension 2 has size `C`. Some utilities in XLA also support Python-like negative indexing: Dimension -1 is the last dimension (equivalent to `N-1` for an `N` dimensional array). For example, for the 3-dimensional array described above, dimension -1 has size `C`, dimension -2 has size `B`, and so on. * Two, three, and four dimensional arrays often have specific letters associated with dimensions. For example, for a 2D array: * dimension 0: `y` * dimension 1: `x` For a 3D array: * dimension 0: `z` * dimension 1: `y` * dimension 2: `x` For a 4D array: * dimension 0: `p` * dimension 1: `z` * dimension 2: `y` * dimension 3: `x` * Functions in the XLA API which take dimensions do so in increasing order of dimension number. This matches the ordering used when passing dimensions as an `initializer_list`; e.g. `ShapeUtil::MakeShape(F32, {A, B, C, D})` will create a shape whose dimension size array consists of the sequence `[A, B, C, D]`. ## Layout The `LayoutProto` proto describes how an array is represented in memory. It includes the following fields: ``` message LayoutProto { repeated int64 minor_to_major; int64 tail_padding_alignment_in_elements; ... } ``` ### Minor-to-major dimension ordering The only required field is `minor_to_major`. This field describes the minor-to-major ordering of the dimensions within a shape. Values in `minor_to_major` are an ordering of the dimensions of the array (`0` to `N-1` for an `N` dimensional array) with the first value being the most-minor dimension up to the last value which is the most-major dimension. The most-minor dimension is the dimension which changes most rapidly when stepping through the elements of the array laid out in linear memory. For example, consider the following 2D array of size `[2 x 3]`: ``` a b c d e f ``` Here dimension `0` is size 2, and dimension `1` is size 3. If the `minor_to_major` field in the layout is `[0, 1]` then dimension `0` is the most-minor dimension and dimension `1` is the most-major dimension. This corresponds to the following layout in linear memory: ``` a d b e c f ``` This minor-to-major dimension order of `0` up to `N-1` is akin to *column-major* (for 2-dimensionals). Assuming a monotonic ordering of dimensions, another way we may refer to this layout in the code is simply "dim 0 is minor". On the other hand, if the `minor_to_major` field in the layout is `[1, 0]` then the layout in linear memory is: ``` a b c d e f ``` A minor-to-major dimension order of `N-1` down to `0` for an `N` dimensional array is akin to *row-major* (for 2-dimensionals). Assuming a monotonic ordering of dimensions, another way we may refer to this layout in the code is simply "dim 0 is major". #### Default minor-to-major ordering The default layout for newly created Shapes is "dimension order is major-to-minor" (i.e. `[N-1, ..., 0]`). ### Padding The `tail_padding_alignment_in_elements` field defines the alignment of the [tiled](tiled_layout.md) array in terms of the number of elements. After applying tiling, padded elements will be added at the end of the layout until the total number of elements is a multiple of this value. ### Indexing into arrays The class `IndexUtil` in [index_util.h](https://github.com/openxla/xla/tree/main/xla/index_util.h) provides utilities for converting between multidimensional indices and linear indices given a shape and layout. Multidimensional indices include an `int64` index for each dimension. Linear indices are a single `int64` value which indexes into the buffer holding the array. See `shape_util.h` and `layout_util.h` in the same directory for utilities that simplify creation and manipulation of shapes and layouts. ## Memory Space Identifiers In HLO, each array may be annotated with a memory space identifier, written as S(n). * `S(0)` (often omitted) denotes device high bandwidth memory (HBM). * `S(1)` represents on device virtual memory (VMEM). * `S(2)`, `S(3)`, etc., correspond to additional device specific memory spaces. * `S(5)` indicates host memory. --- ### Sparsecore (`docs/sparsecore.md`) # A deep dive into SparseCore for Large Embedding Models (LEM) SparseCore is a specialized tiled processor engineered for high-performance acceleration of workloads that involve irregular, sparse memory access and computation, particularly on large datasets stored in High Bandwidth Memory (HBM). While it excels at tasks like embedding lookups, its capabilities extend to accelerating a variety of other dynamic and sparse workloads. ## 1. Introduction to SparseCore Key architectural features: * **Tiled architecture**: Comprises multiple compute tiles (each tile is a complete dataflow unit with its own local memory and processing unit) allowing for parallel processing. * **Dynamic execution**: Natively supports data-dependent control flow and memory accesses, crucial for sparse data. * **Vector processing**: Utilizes small-vector tasks (8-element or 16-element, depending on the hardware version) for efficient computation. * **Centralized control**: A single SparseCore sequencer orchestrates tasks across all tiles, ensuring synchronized operations. * **Data summarization support**: Includes specialized cross-lane operations beneficial for tasks such as sorting, filtering, and prefix sums. * **Memory hierarchy**: Strategically leverages HBM for storing large datasets and local scratchpad memory (SPMEM) for staging frequently accessed data, significantly reducing HBM latency. ### Specifications at a glance: | Attribute | TPU v4 | TPU v5p | Trillium | | :----------------- | :------ | :------- | :------------------ | | **SparseCores/Chip** | 4 | 4 | 2 | | **Tiles/SparseCore** | 16 | 16 | 16 | | **SIMD Width** | 8 | 8 | 8 (F32) 16 (BF16) | | **HBM Capacity** | 32 GiB | 96 GiB | 32 GiB | ## 2. SparseCore host preprocessing Effective data preparation is paramount for SparseCore performance, and this is where host preprocessing plays a vital role. It encompasses several key functionalities: * **Data transformation**: * Apply necessary transformations to raw input data. * Manage ID transformations, which is particularly important when dealing with feature or table stacking. * Convert input data into the Coordinate (COO) sparse format, detailed in the following section. * Partition the data for efficient distribution across the different SparseCores available on the chip. * **Limit validation**: * Ensure that the characteristics of the input data (for example, number of IDs) conform to the predefined operational limits of the SparseCore, such as `max_ids_per_partition` and `max_unique_ids_per_partition`. * If the input data exceeds these limits, your host preprocessing layer can attempt to segment the data into smaller mini-batches that do fit within the constraints. * **Data transfer**: * Efficiently copy the processed and validated data to the TPU's High Bandwidth Memory (HBM), making it ready for SparseCore execution. ### Understanding table stacking: Table stacking is a significant optimization technique where multiple embedding tables are logically combined to enhance embedding lookup efficiency. This process is typically handled automatically by the underlying ML framework. * **Feature stacking**: This occurs when multiple distinct features share the same underlying embedding table. A common example is using a single embedding dictionary for various categorical features like zip codes from different contexts. * **Table stacking**: In this scenario, multiple distinct embedding tables are stacked together. Tables that share the same embedding dimension and optimizer configuration are often grouped. The primary advantage of table stacking is the creation of a larger effective batch size for the operations on these stacked tables. This reduces computational overhead and can be effective in hiding inter-chip communication (ICI) latencies. For optimal performance, a moderate number of stacked tables (generally in the range of 5 to 100) is recommended. ## 3. Conversion to COO tensors Before data can be processed by SparseCore, it's commonly converted into a Coordinate (COO) sparse tensor format. The COO format is a way to represent sparse matrices efficiently typically using three arrays: * `row_ids`: An array containing the row indices for each non-zero element. In the context of batch processing, this often corresponds to the batch dimension. * `col_ids`: An array containing the column indices for each non-zero element. For embeddings, these are often the feature or ID values. * `values` (optional): An array holding the actual values of the non-zero elements at the corresponding (`row`, `col`) coordinates. For limit calculations (discussed later) related to ID counts, these values (gains) are often not considered. ### Illustrative example: Consider an input sparse matrix representing batches of IDs: [ [id_A], // Sample 0 [id_A, id_B, id_C], // Sample 1 [id_B, id_B, id_D], // Sample 2 (note duplicate id_B) ] After conversion to COO format (and potentially after deduplicating IDs within the same sample): row_ids = [0, 1, 1, 1, 2, 2] col_ids = [id_A, id_A, id_B, id_C, id_B, id_D] This conversion is fundamental to how SparseCore processes and distributes work. The `col_ids` in particular, are crucial for determining which specific SparseCore partition an ID belongs to, enabling efficient sharding and lookup. ## 4. SparsecoreConfig: The high-level API ### Framework specific embedding APIs: * **JAX**: https://github.com/jax-ml/jax-tpu-embedding * **TensorFlow**: https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding * **Keras**: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding The `SparsecoreConfig`, or equivalent mechanisms such as XLA flags, serves as a high-level interface for controlling a wide array of SparseCore behaviors. A thorough understanding of these parameters is vital for effective performance tuning and ensuring the correct operation of your models. * **`disable_table_stacking: bool = False`** * **Explanation**: This flag controls whether automatic table stacking is prevent the framework from stacking tables, potentially leading to reduced performance due to increased overheads and a diminished ability to hide Inter-Chip Interconnect (ICI) latency. * **Default**: `False` (implying table stacking is generally enabled by default where the framework supports it). * **`max_ids_per_chip_per_sample: int = 64`** * **Explanation**: This parameter establishes a global upper limit on the total number of embedding IDs that a single chip can process from one sample in the input batch, aggregated across all tables. It is a mechanism for managing resources at the chip level, before more granular per-table or per-partition limits are taken into account. Fine-tuning this value typically depends on specific model characteristics and overall system capacity. * **Default**: `64`. * **`max_ids_per_table: Optional[Dict[str, int]] = None`** * **Explanation**: This parameter specifies the maximum number of embedding IDs (which can include duplicates) that can be processed for each logical table, considering all its partitions across all SparseCores. This is a broader limit than `max_ids_per_partition`. If a table `T` is divided into `P` partitions, this limit applies to the sum of IDs directed to all P partitions. It's often related to `max_ids_per_partition_per_sample` and the overall batch size. * **Setting**: Typically configured using a limits file (for example, using the `xla_sparse_core_max_ids_file` flag), where `max_ids_per_partition` is defined. This table-level concept is a method to set those partition-level limits (`max_ids` and `max_uniques`). * **Default**: `None` (the value may be inferred from per-partition limits or other configurations if not explicitly provided). * **`max_unique_ids_per_table: Optional[Dict[str, int]] = None`** * **Explanation**: Analogous to `max_ids_per_table`, but this parameter specifies the maximum number of unique IDs for each logical table. This is a critical setting for appropriately sizing on-device buffers used in unique ID processing and subsequent vector operations. * **Setting**: Also commonly defined in a limits file or derived from `max_unique_ids_per_partition_per_sample`. * **Default**: `None`. * **`allow_id_dropping: bool = False`** * **Explanation**: This boolean flag controls ID dropping when the number of IDs encountered in the input data (observed limits) surpasses the limits set during compilation (for example, `max_ids_per_partition`). * **If `True`**: IDs that would cause the limits to be exceeded are silently dropped. Typically, IDs within a partition are processed in a sorted order, and any ID that would push the running count over the limit for its designated mini-batch is discarded. This allows the program to continue execution but may have an adverse impact on model accuracy. * **If `False`**: An error is triggered, and the process will likely terminate if observed limits go beyond compiled limits. This approach ensures all data is processed but requires limits to be configured more conservatively. * **Default**: `False` (causing an error on overflow rather than silent data dropping). * **`initialize_tables_on_host: bool = True`** * **Explanation**: This flag determines whether embedding tables are initialized on the host CPU before being subsequently transferred to the TPU's High Bandwidth Memory (HBM). Standard practice is for tables to be initialized on the host. Setting this to `True` follows this convention. If it were set to `False`, it would imply an on-device initialization mechanism, which could have different performance implications or specific initialization prerequisites. * **`enable_fast_table_initialization: bool = False`** * **Explanation**: Initializes the tables directly on the TPU. This can help reduce model startup times. ## 5. Pipelining for performance Pipelining is a performance optimization technique that enables the simultaneous execution of operations on the TensorCore (TC) and the SparseCore (SC). By overlapping these computations, overall throughput can be significantly improved. * **Mechanism**: In a standard training step that involves sparse embedding lookups (handled by SC) and dense layer computations (handled by TC), pipelining allows the SC to work on its part of step `i` (for example, forward or backward pass) while the TC is concurrently processing a different part of the same step `i`, or even parts of adjacent steps like `i-1` or `i+1`. * **Impact on gradients**: The SparseCore might operate on "stale" gradients. For instance, the gradients calculated during the backpropagation phase of step `i` might not be fully updated and visible to the SC until step `i+2`. * **Performance vs. numerics trade-off**: This overlapping execution can lead to substantial speedups, potentially up to a 2x improvement in device step time. However, the subtle changes in numerics (embedding_weights) resulting from the use of stale gradients might influence model convergence behavior or the final achieved accuracy. The acceptability of this trade-off is highly model-dependent and often requires empirical validation. * **Control flag**: Pipelining can be controlled by `tf_xla_disable_full_embedding_pipelining`. Setting this flag to `true` disables full pipelining (overlapping TensorCore and SparseCore computation), whereas setting it to `false` (or if the flag's semantics imply enabling when false) activates it. ### Conceptual pipelining flow: * **Without pipelining (simplified sequential flow)**: `Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i` * **With pipelining (simplified overlapped flow)**: ``` Time -> Step i: SC/F_i | TC/F_i | TC/B_i | SC/B_i Step i+1: SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1 ``` **Note**: The actual pipelining stages implemented in the hardware and compiler can be more intricate, often involving pre-loops, main execution loops, and post-loops to manage data dependencies and ensure correctness. ## 6. The role of XLA XLA (Accelerated Linear Algebra) is the domain-specific compiler that translates high-level computational graphs, typically from frameworks like TensorFlow, into highly optimized machine code tailored for TPUs. This includes generating the instructions for operations destined for the SparseCore. ### Key functions in the SparseCore context: * **Compilation of sparse operations**: XLA is responsible for compiling embedding lookup operations (such as the `SparseDenseMatmulOp`) and other sparse computations into low-level, executable SparseCore programs. * **Integration of limits**: It utilizes the configured operational limits (for example, `max_ids_per_partition`, `max_unique_ids_per_partition`, often provided via a limits file specified by flags like `xla_sparse_core_max_ids_file`) to statically determine the sizes of and allocate on-device memory buffers, particularly within the SPMEM. * **Targeted optimizations**: XLA performs a suite of optimizations specifically designed for the SparseCore architecture. These can include instruction scheduling, memory layout transformations, and fusion of operations to maximize efficiency. * **Control using flags**: Many aspects of SparseCore behavior, tuning parameters, and optimization strategies are exposed and controlled through XLA flags (for example, `xla_sparse_core_estimate_max_ids` for limit estimation, or `xla_sc_detect_nan` for debugging). ### Open source status: Currently Sparsecore Implementation is internal and served using `libtpu.so`. ### Error reporting and diagnostics: Compilation failures related to SparseCore configurations or resource constraints often manifest as `XLA:TPU` compile-time errors. These error messages can provide valuable insights into issues such as limits being set too high for the available SPMEM, or the use of unsupported configurations. ## 7. How limits translate to tables on SparseCore On SparseCore, "limits" are fundamental configuration parameters that primarily refer to two per-partition settings for each table that is sharded (distributed) across the available SparseCores: * **`max_ids_per_partition`**: This defines the maximum number of total IDs (including duplicates) that any single SparseCore is expected to send to, or process for, a specific partition of a given table within a single computational step. * **`max_unique_ids_per_partition`**: This defines the maximum number of unique IDs that any single SparseCore is expected to send to, or process for, a ### Translation to physical table layout and processing: * **Table sharding strategy**: Embedding tables are typically "mod-sharded" across all SparseCores in the system. This means each SparseCore becomes responsible for a distinct subset of the vocabulary (rows) of each table. An ID `j` would generally be assigned to `SparseCore_k` based on a formula like `k = j % num_total_sparse_cores`. * **Definition of a "partition"**: In this context, a "partition" refers to the specific segment of an embedding table for which a single SparseCore handles lookups. * **SPMEM buffer allocation**: These limits are used by the XLA compiler to statically size and allocate buffers within the on-device scratchpad memory (SPMEM). Buffers are dimensioned such that all necessary data related to the IDs for a given partition (up to the specified `max_ids` and `max_unique_ids` limits) can be loaded into SPMEM for processing. This is particularly crucial for non-elementwise computations, such as reducing duplicate IDs within a partition (for example, , when creating a Compressed Sparse Row (CSR) representation), where the entire relevant dataset for that partition's IDs needs to be readily available in fast memory. * **Compiled limits versus observed limits**: * **Observed limits**: These are the actual number of IDs encountered for each partition during runtime, based on the input data being processed. * If observed limits exceed compiled limits, it can lead to ID dropping (if `allow_id_dropping` is enabled) or errors. * **Calculating limits**: The process of determining appropriate limits involves a careful analysis of the input data distribution. For any given table (let's call it `T1`, which might itself be part of a larger stacked table `T`): 1. The input batch (for example, a 2D `SparseTensor` of shape `[BatchSize, MaxSequenceLength]`) is initially split across the available SparseCores. For instance, if a TensorCore is paired with 2 SparseCores, each SparseCore might receive a sub-batch of shape `[BatchSize/2, MaxSequenceLength]`. 2. This sub-batch is then converted into the COO format, yielding `row_ids` and `col_ids`. 3. Duplicate IDs within the same sample (i.e., entries with the same `row_id` and `col_id`) are removed. 4. For each remaining unique `col_id` (within a sample), the target SparseCore responsible for this ID is determined using the mod-sharding rule: `target_sc_id = col_id % num_total_sparse_cores`. 5. A count is maintained of the total number of IDs (`ids_per_sparse_core[target_sc_id]++`) and the number of unique IDs (`unique_ids_per_sparse_core[target_sc_id]++`, after ensuring uniqueness for that specific `target_sc_id`) that are destined for each `target_sc_id`. 6. The `max_ids_per_partition` for table `T1` is then set to `max(ids_per_sparse_core_array)`. 7. Similarly, the `max_unique_ids_per_partition` for table `T1` is set to `max(unique_ids_per_sparse_core_array)`. 8. If table `T1` is a component of a stacked table, additional transformations like rotations or shifts might be applied to the ID distributions before summing statistics from all constituent tables. This helps in balancing the load across chips. Setting these limits correctly is a balancing act: lower limits can potentially lead to higher performance (as less data needs to be processed per step and SPMEM pressure is reduced), but if set too low, they can result in excessive mini-batching or undesirable ID dropping. ## 8. How each SparseCore communicates SparseCore communication, particularly in the context of processing a list of IDs for embedding lookups, relies on several coordinated mechanisms: * **Mod sharding and implicit routing**: * Embedding tables are mod-sharded across all SparseCores in the system. * When the host provides a batch of input data (which is subsequently preprocessed into COO format, including `col_ids`), the `col_id` value is used to determine which SparseCore is responsible for that specific ID: `target_sc_id = col_id % num_total_sparse_cores`. * Each SparseCore effectively receives and processes only the subset of IDs that map to its assigned vocabulary partitions. The host preprocessing stage is crucial for preparing the data in such a way that each SparseCore can readily identify and operate on its relevant IDs. * **Data distribution by host**: * The host preprocessing logic partitions the overall input batch and distributes the pertinent portions of the `row_ids` and `col_ids` (along with any associated features or weights if applicable) either to memory (HBM) directly accessible by each SparseCore or to a shared HBM from which the SparseCores will fetch their required data. * **Intra-SparseCore processing**: * Once a SparseCore has received its designated set of IDs for a given table partition, it performs operations such as deduplication of these IDs and gathering the corresponding embedding vectors. These are primarily local computations executed within the SparseCore's own tiles and utilizing its local SPMEM. * **Inter-SparseCore communication (All-to-All)**: * After the initial processing phase (like embedding lookups), an "all-to-all" communication pattern may be used to combine or redistribute results across SparseCores (for example, before feeding activations into a TensorCore layer that expects input corresponding to all original sample positions). This is vital for reconstructing the complete set of activations if the original input batch was distributed for parallel processing. * **Communication with TensorCores**: * SparseCores communicate with TensorCores to send embedding activations (during the forward pass) and to receive gradients (during the backward pass). This interaction is orchestrated by the XLA-compiled program and frequently involves HBM as an intermediary buffer. The pipelining strategy (discussed earlier) heavily influences the timing and synchronization of this SC-TC communication. In essence, the initial "distribution" of IDs to the appropriate SparseCores is largely handled by the sharding scheme and the host preprocessing steps. Subsequent communication involves SparseCores operating on their local data, potentially followed by collective communication operations like all-to-all if data needs to be globally exchanged or reordered across SparseCores before further processing by the TensorCores. ## 9. SparseCore memory management Each SparseCore efficiently manages several distinct types of memory to perform its computations: * **Scratchpad memory (SPMEM)**: * **Nature**: A relatively small, but very fast, local SRAM that is available exclusively to each SparseCore. It is important to note that SPMEM is not a cache; its usage is explicitly managed and orchestrated by the XLA compiler. * **Purpose**: SPMEM is used to "opportunistically stage data." This includes inputs, outputs, and intermediate results that are required for ongoing SC computations. Staging data in SPMEM significantly reduces the high latency typically associated with accessing HBM. * **Sizing**: As discussed in the section on "Limits," SPMEM buffers are statically sized at compile time. This sizing is based on parameters like `max_ids_per_partition` and `max_unique_ids_per_partition`. This static allocation ensures that for any given operation on a table partition (such as CSR reduction), all the necessary data for that partition's IDs (up to the defined limits) can fit into SPMEM. * **Compiler optimizations**: The XLA compiler incorporates sophisticated optimizations to determine precisely how much data, and which specific data elements, need to be staged in SPMEM to effectively hide HBM latency and maximize performance. * **Dynamic allocation constraint**: The SparseCore compiler does not currently support dynamic scratchpad allocation. This highlights the critical importance of static sizing through the careful configuration of limits. * **High bandwidth memory (HBM)**: * **Nature**: A large, shared memory resource accessible by all SparseCores, TensorCores, and the host system. Primary embedding tables are stored in HBM. * **Stack usage**: SparseCore operations often require temporary storage in HBM for intermediate results that either do not fit into the limited SPMEM or need to be passed between larger stages of the processing pipeline. HBM stack usage during both forward and backward passes can be estimated as follows - * Forward Pass HBM Stack (single table) ≈ (2 \* `feature_width` + 1) \* `max_unique_nz_per_row` \* `logical_replica_count` \* 4 bytes * Backward Pass HBM Stack (single table) ≈ 3 \* `feature_width` \* `max_unique_nz_per_row` \* `logical_replica_count` \* 4 bytes * **Heap usage**: HBM also accommodates the heap, which is managed by the host. The heap stores data such as dense layer weights, constants used by the model, and prefetched input data. Heap usage tends to increase with the number of steps for which the host prefetches data (controlled by the `maximum_parallel_iterations` flag). While more prefetching can improve performance by overlapping host-to-device transfers with device computation, it also consumes more HBM. * **Serialization for HBM optimization**: The flag `xla_sc_num_serialized_tables_to_optimize_hbm` provides a mechanism to control how many tables' data is kept "live" in HBM stack memory at any given time. Increasing this number effectively serializes the processing for more tables, which can reduce peak HBM stack usage but may come at the cost of performance due to reduced parallelism. * **Vector memory (VMEM)**: * VMEM is a local scratchpad memory used exclusively by the TC (TensorCore). While VMEM is not directly managed by the SparseCore, it is an integral part of the memory ecosystem with which the SC interacts, primarily through the TensorCore. ### Overall memory management strategy: The core memory management strategy for SparseCore revolves around using the small, fast SPMEM for the "hot" data that is actively being processed by a SparseCore tile, thereby minimizing accesses to the slower HBM. The configured limits are the primary mechanism for ensuring that SPMEM does not overflow. HBM is utilized for storing large embedding tables and temporary data that either exceeds SPMEM capacity or needs to be shared across different processing units or pipeline stages. The XLA compiler is responsible for orchestrating all data movement and buffer allocation based on these architectural principles and the user-configured limits. ## 10. Performance & memory bottlenecks Achieving optimal performance with SparseCore necessitates a clear understanding of potential bottlenecks and how to address them. These can arise on the host, within the SparseCore itself, or in its interaction with the TensorCores. ### Common performance bottlenecks: * **Host bottleneck**: * **Issue**: The host CPU may fail to preprocess data and feed it to the TPU rapidly enough, leading to underutilization of SparseCores and TensorCores. This is a frequent performance limiter. * **Mitigation**: Monitor host CPU utilization and input pipeline metrics. Optimize host-side data loading and preprocessing routines (refer to COO conversion tips). Adjust the `maximum_parallel_iterations` flag to fine-tune data prefetching. * **Suboptimal TC/SC synchronization (lack of pipelining)**: * **Issue**: If pipelining between the TensorCore and SparseCore is disabled or not operating efficiently, one unit may spend significant time waiting for the other, thereby reducing overall system throughput. * **Mitigation**: Ensure pipelining is enabled (for example, `tf_xla_disable_full_embedding_pipelining = false` or its equivalent). * **Limit-induced bottlenecks**: * **Issue**: * **Limits too low**: Can trigger excessive mini-batching (splitting input batches into numerous smaller sub-batches to meet the tight limits). While this maintains correctness, each mini-batch introduces some processing overhead, potentially slowing down the overall execution. If `allow_id_dropping` is true, overly low limits can also lead to ID dropping, which impacts model accuracy. * **Limits too high (but still fit)**: While very high limits might prevent mini-batching, they could increase SPMEM pressure unnecessarily if the actual data characteristics rarely approach these peak values. They could also lead to larger HBM stack usage than strictly needed. * **Compilation failures**: If the configured limits require more SPMEM or HBM stack than the available physical memory, the compilation will fail. * **Mitigation**: Ensure limits are set correctly. * **Data distribution skew**: * **Issue**: If certain SparseCore partitions consistently receive a disproportionately larger number of IDs compared to others (indicating a poor ID distribution), those overloaded SparseCores will become performance bottlenecks. * **Mitigation**: ID shuffling during the mini-batching process can help alleviate this for stacked tables, especially those with "hot" user tables. Analyze ID distributions carefully to set appropriate and balanced per-table limits. * **Table stacking issues**: * **Issue**: * **Too few tables stacked**: May not be sufficient to effectively hide ICI latency or adequately reduce processing overheads. * **Too many tables stacked**: Could result in the creation of very large logical tables that become unwieldy to manage or might exceed available resource limits. * **Mitigation**: * Ensure optimal number of tables for stacking. A general guideline suggests a "sweet spot" of 5-100 tables for stacking. * **Inefficient numerics/quantization**: * **Issue**: Using full FP32 precision when lower precision formats like BF16 or quantized integers would suffice (and offer faster computation) can be a performance bottleneck. * **Mitigation**: Explore lower precision options. However, be aware that quantization itself has some overhead and might require careful tuning of quantization parameters to maintain model accuracy. * **HBM bandwidth saturation**: * **Issue**: Excessive data movement to and from HBM, potentially caused by very small feature widths (leading to high padding overhead), inefficient memory access patterns, or an extremely large number of lookups, can saturate the available HBM bandwidth. * **Mitigation**: Scaling the number of TPUs can help with HBM bandwidth saturation. ### Common memory bottlenecks: * **SPMEM overflow (compilation failure)**: * **Issue**: If `max_ids_per_partition` and `max_unique_ids_per_partition` are set too high, the XLA compiler may be unable to allocate sufficient SPMEM, resulting in compilation errors like: `"Fixed size allocations (...) do not fit in TileSpmem (...)"`. Additionally, if the term `(sample_count * feature_width) / kNumTiles` (where `kNumTiles` is the number of tiles per SC) is too large for staging gather operands within the tile SPMEM, errors such as `"Gather operand too large..."` can occur. * **Mitigation**: Reduce the batch size or increase the number of chips used for processing. * **HBM stack overflow (runtime or compilation)**: * **Issue**: If the combination of `feature_width`, `max_unique_nz_per_row`, and `logical_replica_count` leads to HBM stack memory requirements that exceed the available HBM, this can cause Out-Of-Memory (OOM) errors either at runtime or during compilation. * **Mitigation**: Tune the `xla_sc_num_serialized_tables_to_optimize_hbm` flag to reduce HBM stack usage by serializing the processing of tables (this usually comes at a performance cost). * **HBM heap exhaustion**: * **Issue**: Primarily caused by very large dense layer weights, numerous constants stored in memory, or overly aggressive input prefetching (high `maximum_parallel_iterations`). * **Mitigation**: Monitor heap usage using tools like XProf Memory Viewer. * **Padding overhead**: * **Issue**: Embedding tables are padded to be 32B-aligned (equivalent to 8 floats) in the feature dimension. Consequently, small feature widths (for example, 1 float) incur significant padding overhead (for example, 7/8 of the allocated buffer space is padding), leading to wasted HBM. The vocabulary dimension of tables is also padded to be a multiple of the number of SparseCores in the system; however, this impact is usually negligible for tables with a sufficiently high vocabulary size. ### General factors impacting performance and memory: * **Topology**: The number of chips available and their interconnection architecture. * **Batch sized**: Directly affects the `sample_count` per SparseCore, which in turn influences memory consumption and compute load. * **Data formatting**: Ensuring an efficient on-device data layout is crucial for optimal performance. ## 11. Analyzing a SparseCore profile Analyzing a performance profile is a key step in identifying bottlenecks and uncovering opportunities for optimization within your SparseCore workloads. 1. **Obtain a trace**: * Utilize profiling tools, such as [XProf](https://github.com/openxla/xprof), to capture a detailed execution trace while your model is training or running inference. This trace will provide a timeline of operations occurring on the host, the TensorCores, and the SparseCores. 2. **Examine the Trace Viewer (for example, in XProf or TensorBoard)**: * **Host activity**: Scrutinize the host's activity. Are there significant gaps in TPU activity? Such gaps might indicate that the host is a bottleneck, failing to feed data quickly enough. Analyze the performance of your input pipeline. * **TensorCore (TC) and SparseCore (SC) activity**: * Look at the execution timelines for both TC and SC. Are they operating in parallel, indicating effective pipelining? Or are there extended periods where one unit is idle, waiting for the other? * Identify the operations that consume the most time (longest-running operations) on both the SC and TC. * Visual trace outputs (often showing colored blocks representing different operations over time, like the `TPU:0 SparseCore 1 (pid 1005)`) are invaluable for visually identifying dominant operations and idle periods. * **Step time analysis**: Observe the overall step time and understand how it is distributed or broken down between host processing, SC computation, and TC computation. 3. **Memory analysis (XProf Memory Viewer)**: * **Heap usage**: Use tools like XProf's "Memory Viewer" tab to inspect HBM heap usage. This can help determine if large model weights, constants, or overly aggressive input prefetching are consuming an excessive amount of HBM. Enabling flags like `--vmodule=best_fit_allocator=1` might provide logs of peak heap usage. * **Stack usage (indirect)**: While direct HBM stack profiling can be complex, if you encounter Out-Of-Memory errors and heap usage appears reasonable, HBM stack exhaustion (often due to overly large limits or feature widths) is a strong suspect. The formulas provided for HBM stack usage can help in estimating this. 4. **Look for specific patterns**: * **Mini-batching**: If limits are frequently being exceeded, you might observe evidence of mini-batching in the trace (for example, a higher number of smaller SC operations than expected for the global batch size). This can often be inferred from logs or by observing the invocation counts of certain operations. * **ID dropping**: If ID dropping is enabled and occurring, system logs might provide indications of this. This would also be a clear sign that the configured limits are too restrictive for the input data. * **Compilation times**: Extended recompilation times, particularly if Feedback Directed Optimization (FDO) is enabled and frequently adjusting limits, can add significant overhead to the overall training time. 5. **Correlate with flags and configuration**: * Relate the observed behavior in the profile back to your SparseCore configurations (settings in limits files, XLA flags). For example, if `xla_sc_num_serialized_tables_to_optimize_hbm` is set to a high value, you might expect slower SC performance but lower HBM stack consumption. 6. **Iterative process**: * Profiling is often an iterative refinement process. Make a specific change (adjust a limit, enable or disable a feature), capture a new profile, and then compare it against the previous profile to see the impact of your modification. ## 12. General debugging flags Several flags can be enabled to assist in debugging issues related to SparseCore execution. It's important to note that enabling these checks often incurs a performance penalty and, therefore, they should typically be disabled for production runs. * **ID checks (out-of-range)**: * **Flag**: `xla_sparse_core_enable_id_bound_check = true` * **Purpose**: Enables checks on the host system to detect if any embedding IDs in the input data fall outside the valid vocabulary range defined for a given embedding table. This helps catch issues related to incorrect or corrupted input data. * **NaN checker**: * **Flag**: `xla_sc_detect_nan = true` * **Purpose**: Enables the detection of NaN (Not a Number) values within floating-point data processed on the SparseCore. If a NaN is detected in the inputs or outputs of various compiler passes, this flag will cause an error to be raised. Such errors typically provide information about where the NaN was encountered. * **Bounds checker (memory access)**: * **Flag**: `xla_sc_assert_level=bounds` * **Purpose**: This flag enables an ASAN (AddressSanitizer)-style tool that rewrites memory-accessing instructions (such as VMEM loads/stores and DMA operations) to include dynamic checks. These checks verify if the memory access is within the allocated bounds of the target memory region. * **Behavior**: If an out-of-bounds memory access is detected, the execution will fail. * **Caution**: It's possible for this checker to produce false positives, for example, due to complex strided access patterns that are not fully comprehended by the checker. This transformation is applied at a late stage in the backend compilation process. * **Buffer checker (memory corruption)**: * **Flags**: * `xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'` * `xla_tpu_verify_launch_id_across_cores=true` * **Purpose**: These flags help ensure that memory buffers are not being inadvertently corrupted or overwritten by unrelated operations. The Buffer Sanitizer checks the contents of buffers to verify that they are not changing unexpectedly. ## 13. Quantization support SparseCore's `SparseDenseMatmulOp` is designed to support operations on embedding tables using both 32-bit floating-point (FP32) and integer data types. While model training is typically performed using FP32 precision for embedding tables, post-training quantization (PTQ) can be applied. PTQ allows the use of lower-precision datatypes (like 8-bit integers) for inference, which can potentially lead to improved performance and a reduced memory footprint. ### Simulated Quantization: The `SparseDenseMatmulOp` can be configured to perform "simulated quantization." In this operational mode, embedding vectors are first quantized to a lower precision and then dequantized back to a higher precision (for example, FP32) before they are used in subsequent computations. This technique allows models to be trained while accounting for the effects of quantization noise. Training with simulated quantization can improve the accuracy of the final model when it is fully quantized for inference. ### Configuration attributes for `SparseDenseMatmulOp` (for quantization): * **`quantization_config_num_buckets = 256`** * This attribute specifies the number of discrete buckets or levels into which a 32-bit floating-point number will be quantized. For example, when quantizing to 8-bit integers, one would typically specify 2^8 =256 buckets. * **`quantization_config_low = -X.X`** * This attribute defines the minimum floating-point value in the quantization range. Any input values below this specified minimum will be clipped to this minimum value during quantization. * **`quantization_config_high = Y.Y`** * This attribute defines the maximum floating-point value in the quantization range. Any input values above this specified maximum will be clipped to this maximum value during quantization. ### Numerics and pipelining interaction: The numerical behavior of the model can change depending on whether pipelining between the TensorCore and SparseCore is enabled. If pipelining is active, gradients processed by the SparseCore might be "stale" (from a previous iteration). This can interact with the quantization process and potentially affect model training dynamics or final accuracy. ## 14. Upcoming features and recent improvements The SparseCore ecosystem is subject to continuous development and enhancement. ### Roadmap: * **Sample-dimension mini-batching**: * This is planned as a complementary feature to the existing vocabulary-dimension mini-batching capabilities. * It would allow for further partitioning of embedding inputs along the sample dimension. This would be achieved by introducing on-device loops that can filter and process lookups from a subset of samples at a time. Such a feature could be beneficial for managing very large per-sample ID counts or for improving load balancing across processing units. * **Improved support for embedding with less than 8 integers per row (small feature widths)**: * The current design often uses significant padding for embedding feature widths that are less than 8 floats (which corresponds to 32 bytes). This padding can lead to wasted HBM and potentially underutilized compute resources. Future improvements aim to mitigate this inefficiency for tables with small feature dimensions. ### Recent improvements: * **Staging of gather operands in HBM**: * This optimization helps to reduce pressure on the shared scratchpad memory (SPMEM) by allowing some gather operation inputs or outputs to be staged in the larger HBM. * **Reduction in stack memory usage**: * Enhancements have been implemented to reduce HBM stack memory consumption during SparseCore operations, ideally without negatively impacting overall performance or throughput. These enhancements are focused on improving SparseCore's performance, memory efficiency, and operational flexibility for an even wider range of sparse workloads. --- ### Symbolic Expression (`docs/symbolic_expression.md`) # Symbolic expressions and maps A symbolic expression (`SymbolicExpr`) is a mathematical abstraction system that enables symbolic tensor computations. A symbolic map (`SymbolicMap`) is a collection of symbolic expressions, that mathematically represents tensor mappings and transformations in the compilation pipeline. They act as the "mathematical bridge" between a high-level HLO operation and the actual memory addresses accessed by the GPU/CPU. `SymbolicExpr` and `SymbolicMap` are XLA's custom implementations that supersede the legacy `mlir::AffineExpr` and `mlir::AffineMap`. ## `SymbolicExpr` A `SymbolicExpr` represents a node in an abstract syntax tree (AST). Unlike the standard MLIR Affine Expression, it supports a wider range of operations necessary for modern GPU tiling. Example of a `SymbolicExpr`: `d0 + s0 * 8` ### Supported types **Constants:** Fixed integer (`int64`) values. **Variables:** Symbolic variables with *dimensions (d)* and *symbols (s)*. All variables (dimensions and symbols) are treated as `VariableID`, and the interpretation of a variable depends on its context within a Symbolic Map. **Operations:** `add`, `mul` `mod`, `floorDiv`, `ceilDiv`, `min`, `max`. *Supported operators: `+`, `-`, `*`, `/` (`floorDiv`), `%` (`mod`)* ### Example usage ```sh v0 = CreateSymbolicVariable(0, context); // 0 is the var_id v1 = CreateSymbolicVariable(1, context); SymbolicExpr expr = (((v0 + 42) * v1.min(2).max(0)) / 2).ceilDiv(2); int64_t result = expr.Evaluate({5, 1}); // Result: 12 ``` ### Key features * **Immutability**: Symbolic expressions are pointer-like handles to internal storage managed by [`mlir::MLIRContext`](https://mlir.llvm.org/doxygen/classmlir_1_1MLIRContext.html). They are automatically deduplicated to ensure uniqueness. * **Canonicalization (`Canonicalize()`)**: Symbolic expressions are algebraically simplified (with constant folding, identity elimination, associative property, distributive property, etc.) to ensure expressions are represented in a standard, minimal form. For example, `(d0 + 1) - 1` will be simplified to `d0`. This is vital for predictable equality checks (`operator==`) in expressions that are mathematically equivalent but structurally different. ## `SymbolicMap` `SymbolicMap` represents a mathematical mapping of transformation between coordinate systems, typically between input and output tensors. Example of a `SymbolicMap`: `(d0, d1)[s0, s1] -> (d0 + s0, d1 * s1)` ### Example usage ```sh SymbolicMap map = SymbolicMap::Get( context, 2, // number of dimensions 1, // number of symbols {d0 + s0, d1 * s1}); // SymbolicExprs // Access components int64_t num_dims = map.GetNumDims(); int64_t num_symbols = map.GetNumSymbols(); auto results = map.GetResults(); ``` ### Key operations **Variable substitution (`ReplaceDimsAndSymbols()`)**: Map dimensions and/or symbols can be substituted with other expressions, this enables re-mapping coordinate spaces. For example: ```sh // c2 and c3: SymbolicConstants // sample_map: (d0, d1)[s0, s1] -> (d0 + s0, d1 * s1) sample_map.ReplaceDimsAndSymbols( {d1, c2}, // New dimensions: Replace d0 with d1, d1 with c2 {c3, d0}, // New symbols: Replace s0 with c3, s1 with d0 2, // New number of dimensions 2) // New number of symbols // Result: (d1 + c3, c2 * d0) ``` **Composition (`Compose()`)**: Maps with compatible dimensions can be composed. This enables chaining transformations, and forms the core mechanism for fusing multiple HLO operations into a single indexing calculation. For example: ```sh map1: (d0, d1)[s0] -> (d0 + s0, d1 * 2) map2: (d0)[s0] -> (d0 - 10, d0 + s0) map1.compose(map2): (new_d0)[new_s0_map1, new_s0_map2] -> ((new_d0 - 10) + new_s0_map1, (new_d0 + new_s0_map2) * 2) ``` **Optimization (`CompressDims()` / `CompressSymbols()`)**: Simplifies maps by identifying and removing unused variables, keeping the generated indexing logic as lean as possible. For example: ```sh map1: (d0, d1, d2)[s0] -> (d0 + d2, s0 * 5) // d1 is unused map1.CompressDims(): (new_d0, new_d1)[new_s0] -> (new_d0 + new_d1, new_s0 * 5) map2: (d0)[s0, s1, s2] -> {d0 + s2, s0 * 5} // s1 is unused map2.CompressSymbols(): (new_d0)[new_s0, new_s1] -> (new_d0 + new_s1, new_s0 * 5) ``` `SymbolicMap` forms the mathematical basis for `IndexingMap`, that describes how tensor elements map to each other in [HLO semantics](./operation_semantics.md). `IndexingMap` consists of symbolic maps with domain-specific constraints. It enables shape and tiling analysis, and collapsing operation chains (like consecutive reshaping, transpose, broadcast, etc.) into optimized indexing calculation. Learn more with concrete examples in [Indexing analysis](./indexing.md). --- ### Terminology (`docs/terminology.md`) # XLA Terminology There are several terms that are used in the context of XLA, MLIR, LLVM, and other related technologies. Below is a partial list of these terms and their definitions. - **OpenXLA** - OpenXLA is an open ecosystem of performant, portable, and extensible machine learning (ML) infrastructure components that simplify ML development by defragmenting the tools between frontend frameworks and hardware backends. It includes the XLA compiler, StableHLO, VHLO, [PJRT](https://openxla.org/xla/pjrt) and other components. - **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. The XLA compiler outputs some code to LLVM, some to "standard" MLIR, and some to [Triton MLIR](https://triton-lang.org/main/dialects/dialects.html) that is processed by (MLIR-based) OpenAI Triton compiler. - **PJRT** - [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) is a uniform Device API that simplifies the growing complexity of ML workload execution across hardware and frameworks. It provides a hardware and framework independent interface for compilers and runtimes. - **StableHLO** - StableHLO is the public interface to OpenXLA, it is a standardized MLIR dialect that may be used by different frameworks and compilers in the OpenXLA ecosystem. XLA supports StableHLO, and immediately converts it to HLO on the input. There are some [StableHLO to StableHLO](https://openxla.org/stablehlo/generated/stablehlo_passes) passes implemented using the MLIR framework. It is also possible to convert StableHLO to other compilers' IR without using HLO, for example in cases where an existing IR is more appropriate. - **CHLO** - CHLO is a collection of higher level operations which are optionally decomposable to StableHLO. - **VHLO** - The [VHLO Dialect](https://openxla.org/stablehlo/vhlo) is a MLIR dialect that is a compatibility layer on top of StableHLO. It provides a snapshot of the StableHLO dialect at a given point in time by versioning individual program elements, and is used for serialization and stability. - **MHLO** - MHLO is a standalone MLIR-based representation of XLA's HLO IR. The dialect is being evaluated for deprecation, and new users of the dialect should prefer to use StableHLO instead. - **HLO** - HLO (High Level Optimizer) is an internal graph representation (IR) for the XLA compiler (and also supported input). It is **not** based on MLIR, and has its own textual syntax and binary (protobuf based) representation. - **MLIR** - [MLIR](https://mlir.llvm.org) is a hybrid IR infrastructure that allows users to define "dialects" of operations at varying degrees of abstraction, and gradually lower between these opsets, performing transformations at each level of granularity. StableHLO and CHLO are two examples of MLIR dialects. - **LLVM** - [LLVM](https://llvm.org/) is a compiler backend, and a language that it takes as an input. Many compilers generate LLVM code as a first step, and then LLVM generates machine code from it. This allows developers to reuse code that is similar in different compilers, and also makes supporting different target platforms easier. XLA:GPU and CPU backends have [LLVM IR emitters](https://github.com/openxla/xla/tree/main/xla/service/llvm_ir) for targeting specific hardware. --- ### Test Hlo Passes (`docs/test_hlo_passes.md`) # Writing unit tests for HLO passes There are different ways to write unit test for HLO passes. This page describes the preferred method to ensure consistency and readability. ## `FileCheck` with `CHECK` lines interleaved Most HLO passes can be tested using [`FileCheck`](https://llvm.org/docs/CommandGuide/FileCheck.html) tests. Interleave `CHECK` lines in input HLO module texts, and make sure to use `// CHECK` instead of `; CHECK` uniformly as the `FileCheck` delimiter. For example, you can re-write the [`fusion cc_test` for a `priotity_fusion` pass](https://github.com/openxla/xla/blob/fe30942a406659bff75399a2a10585bbd1287e07/xla/service/gpu/transforms/priority_fusion_test.cc#L133-L149) as follows: ``` TEST_F(PriorityFusionTest, FuseBroadcastIntoBitcastConsumers) { absl::string_view kHlo = R"( HloModule test_module // CHECK: ENTRY main ENTRY main { // CHECK-NEXT: %[[PARAM:.*]] = f32[96]{0} parameter(0) param_0 = f32[96]{0} parameter(0) broadcast = f32[8,96,128,7]{3,2,1,0} broadcast(param_0), dimensions={1} bitcast.6079.2 = f32[8,24,4,128,7]{4,3,2,1,0} bitcast(broadcast) // CHECK-NEXT: ROOT %{{.*}} fusion(%[[PARAM]]) {{.*}} ROOT transpose.1990.2 = f32[8,24,128,7,4]{4,3,2,1,0} transpose(bitcast.6079.2), dimensions={0,1,3,4,2} } )"; RunAndFilecheckHloRewrite(kHlo, std::move(priority_fusion_)); } ``` Note: Currently, the codebase has some tests where input HLO module and expected module are written separately. Inlining the `CHECK` lines is the preferred method for future tests. It enables better readability and a similar signature as MLIR based tests [like in `stablehlo_aggressive_folder.mlir`](https://github.com/openxla/stablehlo/blob/main/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir#L31-L39). ## `LIT` runner and `hlo-opt` Where feasible, use [`LIT`](https://llvm.org/docs/CommandGuide/lit.html) runner and `hlo-opt`, and place `CHECK` lines locally next to the input IR they correspond to. Again, make sure to use `// CHECK` instead of `; CHECK` as the delimiter. For example, some [GPU tests](https://github.com/openxla/xla/tree/main/xla/backends/gpu/tests) can be written as follows: ``` // RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../../../backends/gpu/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-%{PTX} %s HloModule Test, is_scheduled=true fused_computation { param_0 = f32[100,200]{1,0} parameter(0) ROOT b.1 = f32[200,100]{1,0} transpose(f32[100,200]{1,0} param_0), dimensions={1,0} } ENTRY main { a = f32[100, 200]{1,0} parameter(0) // CHECK-PTX: call void @llvm.nvvm.barrier0 // CHECK-GCN: call void @llvm.amdgcn.s.barrier ROOT wrapped_b = f32[200,100]{1,0} fusion(f32[100,200]{1,0} a), kind=kInput, calls=fused_computation } ``` ## Automated `CHECK`-generation script Writing test checks manually can be a lot of work, so it's often more practical to run an optimizer, read over the results to make sure they match expectations, and then convert the optimized HLO into `CHECK` directives. To simplify this process, you can use [`generate_hlo_test_checks.py`](https://github.com/openxla/xla/tree/main/xla/hlo/tools/generate_hlo_test_checks.py) to automatically insert generated `CHECK` directives above each test case in an HLO file. > IMPORTANT: This tool inherently assumes that the pass's current behavior is > correct, so make sure to look over the generated `CHECK` lines yourself and > confirm that they match the output you expect. ## (Don't) Graph traversal Refrain from writing tests that travel leaf nodes of the result graph and match with expected op. These tests are tedious to write, difficult to quickly read, and more difficult to debug and fix. Use one of the above options instead. --- ### Tiled Layout (`docs/tiled_layout.md`) # Tiled layout XLA TPU uses a variety of tile-based formats due to the two-dimensional nature of its vector registers. ## Tiled formats A tiled format breaks down a shape into tiles (usually 1D or 2D). Tiles are laid out in memory in major to minor order (row major layout). Within a tile, elements are also laid out in major to minor order.
Figure 1 Figure 1 shows how an array F32[3,5] is laid out in memory with 2x2 tiling. A shape with this layout is written as F32[3,5]{1,0:T(2,2)}, where 1,0 relates to the physical order of dimensions (`minor_to_major` field in Layout) while (2,2) after the colon indicates tiling of the physical dimensions by a 2x2 tile. Intuitively, tiles are laid out to cover the shape and then within each tile, elements are then laid out without tiling, as in the example above, where the right part of the example shows the layout in memory, including the white padding elements that are added in order to have complete 2x2 tiles even though the original array bounds are not even. The extra elements in the padding are not required to contain any particular value. ## Linear index formulas for tiling given a shape and a tile Without tiling, an element e=(en, en-1, ... , e1) in an array with array bounds d=(dn, dn-1, ... , d1) (d1 is the most minor dimension) is laid out by major to minor order at position:    linear_index(e, d)
= linear_index((en, en-1, ... , e1), (dn, dn-1, ... , d1))
= endn-1...d1 + en-1dn-2...d1 + ... + e1 For simplicity of notation in this document we assume a tile has the same number of dimensions as the array. In XLA's implementation of tiling, this is generalized to tilings with fewer dimensions by leaving the initial most-major dimensions unchanged and applying the tiling only to the most minor dimensions, so that the tiling that is specified mentions a suffix of the physical dimensions of the shape being tiled. When tiling of size (tn, tn-1, ... , t1) is used, an element in the array with indices (en, en-1, ... , e1) is mapped to this position in the final layout:    linear_index_with_tile(e, d, t)
= linear_index((⌊e/t⌋, e mod t), (⌈d/t⌉, t))     (arithmetic is elementwise, (a,b) is concatenation)
= linear_index((⌊en/tn⌋, ... , ⌊e1/t1⌋, en mod tn, ... , e1 mod t1), (⌈dn/tn⌉, ... , ⌈d1/t1⌉, tn, tn-1, ... , t1))
= linear_index((⌊en/tn⌋, ... , ⌊e1/t1⌋), (⌈dn/tn⌉, ... , ⌈d1/t1⌉))∙tntn-1...t1 + linear_index((en mod tn, ... , e1 mod t1), (tn, tn-1, ... , t1)) The layout can be thought of as having two parts: (⌊en/tn⌋, ... , ⌊e1/t1⌋), which corresponds to a tile index in an array of tiles of size (⌈dn/tn⌉, ... , ⌈d1/t1⌉), and (en mod tn, ... , e1 mod t1), which corresponds to a within-tile index. The ceil function appears in ⌈di/ti⌉ because if tiles overrun the bounds of the larger array, padding is inserted as in Figure 1. Both the tiles and elements within tiles are laid out recursively without tiling. For the example in Figure 1, element (2,3) has tile index (1,1), and within-tile index (0,1), for a combined coordinate vector of (1,1,0,1). The tile indices have bounds (2,3) and the tile itself is (2,2) for a combined vector of (2,3,2,2). The linear index with tile for the element with index (2,3) in the logical shape is then    linear_index_with_tile((2,3), (3,5), (2,2))
= linear_index((1,1,0,1), (2,3,2,2))
= linear_index((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index((0,1), (2,2))
= (1 ∙ 3 + 1) ∙ 2 ∙ 2 + (0 ∙ 2 + 1)
= 17. $$ LinearIndexWithTile((2,3), (3,5), (2,2)) \\ = LinearIndex((1,1,0,1), (2,3,2,2)) \\ = LinearIndex((1,1), (2,3)) \cdot 2 \cdot 2 + LinearIndex((0,1), (2,2)) \\ = (1 \cdot 3 + 1) \cdot 2 \cdot 2 + (0 \cdot 2 + 1) \\ = 17 $$ ## Tiling as pad-reshape-transpose Tiling-based layout operates as follows:
Consider an array of dimensions (dn, dn-1, ... , d1) (d1 is the most minor dimension). When it’s laid out with tiling of size (tn, tn-1, ... , t1) (t1 is the most minor dimension), that tiling can be described in terms of pad-reshape-transpose in the following way. 1. The array is padded to (⌈dn/tn⌉∙tn, ... , ⌈d1/t1⌉∙t1). 2. Each dimension i is broken into (⌈di/ti⌉, ti), i.e. the array is reshaped to
    (⌈dn/tn⌉, tn, ... , ⌈d1/t1⌉, t1).
There is no physical layout change in this reshape by itself, so this reshape is a bitcast. If one is not explicitly thinking of a tiling, this reshape could express any shape with the same number of elements as the padded shape - the example here is of how to express a tile in this way. 3. A transpose happens by moving tn, ... , t1 to the most minor dimensions while keeping their relative order, so that the order of dimensions from most major to most minor becomes
    (⌈dn/tn⌉, ... , ⌈d1/t1⌉, tn, ... , t1). The final shape has the prefix
    (⌈dn/tn⌉, ... , ⌈d1/t1⌉), which describes the number of tiles in each dimension. An element in the array (en, ... , e1) is mapped to this element in the final shape:
    (⌊en/tn⌋, ... , ⌊e0/t0⌋, en mod tn, ... , e1 mod t1). It is easy to see that the linear index of the element follows the formula above as expected. ## Repeated tiling XLA's tiling becomes even more flexible by applying it repeatedly.
Figure 2 Figure 2 shows how an array of size 4x8 is tiled by two levels of tiling (first 2x4 then 2x1). We represent this repeated tiling as (2,4)(2,1). Each color indicates a 2x4 tile and each red border box is a 2x1 tile. The numbers indicate the linear index in memory of that element in the tiled format. This format matches the format used for BF16 on TPU, except that the initial tile is bigger, namely the tiling is (8,128)(2,1), where the purpose of the second tiling by 2x1 is to collect together two 16-bit values to form one 32-bit value in a way that aligns with the architecture of a TPU. Note that a second or later tile can refer to both the minor within-tile dimensions, which just rearranges data within the tile, as in this example with (8,128)(2,1), but can also refer to the major cross-tile dimensions from the prior tiling. ## Combining dimensions using tiles XLA's tiling also supports combining dimensions. For example, it can combine dimensions in F32[2,7,8,11,10]{4,3,2,1,0} into F32[112,110]{1,0} first before tiling it with (2,3). The tile used is (∗,∗,2,∗,3). Here an asterisk in a tile implies taking that dimension and combining it with the next more minor dimension. Multiple adjacent dimensions can be subsumed together into one dimension. A subsumed dimension is represented by a tile value of -1 in that dimension of the tile, which is not otherwise valid in a tile as a dimension size. More precisely, if dimension i of the shape is eliminated via an asterisk in the tile, then before the prior definition of tiling is applied, that dimension is removed from both the shape being tiled and the tile vector, and what was dimension i-1 of the shape has its array bound increased from di-1 to didi-1. This step is repeated for each asterisk in the tile vector. ## Examples of tiling formats This section shows examples of popular XLA formats. 1. **Untiled format** - Most arrays not on the TPU use an untiled linear format, same as in e.g. C++. 2. **TPU tile format** - The most common format in XLA/TPU is tiling by `8x128`, which matches the 32-bit `8x128` vector registers on a TPU. 3. **TPU small tile format** (a.k.a. "Compact 2nd Minor Layout") - When the second most minor dimension size is 1 or 2, XLA/TPU instead tiles by `2x128` to save memory since a `2x128` tile is smaller than an `8x128` tile. When the second most minor dimension size is 3 or 4, XLA/TPU tiles by `4x128`. 4. **TPU 16 bit tile format** - When array element type is BF16, the tiles we use are typically (8,128)(2,1). The second level of tiling does the so-called BF16 packing. See Figure 2 above, one element from an even row and one element from an odd row are laid out together and put in one 32-bit element. This format is used because TPUs work with 32 bit values natively and it is much more efficient to move data across the second most minor dimension than across the most minor dimension, so collecting two 16 bit values to get 32 bits from the same column is much more efficient than doing it in the more obvious fashion of taking two 16 bit values from the same row. 5. **TPU 8 bit tile format** - The format here is very similar to the 16 bit format, the difference is just that we need to collect together 4 elements to get 32 bits instead of just the two, so the tiling becomes `(8,128)(4,1)`. 6. **TPU 1 bit tile format** - TPUs currently use 1 byte for one boolean value, i.e. the size in bytes of the PRED element type is 1. It would be less wasteful to use a tiling by `(32,128)(32,1)` and use only 1 bit per element. --- ### Tools (`docs/tools.md`) # XLA Tooling The XLA development workflow is usually centered around [HLO](./operation_semantics) IR, which represents isolated functional computation given to the compiler. XLA comes with multiple command line tools (described below) which consume HLO and either run it, or provide an intermediate compilation stage. Using such tools is invaluable for a fast `compile->modify->run` iteration cycle, as HLO is both visualizable and hackable, and iteratively changing and running it is often the fastest way to understand and to fix an XLA performance or behavior. The easiest way to obtain the HLO for a program being compiled with XLA is usually to use the `XLA_FLAGS` environment variable: ``` $ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point ``` which stores all before-optimization HLO files in the folder specified, along with many other useful artifacts. ## [`run_hlo_module`] Run HLO Modules ``` bazel run //xla/tools:run_hlo_module -- [flags] ``` The tool `run_hlo_module` operates on pre-optimization HLO, and by default bundles compilation, running and comparison with the reference interpreter implementation. For example, the usual invocation to run an input file `computation.hlo` on an NVIDIA GPU and to check it for correctness is: ``` run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo ``` ### Run Multiple HLO Modules Invocation with multiple HLO modules is supported for `run_hlo_module`. To run all hlo modules from a directory: ``` bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations* ``` ## [`multihost_hlo_runner`] Run HLO Modules With SPMD Support ``` # Note: Binary name is `hlo_runner_main`. bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] ``` Multihost HLO runner is a very similar tool, with the caveat that it supports SPMD, including cross host communication. See [Multi-Host HLO Runner](./tools_multihost_hlo_runner) for details. ### Run Multiple HLO Modules With SPMD Support Similar to `run_hlo_module`, `multihost_hlo_runner` also supports invocation with multiple modules. ``` bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations* ``` ## [`hlo-opt`] Compile HLO Module ``` bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] ``` When debugging or understanding the workings of the compiler, it is often useful to get the expansion for a particular hardware at a particular point in the pipeline (be it HLO, optimized HLO, TritonIR or LLVM), for a given HLO or StableHLO input. `hlo-opt` supports multiple output stages: be it PTX, HLO after optimizations, LLVM IR before optimizations, or TritonIR. The exact set of stages supported depends on the platform (as e.g. PTX is NVIDIA-specific), and can be seen using the --list-stages command: ``` hlo-opt --platform=CUDA --list-stages buffer-assignment hlo hlo-backend html llvm llvm-after-optimizations llvm-before-optimizations ptx ``` After selecting a stage, the user can write the result of the conversion for a given platform to a given stream: ``` hlo-opt --platform=cpu --stage=hlo input.hlo ``` which would print the dump to stdout (or to a given file if `-o` was specified). ### Deviceless Compilation for GPU Deviceless compilation do not need access to a GPU. The Deviceless Compilation provides a way to specify a GPU spec (`--xla_gpu_target_config_filename`) or GPU topology (`--xla_gpu_topology_filename`) on the command line for stages where access to a GPU or multi-GPU topology is required, eliminating a need for GPU device. Example: PTX output without access to a gpu device: ``` hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=xla/backends/gpu/target_config/specs/a100_pcie_80.txtpb input.hlo ``` Specs for popular GPUs are shipped with the compiler, and the provided file is string serialization of `device_description.proto`: ``` gpu_device_info { cuda_compute_capability { major: 8 minor: 0 } threads_per_block_limit: 1024 threads_per_warp: 32 shared_memory_per_block: 127152 shared_memory_per_core: 65536 threads_per_core_limit: 2048 core_count: 6192 fpus_per_core: 64 block_dim_limit_x: 2147483647 block_dim_limit_y: 65535 block_dim_limit_z: 65535 memory_bandwidth: 2039000000000 l2_cache_size: 4194304 clock_rate_ghz: 1.1105 device_memory_size: 79050250240 } platform_name: "CUDA" ``` More GPU specs are located at `xla/backends/gpu/target_config/specs`. #### Multi-GPU Topology and Collectives Passes that lower or optimize collective operations require a GPU topology (number of partitions, hosts per partition, and devices per host). Use `--xla_gpu_topology_filename` to provide either: * A serialized `GpuTopologyProto` file from `xla/backends/gpu/target_config/specs` (such as `h100_1x8.txtpb`, `h100_2x8.txtpb`, `gb200_2x4.txtpb`, or `gb300_4x4.txtpb`), which also automatically populates the target device configuration from `platform_version`: ``` hlo-opt --platform=gpu --stage=hlo --xla_gpu_topology_filename=xla/backends/gpu/target_config/specs/gb200_2x4.txtpb input.hlo ``` * An inline topology string of the form `[:]xx` (or `[:]x`). When `` is omitted, pair it with `--xla_gpu_target_config_filename`: ``` hlo-opt --platform=gpu --stage=hlo --xla_gpu_topology_filename=oberon_b200:1x2x4 input.hlo hlo-opt --platform=gpu --stage=hlo --xla_gpu_topology_filename=1x2x4 --xla_gpu_target_config_filename=xla/backends/gpu/target_config/specs/gb200.txtpb input.hlo ``` #### Autotuning Sometimes compilation may involve autotuning based on a compilation `--stage`. For the deviceless compilation to work, the user either need to \ **disable** autotuning with `--xla_gpu_autotune_level=0`\ or\ **load a pre-existing autotuning results** with `--xla_gpu_load_autotune_results_from=` (obtained with `--xla_gpu_dump_autotune_results_to=`). ``` hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=xla/backends/gpu/target_config/specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo ``` The autotune file is text serialization of `autotune_results.proto`, with example looking like: ``` version: 3 results { device: "CUDA: 8.0, Cores: 108, GPU clock: 1.41 GHz, Memory bandwidth: 1555 GB/s, L2 cache: 40 MB" hlo: "{\n tmp_0 = f16[1,16,17,3]{3,2,1,0} parameter(0)\n tmp_1 = f16[16,51]{1,0} bitcast(f16[1,16,17,3]{3,2,1,0} tmp_0)\n tmp_2 = s8[16,17,3]{2,1,0} parameter(1)\n tmp_3 = s8[51,16]{0,1} bitcast(s8[16,17,3]{2,1,0} tmp_2)\n tmp_4 = f16[51,16]{0,1} convert(s8[51,16]{0,1} tmp_3)\n tmp_5 = f16[16,16]{1,0} dot(f16[16,51]{1,0} tmp_1, f16[51,16]{0,1} tmp_4), lhs_contracting_dims={1}, rhs_contracting_dims={0}\n ROOT tmp_6 = f16[1,16,16]{2,1,0} bitcast(f16[16,16]{1,0} tmp_5)\n}" result { run_time { nanos: 31744 } triton { block_m: 32 block_n: 32 block_k: 32 split_k: 1 num_stages: 1 num_warps: 4 } } } ``` The autotuning database can be serialized using `XLA_FLAGS=--xla_gpu_dump_autotune_results_to=` ## [`hlo-opt`] HLO Pass Development And Debugging ``` # If you are working with hardware independent passes from the # `xla/hlo/transforms/` directory, prefer light-weight version # of the `hlo-opt` tool with fewer dependencies: bazel run //xla/hlo/tools:hlo-opt -- [flags] # Otherwise, for hardware independent and CPU, GPU passes use # the same binary from "Compile HLO Modules" section above: bazel run //xla/tools:hlo-opt -- [flags] ``` The `hlo-opt` tool allows execution of an individual passes independent of the given platform compilation stages. This isolation helps to quickly run passes on input hlo module and pinpoint the root cause of failures. ``` hlo-opt --passes=schedule-aware-collective-cse input.hlo ``` Note: `--platform` option is not required. `hlo-opt` tool also supports [`DebugOptions XLA_FLAGS`](https://github.com/openxla/xla/blob/5bf1e6420d250dce5eb840889096bdf8aad6f432/xla/xla.proto#L40-L1197). ``` hlo-opt --passes=schedule-aware-collective-cse --xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo ``` Use`--list-passes` option to get the pass name string. ``` hlo-opt --list-passes ``` Users can create their own custom pipeline by specifying more than one passes to `--passes` option. ``` hlo-opt --passes=pass1,pass2,pass3 input.hlo ``` ### Assist New HLO Pass Development 1. First, write your pass. 1. Register the new pass to the `hlo-opt` tool pass registry. ``` RegisterPass(FooPassInputOptions) ``` Based on the pass type, choose one of the following locations for registration:\ [`opt_lib.cc`](https://github.com/openxla/xla/blob/5d015a2ddfcf4f40934a33891dc63471704f221d/xla/hlo/tools/hlo_opt/opt_lib.cc) Hardware-independent passes.\ [`cpu_opt.cc`](https://github.com/openxla/xla/blob/5d015a2ddfcf4f40934a33891dc63471704f221d/xla/tools/hlo_opt/cpu_opt.cc) CPU specific passes.\ [`gpu_opt.cc`](https://github.com/openxla/xla/blob/5d015a2ddfcf4f40934a33891dc63471704f221d/xla/tools/hlo_opt/gpu_opt.cc) GPU specific passes.\ [`compiled_opt.cc`](https://github.com/openxla/xla/blob/5d015a2ddfcf4f40934a33891dc63471704f221d/xla/tools/hlo_opt/compiled_opt_lib.cc) Passes common to CPU, GPU, XPU.\ Don't forget to add build dependency. Include pass registration as part of your PR([example](https://github.com/openxla/xla/pull/22968/files#diff-e37a0ea999dfc5764d624240cd2edebb8b7ee4e6d91686be89c632dd7203b823)) so that the pass will be available to use for all `hlo-opt` users. 1. Rebuild the `hlo-opt` tool, validate successful pass registration using `--list-passes` option and then use `--passes` option to run the pass. ``` $ hlo-opt --passes=foo-pass input.hlo ``` 1. Writing unit tests for the pass? refer https://openxla.org/xla/test_hlo_passes for more details. ### Pass Runtime Measurement For large models, full compilation runs can take upto few minutes, making it challenging to detect subtle performance regressions. In contrast, individual pass runs using `hlo-opt` allow for precise performance measurement and the easy detection of even small increases in execution time caused by new code changes. ``` time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier --xla_reduce_window_rewrite_base_length=128 input.hlo ``` ## [`hlo-opt`] Convert HLO Module Formats ``` # Use the light weight version of the `hlo-opt` tool. bazel run //xla/hlo/tools:hlo-opt -- [flags] ``` #### Convert `HLO Text` -> `HLO Proto` ``` hlo-opt --emit-proto input.hlo ``` #### Convert `HLO Proto` or `HLO Proto Binary` -> `HLO Text` ``` hlo-opt input.pbtxt or input.pb ``` ## [`ptx-opt`] Compiler LLVM Module down to PTX The tool will run LLVMIR optimization pipeline and then call CompileToPtx. ``` bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 ``` The tool can also dump LLVMIR after every path. ``` bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to= --xla_gpu_dump_llvmir ``` ## [`hlo_isolation_test`] Isolate and Verify HLO Stability For automated debugging and verifying numeric stability or mismatches across compiled HLO modules, use the `hlo_isolation_test` CLI. For detailed usage, installation, and API instructions, see the [HLO Isolation User Guide](./hlo_isolation.md). ## [`isolate_hlo`] Isolate Problematic HLO Instructions If you have a large HLO dump and suspect a specific instruction or section within the HLO module is causing the crash, you can use the `isolate_hlo` tool. This tool extracts a single HLO instruction (and its necessary context) into a new, smaller HLO module. This is extremely helpful for creating a minimal, compiler-level reproducer. * **Documentation & Source:** The `isolate_hlo` tool is available in the OpenXLA repository. See the `xla/tools` directory in the XLA source code. * **Usage:** Build the tool from the XLA source tree. It typically takes an input HLO module file (text or proto), the name of the instruction to extract, and output file path. ```bash # Example usage after building XLA: # ./build/tools/isolate_hlo --input=module.hlo --instruction_name=fusion.123 \ # --output=isolated_fusion.123.hlo --input_format=txt --output_format=long_txt ``` Refer to the tool's help message (`--help`) for specific flags and format options. --- ### Tools Multihost Hlo Runner (`docs/tools_multihost_hlo_runner.md`) # Multi-Host HLO Runner [TOC] This tool lets you run an HLO module on one or more GPUs. It also allows compiling code targeting multiple GPUs without running it. ## Running multi-GPU (sharded) HLOs We can identify these HLOs by seeing `sharding=` annotations. For example `sharding={devices=[1,1,2,1]0,1}` means that the annotated tensor should be sharded to 2 GPUs (GPU0 and GPU1) along the 3rd dimension. The following instructions assume the working directory is the XLA Git repository and that `./configure.py` has been run. If we have enough GPUs, we can replay these HLOs like this: ``` bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- my-hlo.txt ``` Tip: If the input generation takes too long or uses too much host memory, consider using `--hlo_argument_mode=uninitialized`. It is also possible to compile the same HLO without running it by setting `--compile_only=true`: ``` bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- --compile_only=true my-hlo.txt ``` In that case, a single GPU is necessary, unless the [autotuning cache](./persisted_autotuning) is used. ### Troubleshooting - Errors such as `Check failed: result.replicas >= 1 (0 vs. 1)`: - We have to make sure that we have enough GPUs. - `CUDA_VISIBLE_DEVICES` must be set correctly or not set at all. - Crashes: - We may want to use `--dynamic_mode=off`. - CUDA and cuDNN should be set up correctly. # Examples ## Single process, multiple GPU example ### Setup and get the HLO ``` {note} You can use a container with the following instructions: docker run -it --shm-size=1g --gpus all ghcr.io/nvidia/jax:pax-2024-06-03 cd /opt/xla/ Note, those instructions can be outdated more quickly. Adjust as needed. ``` ``` # The 8 below is the number of GPUs you have. # test-pax.sh --help for more details on the parallelization options (export XLA_FLAGS="--xla_dump_to=/tmp/dump"; test-pax.sh --fsdp 8 --batch-per-gpu 1) ls -lSh /tmp/dump/*before_optimizations.txt # The biggest file one is normally the one you care about. # I picked one, for the rest of the scripts, but the name could change when you change the JAX or XLA version. ``` ### Build XLA multihost runner ``` cd /opt/xla/ ./configure.py --backend CUDA --nccl bazel build //xla/tools/multihost_hlo_runner:hlo_runner_main ``` ### Single process example: Before optimization graph replay ``` bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \ /tmp/dump/module_0023.pjit__wrapped_step_fn.before_optimizations.txt ``` ### Single process example: After optimization graph replay To replay an optimized HLO, you must use either `--xla_disable_all_hlo_passes` or `--run_xla_backend_only`. Otherwise, XLA will try to recompile the HLO and this isn't supported. So it will give you many strange errors. Full command: `bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- --run_xla_backend_only /tmp/dump/module_0023.pjit__wrapped_step_fn.sm_8.0_gpu_after_optimizations.txt` ## Multi-processes, single-node ### Launch container Also install some missing libraries. (Note, that can be outdated more quickly. Adjust as needed.) ``` docker run -it --shm-size=1g --gpus all ghcr.io/nvidia/jax:pax-2024-06-03 apt-get update && apt-get install -y openmpi-bin openmpi-common libopenmpi-dev ``` ### Run original model and dump HLO For this example, we will use an 8-GPU PAXML model from `test-pax.sh`. (Note this will be the same dump as the single process case. So you can do `cp -r /tmp/dump /tmp/dump_multi_process` if you already have it. `export XLA_FLAGS="--xla_dump_to=/tmp/dump_multi_process" mpirun --allow-run-as-root -np 8 test-pax.sh --fsdp 8 --batch-per-gpu 1 -o /tmp/checkpoint --multiprocess` The HLO dump will be saved to `/tmp/dump_multi_process/`. For PAX specifically, the main module will have "pjit__wrapped_step_fn" in the name. For this example we will use `/tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt`. ### Run on a single node using MPI Create a bash script called `run.sh`: ``` #!/bin/bash export CUDA_VISIBLE_DEVICES=${OMPI_COMM_WORLD_LOCAL_RANK} bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \ --task_id=${OMPI_COMM_WORLD_RANK} \ --num_nodes=${OMPI_COMM_WORLD_SIZE} \ --address=127.0.0.1:12345 \ /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt ``` Now, you can execute it using mpirun: ``` chmod a+x run.sh mpirun --allow-run-as-root -np 8 run.sh ``` ### Run on multiple nodes with SLURM When running on multiple nodes using SLURM, you can forward the SLURM env variables to the HLO runner like so in your SLURM job: ``` bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- \ --task_id=${SLURM_PROCID} \ --num_nodes=${SLURM_NTASKS} \ --address="${SLURM_LAUNCH_NODE_IPADDR}:12345" \ /tmp/dump_multi_process/module_0023.pjit__wrapped_step_fn.before_optimizations.txt ``` --- ### Errors Error 0100 (`docs/errors/error_0100.md`) # Error code: E0100 **Category:** Runtime: Buffer Allocation Failure This error indicates that XLA:TPU runtime’s memory allocator failed to find a suitable block of memory on the accelerator’s HBM for the requested allocation. **Sample error message:** ``` ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0) ``` **XLA backends:** TPU ## Overview This error is thrown on: - Failures of user-initiated buffer allocation via [`jax.device_put`](https://docs.jax.dev/en/latest/_autosummary/jax.device_put.html) or - Failures of user-scheduled program's output allocations. These failures are typically caused due to a couple of reasons: - **Out of Memory (OOM):** The user is trying to allocate a chunk of memory that is larger than the total amount of free memory available on the TPU’s HBM. - **Memory fragmentation:** The allocation fails because **no single contiguous free block** in the memory space is large enough to satisfy the requested size. The total amount of free memory is sufficient for the allocation, but it is scattered across the memory space in small, non-contiguous blocks. The TPU runtime has a number of mechanisms in-place to retry allocation failures including: - If there are queued deallocations, the runtime retries failed allocations. - On OOMs caused by a fragmentation, the runtime can automatically trigger a defragmentation and a retry. - The TPU runtime prioritizes buffer allocations over keeping programs loaded. If a buffer allocation fails due to insufficient HBM, the system will evict loaded TPU programs until enough memory is available for the buffer. So an error encountered after the above mitigations typically require user action. ## Debugging - **Reduce your model's memory footprint** - **Decrease batch size:** Reducing the batch size directly lowers memory usage. - **Parameter sharding:** For very large models, use techniques like model parallelism or sharding to distribute parameters across the HBM of multiple TPU cores or hosts. - **Shorten sequence/context length:** For models that operate on sequences (like language models), reducing the input sequence length can significantly decrease the memory footprint. - **Buffer donation:** Utilize framework features (such as: `jax.jit(..., donate_argnums=...)`) to signal to XLA that certain input buffers can be overwritten and reused for outputs. Read [Buffer donation](https://docs.jax.dev/en/latest/buffer_donation.html) for more details. - **Optimize checkpoint strategy:** Instead of saving the entire model state at once, consider saving only the model weights or using a sharded checkpointing strategy. - **Address memory layout and padding** - TPU memory is allocated in chunks, and padding can increase the actual size of tensors. - **Ensure no memory leaks** - Ensure references to `jax.Array` objects are not being held longer than intended. Holding on to `jax.Array` objects might prevent automatic de-allocation even after program compilation is completed. See also [Error code: E1000](https://openxla.org/xla/errors/error_1000) for other strategies you can use to reduce the amount of memory each program uses. ### Tooling Enable the `tpu_log_allocations_on_oom` flag for which the allocator will dump a detailed report of all current allocations when an OOM occurs, which can be invaluable for debugging. --- ### Errors Error 0101 (`docs/errors/error_0101.md`) # Error code: E0101 **Category:** Runtime: Program Allocation Failure This error indicates that the XLA runtime on a TPU device failed to load a compiled XLA program executable into the TPU's HBM. **Sample error message:** ``` XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well). ``` **XLA backends:** TPU ## Overview This error is typically caused by one of the following reasons: - **Program size exceeds available HBM:** The compiled XLA program, including its instructions, static data, and any embedded constants, is larger than the total amount of free HBM currently available on the specific TPU core(s) where the program is being loaded. - **HBM fragmentation:** While the total free HBM on the device might be sufficient in aggregate, it is not available in a single, contiguous block large enough to fit the entire program. It's important to understand how the TPU runtime prioritizes memory. Buffer allocations are privileged over loaded programs. If a buffer allocation fails, the runtime will evict already loaded programs from HBM to free up space. This can lead to a situation where a program that loaded successfully before now fails with an OOM error, because the HBM is now occupied with more data buffers. ## Debugging - **Reduce buffer memory footprint:** Freeing up memory used by data buffers will leave more room for the program itself. - **Decrease batch size:** This is one of the most effective ways to reduce the amount of memory used for activations. - **Parameter sharding:** For very large models, use model parallelism or sharding techniques (like FSDP or Megascale) to distribute the model's parameters and computation across multiple TPU cores or hosts. Note that this can increase the network communication overhead due to splitting tensors across multiple chips. - **Shorten sequence/context length:** For models processing sequential data (e.g., NLP models), reducing the sequence length can significantly decrease memory usage. - **Buffer donation:** Use framework features (e.g., `jax.jit(..., donate_argnums=...)`) to allow XLA to reuse the memory of input buffers for storing output, reducing peak memory usage. - **Reduce program’s memory requirements for temporaries** - Reduce programs memory usage for temporaries by using the `tpu_shared_memory_percent` flag. Note that this might negatively affect performance. - **Optimize execution strategy/reduce serving load** - **Manage program loading:** If you are JIT-compiling multiple functions, be aware that each function can result in a program being loaded. Try to structure your workload to minimize the number of concurrently loaded programs. - **Ensure no memory leaks** - Ensure references to `jax.Array` objects are not being held longer than intended. Holding on to `jax.Array` objects might prevent automatic de-allocation even after program compilation is completed. See also [Error code: E1000](https://openxla.org/xla/errors/error_1000) for other strategies you can use to reduce the amount of memory each program uses. ### Tooling - Enable the `tpu_log_allocations_on_oom` flag for which the allocator will dump a detailed report of all current allocations when an OOM occurs, which can be invaluable for debugging. - Profile Your Program: Use the JAX memory profiler or the TensorFlow profiler to get a detailed view of your program's memory usage over time. This can help identify unexpected peaks in memory consumption. --- ### Errors Error 0102 (`docs/errors/error_0102.md`) # Error code: E0102 **Category:** Runtime: Program Input Mismatch This error occurs when the XLA runtime detects that an input buffer provided at execution time does not match the metadata expected by a compiled program, such as the buffer size or PJRT memory space kind. **Sample error message:** ``` XlaRuntimeError: INVALID_ARGUMENT: Executable(jit_embedding_pipeline_step_fn) expected parameter 2482 of size 5242880 (bf16[16,1280,40]{2,1,0:T(8,128)(2,1)}) but got buffer with incompatible size 1638400 (bf16[16,1280,40]{1,2,0:T(8,128)(2,1)}): while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well). ``` **XLA backends:** TPU ## Overview For buffer-size mismatches, the error message indicates both the expected and actual sizes, as well as the tensor shapes and layouts. Note that these errors might occur even if two tensors have the same shape but their size in memory can be different if their physical layout (how the data is tiled and arranged on the hardware) is different. For memory-space mismatches, the error message reports that the parameter buffer is in an unexpected memory space. These errors are predominantly caused by: - **Checkpoint and XLA configuration mismatch:** A model is trained and a checkpoint is saved. The physical layout of the weights in that checkpoint is determined by the exact XLA version and configuration (e.g. XLA flags) at that time. Later, this checkpoint is loaded in a different environment where the configuration has changed. A new flag, a different default value, or a change in the model/XLA code can cause the runtime to expect a different physical layout for the weights. When the old buffer from the checkpoint is passed to the new compiled XLA program, the runtime throws an error. - **Hardware/topology-specific layouts:** The XLA compiler is free to choose different physical layouts for tensors to optimize performance on different hardware. A layout that is optimal for v4 TPU might be different from a v5 TPU, or even for different pod slices of the same chip (e.g., 4x4x4 vs 4x8). The error occurs when a model is compiled with an assumption about one topology's layout, but at runtime it is scheduled on a different topology, or there is a bug in the compiler's layout logic for a specific piece of hardware. ## Debugging - **Ensure configuration consistency between model export and re-runs from checkpoints** - Avoid using old checkpoints with new code unless you are certain that no layout-affecting changes have been made. - If you suspect a checkpoint/configuration mismatch, the most reliable solution is to re-export the saved model using the exact same (and current) codebase and configuration that you are using for inference or fine-tuning. - Check for configuration changes (e.g. XLA flags) between the two runs. - **Hardware/topology-specific layouts** - Check for hardware version and topology mismatches if switching hardware or topologies. --- ### Errors Error 0200 (`docs/errors/error_0200.md`) # Error Code: E0200 **Category:** Runtime: Core Halted Unexpectedly This error indicates that a TPU core stopped executing instructions prematurely. This is a fatal error state where the hardware forces a halt due to an unrecoverable fault, a violation of hardware constraints, or a deliberate interrupt triggered by compiler-generated runtime assertions. **Sample error message:** ``` INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:1:0x1d9 ... ``` **XLA backends:** TPU ## Overview XLA compiles JAX programs into a sequence of low-level assembly instructions. At runtime, the TPU device executes these instructions sequentially. A "Core Halted Unexpectedly" error occurs when the TPU hardware encounters an unrecoverable condition that prevents further execution, forcing the core into a fatal "HALTED" state. Because this error can stem from physical hardware failures, compiler bugs, or user code issues (particularly in custom kernels), you must carefully analyze the log messages to identify the specific cause. ## Debugging To resolve this error, you must first identify which of the three specific scenarios caused the unexpected halt. Check your logs for the specific text signatures described below. - **"observed errors are: [Hardware/Network/Power]"**: This indicates a physical infrastructure failure. → Jump to [Scenario 1: Infrastructure failures (hardware/newtwork/power)](#scenario_1_infrastructure_failures_hardwarenetworkpower) - **"observed errors are: [User]"**: This indicates a hardware constraint violation. → Jump to [Scenario 2: Hardware constraint violations](#scenario_2_hardware_constraint_violations) - The error message contains specific details such as the following keywords: `BoundsCheck`, `scheckne`, `scheckeq`, `schecklt`, `scheckge`, `scheckbetween` This indicates that a compiler-generated assertion in the compiled program failed during execution. → Jump to [Scenario 3: XLA compiler-generated assertion failures](#scenario_3_xla_compiler-generated_assertion_failures) ### Scenario 1: Infrastructure failures (hardware/network/power) **Signature:** The logs explicitly state `observed errors are: [Hardware]` or `observed errors are: [Network]` or `observed errors are: [Power]`. This indicates a physical infrastructure failure unrelated to your software or model logic. The TPU chip, the network fabric connecting the chips, or the power supply has failed. - **Retry the job:** If the issue was a transient voltage dip or network flap, a simple retry may work. - **Identify and remove bad nodes:** If the error persists on the same specific task or host, the hardware is likely defective. Use your cluster management tools to "drain"/"cordon" the affected node and restart your job on healthy nodes. ### Scenario 2: Hardware constraint violations **Signature:** The logs state `observed errors are: [User]`. This indicates that the XLA compiler generated an instruction that violated a inviolable hardware constraint (e.g., an instruction attempting to access an out-of-bounds memory address on HBM or Scratchpad memory). While labeled "User", this is rarely caused by high-level user code. - **File an XLA bug:** This is likely a compiler bug, the compiler should never emit instructions that violate hardware specs. [Please file a bug report](https://github.com/openxla/xla). ### Scenario 3: XLA compiler-generated assertion failures **Signature:** The error message contains specific details on the compiler-generated assertion that is failing. Look for the for the following keywords: - `BoundsCheck`, `scheckne`, `scheckeq`, `schecklt`, `scheckge`, `scheckbetween` This indicates that a compiler-generated assertion in the compiled program failed during execution. Analyze the specific error message to determine the sub-type. #### Scenario 3.A: Launch group mismatch **Sample error message:** ``` Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:1:0x1d9 (from TensorCoreSequencer:1:0x309): scheckne: An unexpected leader shows up in the launch group with a different launch id than the current group leader. ``` **Cause:** This error typically occurs in multi-host TPU environments. It indicates that the TPU cores, which are expected to execute the same program in a synchronized manner (as part of a "launch group") have become out of sync. Specifically, a TPU core joined a synchronization group with a different program identifier than the current group leader, suggesting inconsistent programs across hosts. - **Verify XLA flags:** Ensure all hosts use the exact same `XLA_FLAGS`. - **Check for consistent JAX programs:** Check that all hosts are executing identical JAX programs. Verify docker images, libtpu versions, etc. #### Scenario 3.B: Bounds check failure **Sample error message:** ``` Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:23:0x292 (from TensorCoreSequencer:23:0xd74a): BoundsCheck 92 [deref of %s931] for %937 = dma.hbm_to_vmem [thread:$0] /*hbm=*/%s931, /*size_in_granules=*/16384, /*vmem=*/%s935, /*dst_syncflagno=*/%s860, /*src_stride=*/512, /*dst_stride=*/128, /*steps_per_stride=*/8 ``` **Cause:** The program tried to access memory outside of allocated bounds. The error message often includes details about the memory access type (e.g., `dma.hbm_to_vmem`) and the address calculation. - **Debug custom kernels:** If using Pallas, check your index calculations. Use [`pl.debug_print`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.debug_print.html) or [`checkify`](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.checkify.check.html) to validate tensor indices. - **Check sharding:** Ensure sharding annotations are consistent with tensor shapes. #### Scenario 3.C: Mosaic/Pallas synchronization **Sample error message:** ``` Core halted unexpectedly: INTERNAL: Accelerator device halted prematurely, perhaps due to an on-device check-failure. Node 0 halted unexpectedly at tag:pc TensorCoreSequencer:21:0xae5 (from TensorCoreSequencer:21:0x54c5): Semaphore (scratch argument 1) has a nonzero value upon exit from a Mosaic kernel. Make sure every DMA is awaited, and every semaphore signal is paired with a wait. ``` **Cause:** This error is specific to code generated by the Mosaic compiler (used by Pallas JAX). It indicates a synchronization issue within a custom kernel. TPUs use semaphores to manage dependencies (e.g., ensuring a DMA is complete before use). This error suggests a signal on a semaphore was not properly waited upon. - **Audit synchronization:** Ensure every `dma_start` has a corresponding `dma_wait`. - **Check semaphores:** Verify that semaphore signals and waits are strictly paired. ### Uncategorized issues If your error log does not match Scenario 1, 2, or 3 (i.e., no "observed errors", no "scheck" tags, and no specific bounds/semaphore messages): - **Action:** This is likely an internal XLA bug. [Please file a bug report](https://github.com/openxla/xla). --- ### Errors Error 1000 (`docs/errors/error_1000.md`) # Error code: E1000 **Category:** Compile Time: HBM OOM This error indicates that the program requires more High Bandwidth Memory (HBM) than is physically available on the TPU device. **Sample error messages:** ``` RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G ``` ``` 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. ``` **XLA backends:** TPU ## Overview XLA performs checks to ensure that the aggregate size of all necessary static allocations fit in the device's HBM. The compiler manages the TPU's fixed HBM capacity for several types of allocations: - **Program inputs and outputs:** Training batches, optimizer states etc. - **TPU temporaries:** Dynamic memory required for intermediate calculations (e.g. activations, gradients, etc). - **Compiled binary:** The machine code for both TensorCore (TC) and SparseCore (SC). - **System overhead:** Reserved space for the XLA Runtime (e.g. infeed buffers on older TPU generations). - **Constants:** Constant values embedded in the HLO IR are allocated on HBM. - **Compiler internals:** Program level and per-HLO allocations (e.g. routing info for nodes in the mesh). This error occurs when the XLA compiler cannot fit all of the above allocations into the device HBM. ## Debugging Carefully analyze the error message and logs to determine which category of HBM OOM below best describes your error: - **"TC Hbm usage: X, SC Hbm usage Y":** If the error explicitly breaks down HBM usage, the aggregate TensorCore (TC) + SparseCore (SC) usage exceeds the HBM limit. → Jump to [Scenario 1. Balance TC and SC HBM usage](#scenario_1_balance_tc_and_sc_hbm_usage). - **"Ran out of memory in memory space HBM"**: Check the logs for an enumeration of the largest allocations on HBM. - In case one or more unexpectedly large tensors (e.g. > 50% of HBM limit) are present → Jump to [Scenario 2. Out of memory due to unexpectedly large allocations](#scenario_2_out_of_memory_due_to_unexpectedly_large_allocations). - If no unexpectedly large tensors are present in the logs → Jump to [Scenario 3. Out of memory due to aggregate allocations](#scenario_3_out_of_memory_due_to_aggregate_allocations). --- ### Scenario 1. Balance TC and SC HBM usage If the error explicitly breaks down usage, e.g., *"TC Hbm usage: X, SC Hbm usage Y"*, this means the aggregate TensorCore (TC) + SparseCore (SC) usage exceeds the HBM limit. Compare the two values to identify the bottleneck: - **High SparseCore usage** - **Optimize HBM stack usage:** HBM stack memory consumption scales with `feature_width`, `max_unique_nz_per_row` and `logical_replica_count`. You can reduce peak stack usage by tuning the `--xla_sc_num_serialized_tables_to_optimize_hbm` flag which serializes the processing of tables. This comes at the cost of reduced parallelism. - **Check padding overhead:** SparseCore aligns embedding tables to 32B (8 floats). Tables with small feature widths (e.g., < 8 floats) incur significant padding overhead, wasting HBM. - **Reduce heap usage:** High values for `maximum_parallel_iterations` increase the amount of input data prefetched into the HBM heap. Lowering this value can free up significant memory. - **Verify sharding:** Ensure embedding tables are properly mod-sharded across all chips. See [How limits translate to tables](https://openxla.org/xla/sparsecore). - Check out [SC: Performance and memory bottlenecks](https://openxla.org/xla/sparsecore#10_performance_memory_bottlenecks) for more ideas. - **High TensorCore usage** - Proceed to [Scenario 2](#scenario_2_out_of_memory_due_to_unexpectedly_large_allocations). - **Balanced** - If neither is individually excessive but the sum is too high, you are at the chip's capacity. You must try lowering usage of both components. Follow recommendations in all three sections. ### Scenario 2. Out of memory due to unexpectedly large allocations If you observe the error message *"Ran out of memory in memory space HBM"* and one or more unexpectedly large allocations are present in the logs (> 50% of HBM limit), it is almost never a hardware capacity issue. It is typically a configuration error. Check the XLA label (if present) of the large allocations for hints on their JAX source code. - **Remove debugging artifacts** - Using [jax.debug.print()](https://docs.jax.dev/en/latest/_autosummary/jax.debug.print.html) in large-scale runs can force the compiler to materialize the full tensor in HBM to transfer it to the CPU, breaking fusion and increasing peak memory usage. Remove any left-over `jax.debug.print()`s. - **Fix inefficient mesh shapes or sharding** - Incorrect mesh shapes or missing sharding annotations can cause the compiler to default to **replication** - forcing the compiler to try to fit really large tensors on a single chip. - Check the shapes of the large allocations and verify sharding is correctly specified and propagated by XLA. ### Scenario 3. Out of memory due to aggregate allocations If you observe the error message *"Ran out of memory in memory space HBM"* and no unexpectedly large tensors are present in the logs, the program runs out of capacity due to the aggregate sum of allocations exceeding the HBM limit. In this case, it is often helpful to visualize the memory profile to identify the specific buffers contributing to the peak usage. See [Debug OOM errors with XProf](https://openxla.org/xla/oom_debugging) for a step-by-step guide on identifying peak memory contributors. Once you have identified some of the top contributors, use the following steps to optimize the memory footprint. #### Scenario 3.A Adjust configuration You can often resolve OOMs with these configuration adjustments: - **Reduce batch size:** The memory needed for intermediate activations and gradients is directly proportional to the batch size. Reducing the batch size can often help reduce memory usage, although you may need to retune your learning rate, momentum, or optimizer hyperparameters to maintain model stability. - **Donate input buffers:** When JAX executes a computation it uses buffers on the device for all inputs and outputs. If you know that one of the inputs is not needed after the computation, and if it matches the shape and element type of one of the outputs, you can specify that you want the corresponding input buffer to be donated to hold an output. This will reduce the memory required for the execution by the size of the donated buffer. You can achieve this by specifying [donate_argnums](https://docs.jax.dev/en/latest/buffer_donation.html) parameter as an argument when using `jax.jit`. - **Enable mixed precision (bfloat16):** Use bfloat16 or quantization (int8 etc) for the largest tensors in the program if the model architecture and quality requirements allow. Note that this change can affect model behaviour and should be considered carefully. ##### Micro-batching (optional) If reducing the global batch size or increasing the chip count is not viable, and the batch size per chip is not already minimized, you can try a micro-batching strategy: - Split each batch into `n` micro-batches; - For each micro-batch, process the forward and backward pass; - Once this is done, accumulate the gradients and update the weight as a whole. This process reduces the activation memory as we divided each batch into `n` micro-batches, so that the if the original batch had size `M`, the activation memory size becomes `M/n`. **Potential issues:** - This process increases step time as we have multiple forward and backward passes. - If the sizes of the model and micro-batch are too different, you may face convergence issues in your model. #### Scenario 3.B Optimize architecture and sharding If configuration changes are insufficient, the model topology might be too large for the current hardware setup. - **Use newer TPU generations:** Newer TPUs generally offer more HBM per chip; switch to newer TPU generations if available. - **Run on a larger chip topology:** If the model weights are too large for the existing topology, you can try sharding them across more chips. - **Implement advanced sharding techniques:** - Explore more advanced data, tensor, or pipeline parallelism approaches. - Specify [sharding hints](https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#constraining-shardings-of-intermediates-in-jitted-code) for intermediate values and outputs. Note that this may cause an increase in network communication overhead due to splitting tensors across multiple chips. - **Use JAX host offloading:** Host offloading techniques allow the user to offload large tensors to the host CPU memory (e.g. [activation offloading](https://docs.jax.dev/en/latest/notebooks/host-offloading.html#activation-offloading) and [optimizer state offloading](https://docs.jax.dev/en/latest/notebooks/host-offloading.html#optimizer-state-offloading)). Note that host offloading techniques can severely impact performance, since these operations will force the system to constantly move large tensors back and forth between TPU HBM and CPU RAM. #### Scenario 3.C Check tensor padding and alignment Inefficient tensor shapes are a common, silent cause of OOMs on TPUs. To get peak performance on TPU's, XLA pads tensor dimensions—typically to multiples of 128 for the minor-most dimension and 8 for the second-minor. This padding affects both input arrays and intermediate tensors (HLO temporaries), potentially inflating memory usage significantly, especially with small dimension sizes. See [Array Layouts](https://docs.jax.dev/en/latest/pallas/tpu/details.html#array-layouts). - **Audit shapes of large buffers:** (On TPU v5 with default layouts) - Hovering over a buffer in [Xprof Memory Viewer](https://openxla.org/xprof/memory_viewer#memory_viewer_components) brings up the buffer details card which contains buffer details including padding information. - *Example*: A shape of `(129, 1024)` might be padded to `(256, 1024)`, resulting in nearly 50% memory waste. - *Correction:* A shape of `(128, 1024)` requires no padding and incurs 0% memory waste. - **Align dimensions:** Ensure all large tensor dimensions (batch size, embedding dimension, hidden size) are multiples of 128. Note that this change can affect model behaviour and should be considered carefully. #### Scenario 3.D Tune key memory impacting XLA flags [Key memory flags](https://openxla.org/xla/flags_guidance#memory_flags) can be tuned to trade-off performance for lower memory usage. However, this strategy should be used as a last resort measure since it can adversely affect performance. #### Scenario 3.E Tune XLA rematerialization pass/manual checkpointing If the model is close to fitting into memory, you can use the [jax.checkpoint](https://docs.jax.dev/en/latest/_autosummary/jax.checkpoint.html) decorator with `jax.grad` to manually control which intermediates are saved on the forward pass versus recomputed on the backward pass. Note that this operation may impact performance, since you are explicitly trading compute cycles for HBM. Check out the JAX documentation for more information: - [Gradient checkpointing with `jax.checkpoint` (`jax.remat`)](https://docs.jax.dev/en/latest/gradient-checkpointing.html) - [Control autodiff’s saved values with `jax.checkpoint` (aka `jax.remat`)](https://docs.jax.dev/en/latest/notebooks/autodiff_remat.html) - [JAX Memories and Host Offloading](https://docs.jax.dev/en/latest/notebooks/host-offloading.html) Alternatively, you can force the `XLA::Rematerialization` pass to prioritize memory savings, potentially at the cost of slower compilations: | Flag | Description | Impact / Trade-off | | --- | --- | --- | | `--xla_tpu_max_hbm_size_mib` | Manually sets the limit on HBM size used by the Rematerialization pass. | Forces the compiler to work harder to fit the program into a limit smaller than the actual physical HBM. | | `--xla_tpu_rematerialization_algo=PEAK_PRIORITY` | Focuses efforts at the points of peak memory usage. | Can be more efficient for aggressive memory reduction than the default algorithm. | | `--xla_tpu_rematerialization_max_block_size_limit=32` | Controls the maximum number of instructions in a block that can be rematerialized at once. | Increasing this allows for memory savings at the cost of **significantly increases compile time**. | | `--xla_tpu_rematerialization_block_effort_factor=10.0` | Defines the amount of effort (compile time) spent searching for blocks to rematerialize. | Higher values allow a more exhaustive search for memory savings at the cost of **increased compile times**. | | `--xla_tpu_pre_fusion_remat=true` | Enables an additional Rematerialization pass *before* the fusion pass. | Can find more memory savings, but increases compile times and may **potentially impact numerical stability**. | Note that making changes to XLA flags should be used as a last resort measure, since it can adversely affect performance. #### Scenario 3.F Use advanced profiling tools [Debug OOM errors with XProf](https://openxla.org/xla/oom_debugging.md) provides a tutorial on using the [XProf Memory Viewer](https://openxla.org/xprof/memory_viewer) to visualize the compiler's view of HBM usage. This tool allows you to see peak memory allocation and buffer lifetimes, which is crucial for understanding exactly what consumes HBM at the point of peak utilization. For general profiling setup, see [Getting started with Xprof](https://openxla.org/xprof#getting_started) and [TensorBoard Profiling](https://docs.jax.dev/en/latest/profiling.html#xprof-tensorboard-profiling). ## Summary table The following table contains a summary of potential interventions to solve OOM errors and information that will help you decide what to do. Intervention | Safe to do? (Will it change the behavior of the program?) | Potential gains | Telltale signs (is this actually the bottleneck that you're experiencing?) ----------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- Using advanced sharding techniques | **Yes.** It almost never changes the numerical correctness of the experiment, although it can cause network communication overhead due to splitting tensors across multiple chips. | **Massive gains** (up to a 256x reduction) | Unexpectedly large individual allocations in the memory viewer (e.g., a single tensor replicated across all TPUs that is 256x bigger than the others). Active arrays showing as un-sharded in TensorBoard hooks. Reducing batch size | **No.** It changes the training dynamics and usually requires retuning the learning rate. (Note: **Microbatching** is a safe alternative that reduces memory without changing behavior). | **Massive gains** (can save a factor of thousands). | "Temporaries" failing to allocate during gradient calculations. Seeing "JVP" in the operation name, and encountering many batch-size-shaped tensors in the memory profile. Enabling mixed Precision (e.g., Bfloat16) | **Risky.** It alters numerical precision, which can change experiment results or cause the model to fail to converge entirely. | **Moderate gains** (typically a factor of 2x, as it halves memory usage). | The memory viewer confirms that the largest tensors are currently utilizing 32-bit floats (`float32`). Manual checkpointing (`jax.checkpoint`) | **Yes.** It does not alter behavior; it merely trades computation time (flops) to save memory by recomputing tensors instead of storing them. | **Large gains** (e.g., can result in only half of the activations needing to exist in memory at the same time). | Multiple tensors of the exact same size filling up memory during a backward pass. Often accompanied by "JVP" in the operation name. Donating input buffers (`donate_argnums`) | **Yes.** Maintains experiment integrity. If applied incorrectly, it will not corrupt data but will simply throw a clear error message. | **Marginal gains** (~1% memory savings). | There is no specific telltale sign, but it is considered a "free win" that is always worth trying. Changing model dimensions | **No.** Directly alters the model's behavior. Modifying input or output dimensions can completely break compatibility with the dataset. | **Variable gains** depending on how drastically hidden dimensions or layers are reduced. | The Xprof memory viewer shows a large amount of memory wasted on "padding" because array dimensions are not powers of two or multiples of 128 (e.g., a dimension of 2050 instead of 2048). Host offloading (CPU) | **Yes (numerically)**, but **No (performance)**. While mathematically safe, it is considered a "foot gun" that may cause severe speed bottlenecks due to data transfer between the CPU and TPU. | **Marginal gains** (the CPU only has about 3x the memory of the TPU). | Typically a last resort for massive optimizer states or for memory-heavy data preparation/pre-processing steps. --- ### Errors Error 1001 (`docs/errors/error_1001.md`) # Error code: E1001 **Category:** Compile Time: Scoped Vmem OOM This error indicates that the program requires more Scoped Vector Memory (Vmem) than what was allocated. **Sample Error Messages:** ``` RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ... ``` **XLA Backends:** TPU ## Overview TPUs have Vector Memory (VMEM) which is a local scratchpad memory used exclusively by the TensorCore (TC). The compiler manages Vmem for different types of allocations: * **Instruction-scoped allocations:** Temporary storage in Vmem while executing a single HLO instruction. This includes operand span buffer (e.g. for double buffering) and register spills. * **Program-scoped allocations:** Allocations that live beyond the scope of a single HLO instruction. These are usually HLO temporaries and intermediate results that are inputs and/or outputs of HLO instructions. A Compile Time Scoped Vmem OOM occurs when the instruction-scoped allocations exceed the allocation limit for that instruction. This limit is controlled - globally for the entire program via the flag [--xla_tpu_scoped_vmem_limit_kib](https://openxla.org/xla/flags_guidance) and - per custom kernel via [vmem_limit_bytes param](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.tpu.CompilerParams.html#jax.experimental.pallas.tpu.CompilerParams.vmem_limit_bytes). These errors are typically caused by an internal compiler bug or by a custom kernel exceeding its allocation limit. ## Usable Limits {#usable-limits} Different TPU hardware have different Vmem sizes. In addition, different hardware may reserve different amounts of Vmem for various activities, thereby limiting the amount that can be allocated to Scoped Vmem. Below are the typical limits on various hardware. | Hardware | Limit (human-friendly) | Limit (KiB) | Limit (bytes) | | :--------- | ---------------------: | ----------: | ------------: | | TPU v2 | 16MB | 16384 | 16777216 | | TPU v3 | 16MB | 16384 | 16777216 | | TPU v4 | 16MB | 16384 | 16777216 | | TPU v4i | 16MB | 16384 | 16777216 | | TPU v5e | 128MB | 131072 | 134217728 | | TPU v5p | 63.9MB | 65472 | 67043328 | | TPU v6e | 127.9MB | 131008 | 134152192 | | TPU 7x | 63.9MB | 65472 | 67043328 | ## Debugging Carefully analyze the error message to identify if the error stems from a custom kernel or a standard HLO. An error due to a custom kernel should have the following signature: ``` Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = custom-call(), custom_call_target="tpu_custom_call" ... ``` - **Custom kernel scoped Vmem OOM**: If the error points to a custom kernel → Jump to [Retune the Kernel](#retune_the_kernel). - **Non-kernel Vmem issues**: If the Vmem OOM occurs due to a non-custom-kernel op, it is likely an internal compiler bug. [Please file a bug report on XLA](https://github.com/openxla/xla) with an HLO dump. ### Retune the kernel If the error originates from a custom kernel, use the following techniques to reduce the kernel's memory requirement: * **Adjust Block Sizes:** Reduce the block sizes (tile sizes) in your kernel configuration, to lower Scoped Vmem usage. * **Set Per-Kernel Scoped Vmem Limits:** Explicitly request the required amount of memory for that specific kernel using the [vmem_limit_bytes param](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.tpu.CompilerParams.html#jax.experimental.pallas.tpu.CompilerParams.vmem_limit_bytes) * **Modify Memory Coloring:** Explicitly color/constrain the kernel's inputs/outputs to Vmem using [pallas.tpu.with_memory_space_constraint](https://docs.jax.dev/en/latest/_autosummary/jax.experimental.pallas.tpu.with_memory_space_constraint.html). Be careful not to color too many inputs outputs to Vmem, as that might cause an overall Vmem OOM. * **Adjust Vmem limit:** If kernel specific retuning is difficult or the issue affects many kernels, you can adjust the global Vmem limit using the flag [--xla_tpu_scoped_vmem_limit_kib](https://openxla.org/xla/flags_guidance). --- ### Errors Error 1200 (`docs/errors/error_1200.md`) # Error code: E1200 **Category:** Compile Time: Host Offload Output Mismatch This error occurs when a tensor explicitly offloaded to host memory is returned as a program output, but the program's output signature is not configured to expect host memory. **Sample Error Messages:** ``` INVALID_ARGUMENT: Tensor which is moved to host (starting from tuple.64) is returned from the entry computation but the layout for this output is not set to host memory. ``` **XLA Backends:** TPU, GPU ## Overview When the compiler encounters an annotation to offload a tensor to the host (CPU), it tracks that tensor's location through the computation graph until one of three events occurs: 1. **Move to Device:** A matching annotation moves the tensor back to the accelerator. 2. **Host Computation:** The tensor is consumed by a host-side operation. 3. **Program End:** The tensor reaches the end of the program and becomes an output. This error is triggered in scenario **#3**. The tensor is physically located in host memory at the end of the execution, but the XLA program's entry computation signature defines that specific output as residing in **Device Memory**. Because the compiler cannot implicitly change the entry computation's interface, it raises an error. ## Debugging To resolve this error, determine whether you intended for this tensor to be an output on the Host or if it should have been moved back to the Device before returning. * **Intended to return on Host:** If you explicitly want this tensor to be returned in host memory (avoiding a transfer back to device), you should explicitly set the output memory space of the entry computation to **Host Memory** for this specific output. * **Intended to return on Device:** If the tensor was meant to stay on the device or return to it before the program ends, you likely missed an annotation. Insert a matching annotation to move the tensor back to the device. If the source of the offloaded tensor is unclear, or you cannot find where the "move to device" annotation is missing, use XLA logging to trace the instructions. * **Enable logging:** If you are on Google Cloud TPU, rerun your program with the following flag: `--vmodule=host_offloader=1`. * **Analyze Logs:** Look for the "trace" output in the logs. This will show the path of the tensor starting from the offload instruction. Use this to pinpoint exactly where the tensor reaches the program boundary without being moved back to the device. --- ### Errors Error 2001 (`docs/errors/error_2001.md`) # Error code: E2001 **Category:** Compile Time: Unsupported RHS DataType on Hardware This error occurs when the data type used for the **Right-Hand Side (RHS)** operand in a matrix multiplication (e.g., `jax.lax.dot_general`, `jax.lax.conv`, `jax.numpy.matmul`, or the `@` operator) is not natively supported by the specific TPU generation being used. **Sample error messages:** ``` INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>' ... The MLIR operation involved: %13440 = "tpu.matmul"(%13435, %13437, %13439) ``` **XLA backends:** TPU ## Overview The TPU's Matrix Multiply Unit (MXU) natively supports `Float32` operations on all hardware generations. However, native support for `BFloat16` and other quantized data types (e.g., Int4, Int8, or Float8) varies by hardware generation. This error is triggered when your kernel attempts to map a matrix multiplication to the MXU using a data type that your specific TPU generation does not have the physical circuitry to execute. This error typically indicates that the compiler's **Canonicalization** pass— which attempts to automatically convert unsupported types into supported ones (e.g., via software emulation)—was unable to find a valid conversion rule or was prevented from doing so because **Compatibility Mode** was disabled. ## Debugging To resolve this error, you must align your data types with the capabilities of your hardware. You have following options: ### 1. Cast to native types The most reliable fix is to manually cast your operands to a hardware-supported datatype (like `Float32` or `BFloat16` on TPU v4+) inside your kernel before the matmul operation. * **Why:** `Float32` is the universal data type supported natively by the MXU on all TPU generations. * **Trade-off:** This comes with a VPU (Vector Processing Unit) cost - the cycles required to perform the cast, but it guarantees your kernel will run on the current hardware. ### 2. Check Compatibility Mode Typically the compiler can automatically handle these type mismatch issues in **Compatibility Mode** which is enabled by default. Double check XLA configs to make sure `--xla_mosaic_compat_mode` is not set to false. This acts as a "polyfill," injecting software emulation sequences for operations your hardware does not natively support. **What Compatibility Mode enables:** * **Mixed-precision MatMuls:** Allows mixing Integer operands with Float accumulators by automatically inserting cast operations (e.g., extending integers to `Float32` before the matmul). * **Low-precision emulation:** On certain hardware generations, emulates unsupported types like `4-bit` floating point (`4E2M1FN`) or `8-bit` floating point (`8E4M3FN`) by extending them to supported types like `BFloat16` or `Float32` before execution. Note that this mode prioritizes compatibility over peak performance as emulation requires additional instructions to convert data formats before the MXU can operate on them. ### 3. Upgrade hardware or request support If your algorithm strictly requires native performance for types like `Int4` or `Float8` without the overhead of casting or emulation, you will need to run on a newer TPU generation with native support. **Feature request:** If you believe your hardware supports this operation, or if the compiler is missing a valid emulation path even in Compatibility Mode, please file a feature request. We usually guarantee that operations are forward compatible. So if your kernel runs on a TPU generation then it should run on all future generations, but it is not guaranteed to have emulation for older generations (for some of which the casts would be very expensive). --- ### Errors Error 2002 (`docs/errors/error_2002.md`) # Error code: E2002 **Category:** Compile Time: Mosaic Input/Output Misaligned Block and Tiling This error occurs when the block shape of a kernel input or output does not align with the default tiling of the datatype on the specific TPU hardware being used. **Sample error messages:** ``` UNIMPLEMENTED: Mosaic failed to compile TPU kernel: Failed to set window params for input 0: Operand of shape (..., 256, 8192) has tiling (16, 128), but its block shape (..., 8, 8192) is not divisible by tiling evenly nor matches the full shape. ``` **XLA backends:** TPU ## Overview Tensor cores (TC) in TPUs have two-dimensional vector registers. The two dimensions are called **sublane** and **lane**. Since TPU compute units (e.g., MXU) operate at the granularity of vector registers, XLA arrays are laid out in TPU memory in tiles. This tiling (e.g., `8x128`) minimizes data transformations when feeding the compute units. The exact tiling dimensions (sublanes × lanes) depend on the hardware generation and the data type. For example, a common tiling for most types is **8×128**. At compile time, XLA enforces the following constraints for the minor and 2nd-minor dimensions of each kernel input/output: 1. **Divisibility:** The block dimension must be a multiple of the tile dimension in the underlying tensor, or 2. **Full shape exception:** If the block dimension is not divisible, it must be equal to the **full size** of that dimension in the underlying tensor. This error is triggered when a block violates both conditions. For example, loading a block of shape `(8, 100)` from a input of shape `(8, 1024)` on hardware with shape `(8, 128)` tiling fails because `100` is not divisible by `128` and `100 != 1024`. But, it would be allowed if the input shape was `(32, 100)`. ## Debugging To resolve this error, ensure your kernel's block shapes align with the current hardware tiling. Modify your kernel code to align the block size such that it is a multiple of the required tiling. For example, if the error states the tiling is `(16, 128)` but your block shape is `(8, 128)`, change the block spec such that the shape matches `(16, 128)`. --- ### Errors Error 2003 (`docs/errors/error_2003.md`) # Error code: E2003 **Category:** Compile Time: Mosaic Unproven Memory Access Alignment This error occurs when the compiler analyzes a memory access operation (such as `vector.load`, `vector.store`, `tpu.load`, or `tpu.store`) and cannot statically prove that the dynamic index used for a specific dimension is a multiple of the required **tiling size**. **Sample error messages:** ``` INTERNAL: Mosaic failed to compile TPU kernel: cannot statically prove that index in dimension 1 is a multiple of 128 at location: ... The MLIR operation involved: %14372 = "vector.load"(%14371, %93, %14363) : (memref<4x256xf32, #tpu.memory_space>, index, index) -> vector<1x32xf32> ``` **XLA backends:** TPU ## Overview When your kernel loads or stores a vector, the memory address (calculated from the base pointer plus the dynamic index) must align with the vector's **tiling size** on the hardware. For example, if a dimension is tiled by 128 elements, the dynamic index used to access it must be `0`, `128`, `256`, etc. Note that many operations (like vector loads and stores) have no such requirements for static indices. The compiler enforces this requirement using **static analysis**. It traces the history of the index variable back through the arithmetic operations that produced it (e.g., multiplications, additions). If the compiler cannot guarantee (at compile time) that the resulting value will always be divisible by the tiling size, it raises this error. The compiler treats "proven misalignment" and "unknown alignment" identically. So if you use an index that is mathematically guaranteed to be misaligned (e.g., `i * 128 + 32`), the compiler will raise the same error. So this error can occur when 1. You use a runtime variable (dynamic index) to access memory. 2. The index calculation logic is too complex for the compiler to analyze. 3. The index is mathematically valid but lacks an explicit proof in the code. 4. Static analysis determines "proven misalignment". ## Debugging To resolve this error you have the following options: ### 1. Assert alignment explicitly If you know your index is valid but the compiler cannot prove it, use the `tpu.assume_multiple` operation. This acts as a promise to the compiler that a value is divisible by a specific factor. ### 2. Use aligned loads and rotate In scenarios where the misalignment is intentional, instead of loading a small, unaligned vector segment: * Load a larger, fully aligned tile and then rotate the values by a dynamic amount to shift the desired data into position (since vector slices with dynamic start indices are not supported). or * Reshape or pad the tensor so that the data starts at index 0 and the stride between accesses matches the hardware alignment. * Example: If you iterate over chunks of size 32 starting at offset 1, your offsets are 1, 33, 65... (unaligned). * Fix: Re-pack the data into a new tensor where the first chunk is at 0 and the dimension is padded to 128. Your offsets become 0, 128, 256..., which satisfies the alignment requirement. These methods consume more memory but often simplify kernel logic and eliminate the need for manual alignment assertions. --- ### Errors Error 3000 (`docs/errors/error_3000.md`) # Error code: E3000 **Category:** Compile Time: SparseCore Allocation Failure This error occurs when the `XLA:SparseCore` compiler is unable to allocate a contiguous block of memory in the specified memory space required by the current SparseCore program. **Sample error messages:** ``` INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array}> : () -> memref<23440xf32, 201> ``` **XLA backends:** TPU ## Overview The SparseCore (SC) is a specialized processor for sparse workloads. It relies on specific memory hierarchies to manage data movement efficiently. The XLA compiler attempts to statically size and allocate buffers based on hardware limits and user-defined shapes. This error indicates an **Out of Memory (OOM)** condition during this allocation phase. The error message typically specifies a memory space ID. Below are the common memory spaces and their integer encodings: | Memory Space ID | Name | Description | | --- | --- | --- | | **0** | **Smem** | Local Scalar Memory. Used for scalar registers and control flow. | | **201** | **TileSpmem** | Tile-specific Scratchpad Memory. Fast, local SRAM available to a specific SC tile. | | **202** | **Spmem** | Shared Scratchpad Memory. Used to opportunistically stage data (inputs, outputs, intermediates) to hide HBM latency. | | **203** | **HBM** | High Bandwidth Memory. Large, shared memory used for embedding tables, heaps, and stacks. | | **204** | **Sync Flags** | Synchronization primitives used for coordination. | For a deep dive into SC and its memory hierarchy, refer to the [SparseCore documentation](https://openxla.org/xla/sparsecore). ## Debugging The resolution depends on which memory space failed the allocation. ### Scenario 1. HBM allocation failures (Memory Space ID: 203) This error occurs if a single temporary allocation requested by the SparseCore program is too large to fit in the available HBM. In standard embedding workloads and SC offloaded collectives extremely large per-core partitions or incorrect sharding specifications can force the compiler to request massive buffers. **Recommended actions:** - **Check sharding:** Ensure your embedding tables and SC input/output tensors are partitioned/sharded correctly. If a single core is responsible for too much data, the allocation might fail. - **Adjust limits:** Review `max_ids_per_partition` and `max_unique_ids_per_partition`. If these are set unnecessarily high, the compiler reserves more memory than needed. Refer to [How limits translate to tables](https://openxla.org/xla/sparsecore#7_how_limits_translate_to_tables_on_sparsecore). ### Scenario 2. Internal memory failures (Memory Space IDs: 0, 201, 202, 204) Allocation failures in **Smem**, **TileSpmem**, **Spmem**, or **Sync Flags** typically occur due to compiler bugs or limitations in the allocation strategy, where the compiler fails to account for all memory requirements. **Recommended actions:** 1. **Isolate the failing XLA operation:** To identify the specific SC HLO or Mosaic kernel causing the failure, generate the intermediate compiler representations: - **Dump SparseCore MLIR:** Set the flag `--xla_sc_dump_mlir_to=/path/to/dump`. This generates the MLIR of the SparseCore program, allowing you to see which allocation size matches the error message. - **Dump Mosaic LLO:** For custom kernels, use `--xla_mosaic_dump_to=/path/to/dump` to inspect all Low Level Optimizer (LLO) programs emitted by Mosaic. 2. **Reduce scratch sizes (Pallas users):** If the failure occurs within a Mosaic kernel, review your `scratch_shapes` configuration. Ensure that your `pltpu.SMEM` requests fit within the hardware specifications for your specific TPU generation. 3. **Disable collective offload:** If the error arises from a SC offloaded collective operations, try disabling the SC offloading features: - `--xla_tpu_enable_sparse_core_collective_offload_all_gather=false` - `--xla_tpu_enable_sparse_core_collective_offload_all_reduce=false` 4. **File a bug:** If the above steps do not resolve the issue, it is likely a compiler bug. [Please file a bug report](https://github.com/openxla/xla). --- ### Errors Error 3001 (`docs/errors/error_3001.md`) # Error code: E3001 **Category:** CompileTime: SparseCore No Viable Logical Replica Count This error occurs when the `XLA:SparseCore` compiler fails to determine a valid logical replica count configuration that allows the workload to fit within the SparseCore's local scratchpad memory (Tilespmem). **Sample error messages:** ``` XLA:TPU compile permanent error. Compilation failure: No viable logical replica count for the embedding table with metadata: max_nz_per_row = 141352, max_unique_nz_per_row = 8, feature_width = 8, sample_count = 204800 (last tried split factor for vector splitting = 1, last tried split factor for sample dimension splitting = 1, fixed_size_allocation_bytes = 410880, row_dependent_size_allocation_bytes = 1696224, total_spmem_size_bytes = 524288) ... ``` **XLA backends:** TPU ## Overview This error is specific to **SparseCore** use cases, particularly for Large Embedding Models (LEMs). The **logical replica count** is an internal compiler parameter that determines how input batches are partitioned to manage scratchpad allocation pressure. The compiler attempts to split the workload into smaller chunks (replicas) so that the intermediate buffers required for each chunk fit into the SparseCore's limited **Scratchpad Memory**. Generally, a higher logical replica count reduces allocation pressure by processing smaller batches of data at a time. This error indicates that even after attempting various splitting configurations, the compiler could not find a setup where the required buffers fit in the Tilespmem memory. The allocation size is determined by a combination of: - **`sample_count`**: The number of embedding lookup IDs assigned to each SparseCore (derived from batch size). - **`feature_width`**: The size of the embedding dimension. - **`max_nz_per_row`**: The maximum number of non-unique embedding lookup IDs across all SparseCores. - **`max_unique_nz_per_row`**: The maximum number of unique embedding lookup IDs. ## Debugging To resolve this error, you need to reduce the memory pressure on the SparseCore scratchpad. ### 1. Improve metadata estimations The compiler allocates memory based on `max_nz_per_row` and `max_unique_nz_per_row`. If these values are estimated conservatively (i.e., set much higher than the actual data requires), the compiler will reserve unnecessary space, causing this error. Ensure these parameters accurately reflect the actual ID distribution of your dataset. You can consider applying **Feedback-Directed Optimization (FDO)** to determine optimal values for these parameters. ### 2. Reduce batch size The `sample_count` is directly derived from your global batch size. Reducing the batch size decreases the amount of data each SparseCore must process per step, thereby reducing the size of the required scratchpad buffers. --- ### Images Hlo To Thunks.Mermaid (`docs/images/hlo_to_thunks.mermaid.txt`) flowchart TD InitialHlo("HLO before optimizations") OptimizedSheduledHlo("Optimized and scheduled HLO") BufferAssignment("BufferAssignment") ThunkSequence("ThunkSequence") LlvmIr("LLVM IR") MachineCode("Machine Code
PTX, obj, ..") Executable("Executable") %% Stages (Processes - Subroutine Shape) HloPasses[[HLO Optimization]] Scheduling[[Scheduling]] BufferAssigner[[Buffer Assignment]] ThunkEmitter[[Thunk Emission]] LlvmCompiler[[LLVM Compilation]] ExecutableBuilder[[Build Executable]] %% Edges (Flow of Artifacts through Stages) InitialHlo --> HloPasses HloPasses --> Scheduling Scheduling --> OptimizedSheduledHlo OptimizedSheduledHlo --> BufferAssigner BufferAssigner --> BufferAssignment OptimizedSheduledHlo --> ThunkEmitter BufferAssignment --> ThunkEmitter ThunkEmitter --> ThunkSequence ThunkEmitter --> LlvmIr --> LlvmCompiler LlvmCompiler --> MachineCode MachineCode --> ExecutableBuilder ThunkSequence --> ExecutableBuilder BufferAssignment --> ExecutableBuilder ExecutableBuilder --> Executable --- ### Megascale Debugging Workflow (`docs/megascale/debugging_workflow.md`) # Debugging Workflow This document describe a general workflow for debugging MXLA issues. ## Prerequisite 1. Use JAX 0.6 or up, and enable JAX distributed service. This version of JAX contains additional logging that can help identify which workers are experiencing issues. 2. (Optional) Generate an HLO dump using the --xla_dump_to flag when initializing your workload. This is discussed in the [XLA documentation](https://openxla.org/xla/hlo_dumps). 3. (Optional) Set --vmodule=real_program_continuator=1 to enable verbose logging for the TPU program execution status. ## Flow chart The flowchart below illustrates the debugging process. To access detailed playbooks for each step, click on the corresponding item in the chart. ## Hangs ### Locate the Megascale Hang Detected Error If you see the following error message in your TPU worker logs, this means that MXLA timed out after detecting no progress: ```text Megascale hang detected: Timed out waiting for 4 graphs to complete at launch_id 13650. Already completed: 100. StepGloballyInProgress: true. Timeout: 1m ``` 1. Workers will report errors to a coordinator for processing. * For **Pathways** jobs: the digest can be found in the logs of `resource_manager` job. * For **McJAX** jobs: the logs can be found on MXLA Coordinator. This is typically task 0 of slice 0. 2. Check logs around time of error detection, and look for `Megascale detects a hang`. 3. Follow steps below to diagnose the issue based on the identified cause. ### Diagnosis #### Interpreting TPU States Before diagnosing MXLA hangs, it is important to understand the TPU states report format. Below is a sample report: ```text Full error digest: Potential cause: Potential culprit workers: First error timestamp: First error type: TPU states: Launch ID: Module: jit.step_fn Fingerprint: Sample worker: @:: Tag:PC breakdown: @(HLO): [:@:, ...] ... ``` #### Bad TPU Chip (tensor core or sparse core) ```text Megascale detects a hang that is likely caused by bad TPU chips on the following hosts. Please remove the hosts from the fleet and restart the workload. If problem persists please contact Megascale XLA team. The host that have bad TPUs are: Full error digest: Potential cause: Bad TPU chips Potential culprit workers: /: ``` This error means that the issue is potentially caused by a faulty TPU chip. The error message should include the job information and host name of the faulty chip. In the example above, the faulty chip is on host ``, affecting task `` of the job ``. You can configure your job to avoid that host. **Note:** There are some cases that the hang was caused by a XLA or custom pallas kernel bug, but if you see the same host appearing multiple times (for example more than 3 times) in multiple hang events, the TPU on that host is very likely faulty. #### Networking issue ```text Megascale detects a hang that is likely caused by a networking issue. Please examine the underlying networking stack for the following hosts. The hosts are: Full error digest: Potential cause: Networking issue Potential culprit workers: , ``` This error indicates that your job has encountered a failed network link. The error message should include a single or pair of job name, task id, host name of the faulty network link. In the example above, the faulty network link is between host `host_name_1` and `host_name_2`. Sometimes RapidEye can further localize the faulty host if a single host appears in multiple broken network links. #### Different modules ```text Megascale detects a hang that is likely caused by running different modules on different devices. Please confirm that all workers is running the exact same program. It can also be caused by a hang in a subset of devices and the unaffected devices have moved on to the next program. Please inspect the digest below to further root cause the hang. Example hosts that have different HLO modules: Full error digest: Potential cause: Different module Potential culprit workers: TPU stats: : TPU states: Module: jit_loss_and_grad Fingerprint: Launch ID: 193 :(): Module: jit_optimizer_apply Fingerprint: Launch ID: 0 :(): ``` This error may indicate that a hang has occurred in a subset of workers, causing those workers to be stuck at the current module while unaffected workers advance to the next module. To identify the root cause, inspect the digest printed by RapidEye in the logs. The `TPU states` section of the logs shows which modules are running on which workers. In the example above, workers are running different modules: `jit_loss_and_grad` and `jit_optimizer_apply`. #### Fingerprint mismatch for HLO module ```text Megascale detects a hang that is likely caused by inconsistent HLO module compilation across workers. This is likely a bug in JAX tracing or XLA compiler. Please inspect the HLO dumps to confirm the root cause. Example hosts that have different HLO fingerprints: Full error digest: Potential cause: Fingerprint mismatch Potential culprit workers: TPU stats: Module: reduce.31 Fingerprint: Launch ID: 37 :(): Module: reduce.31 Fingerprint: Launch ID: 40 :(): ``` This log message indicates the hang was likely caused by inconsistent HLO module compilation across workers, possibly due to an issue with JAX tracing or the XLA compiler. If you see this log, follow [these steps](https://openxla.org/xla/hlo_dumps) to collect HLO dumps from the culprit workers for further debugging. #### Data input stall **Note:** This error is not yet implemented. ```text Megascale detects a hang that is likely caused by data input stall on the following hosts. Please check the workers to make sure the data input pipeline is working properly. The host that have data input stalls are: ``` This error means that all devices launched the same program, but that input data was not provided to the program before the system timed out. To fix this issue, confirm that: 1. The identified hosts can access the input datasource. 2. The identified hosts are properly loading/parsing the input datasource. 3. Confirm identified hosts are not throttled on reads to the input datasource. #### Unrecoverable error ```text Some workers have halted with an unrecoverable error: : {some error} Please inspect the error log of these workers: ``` This error means that there was an issue that prevented the program from properly executing and could not be recovered automatically. This error was unable to be specifically categorized. Further information can be obtained from checking the logs of the worker(s) mentioned in the error report. If the error appears to be specific to the given machine (ex. failure to copy data from TPU to host), then you can configure your job to avoid those hosts. #### Unknown Error ```text Megascale detects a hang but cannot determine the root cause. Please inspect the full digest below. ``` This error means that there was an issue that prevented the program from properly executing and could not be recovered automatically. This error was unable to be specifically categorized and there is no further error information available. If RapidEye cannot identify the culprit, follow the [Manual Diagnosis ](#manual-diagnosis) instructions below. ### Manual Diagnosis Always inspect the [RapidEye diagnosis](#diagnosis) first when investigating a hang. Only use the manual diagnostic flags below as a fallback when RapidEye cannot identify the culprit—such as unknown MXLA hangs, single-slice hangs, or exceptionally slow step times. #### Diagnostic Flags Reference When diagnosing stalls, timeouts, or program stragglers that RapidEye cannot automatically classify, the following flags enable detailed low-level execution logs and prevent premature job crashes: ```bash # Enable low-level logging for TensorCore and SparseCore execution states --xla_tpu_enable_log_recorder=true --xla_tpu_enable_sc_log_recorder=true # Extend synchronization wait timeouts to prevent false-positive stall # detections --xla_tpu_debug_sflag_wait_timeout_ms=150000 --xla_tpu_debug_sc_sflag_wait_timeout_ms=150000 # Prevent immediate hard crash upon timeout detection so logs can be gathered --xla_tpu_debug_sflag_wait_shalt_on_detection=false # Enable hierarchical tracking of HLO progress --xla_tpu_enable_progress_tracker=16 # Prevent the coordinator from aborting all workers prematurely on hang --megascale_error_reporter_abort_on_hang=false ``` #### Purpose and Usage - **Preventing Premature Aborts**: By default, when a hang is detected, the coordinator immediately aborts all workers (`megascale_error_reporter_abort_on_hang=true`). Setting this flag to `false` keeps workers alive long enough to flush thread stack traces and detailed TPU states to Cloud Logging. - **Investigating Slow Steps vs. Real Hangs**: For exceptionally large models or initialization steps that take longer than the default timeout (typically 60 seconds), increasing `--xla_tpu_debug_sflag_wait_timeout_ms=150000` (150s) rules out false positives. - **Inspecting Timeout Logs**: In Cloud Logging, search for: ```text Wait timeout on sflag ``` Group matching logs by HLO opcode to isolate the first worker or core that failed to make progress, distinguishing the culprit from bystander workers waiting on dependencies. ## Performance For an in-depth reference on configuring and tuning Megascale flags (such as collective buffer sizing and zero-copy memory premapping), see the [Megascale Performance Tuning Flags guide](performance_tuning.md). ### Get an XProf session Follow the instructions in the [XProf documentation](https://openxla.org/xprof/capturing_profiles) to generate an XProf trace for your problematic run. ### Check for shortage of mapped DMA buffers The Megascale XLA runtime needs to register host memory before it can be used for DMAs to and from the TPU. This happens early after the process starts. If you see these registrations (`MapDmaBuffer` calls) at steady-state then it indicates that something is wrong. Look for the presence of these calls in XProf Trace Viewer. See the screenshot below for reference. **Tip:** Search for the exact worker name, because there can be other workers with similar or close names. You also search for the term “MapDmaBuffer” on the page. ![Example xprof trace showing MapDmaBuffer calls](./images/map_dma_buffer_example_trace.png) If the issue is observed then try to increase the size of the premapped memory region by increasing the value of `--megascale_grpc_premap_memory_bytes`, restarting the job, then checking again. Please also check the checkpoint size: the general guidance is that the premap buffer should be at least twice the checkpoint size. ### Check for memory copies during network transfers Megascale XLA network transfers are zero-copy by design. However, there are cases where memory copies will occur and cause degraded performance. Look for memory copies in Megascale's "Communication Transport" traces as shown in the example screenshot below. ![Example "Megascale: Memory Copy" trace during network receive](./images/memory_copy_example_trace.png) If the issue is observed then try to increase the size of the premapped memory region by increasing the value of `--megascale_grpc_premap_memory_bytes`, restarting the job, then checking again. Please also check the checkpoint size: the general guidance is that the premap buffer should be at least twice the checkpoint size. ### Network Analysis Megascale also provides a Colab [notebook](https://github.com/openxla/xla/blob/main/xla/megascale/tools/network_analysis_oss.ipynb) to help analyze network performance using an XProf trace. This tool can be used to do the following: * Examine transfer latencies to identify potential network slowdowns or host slowness. * Examine transfer sizes to identify if your workload is optimized to use a smaller number of larger transfers as opposed to a large number of small transfers. * Determine if your workload is unoptimized and produces bursty collectives/always has a large number of pending collectives. * Visualize the network throughput timeline to see if the workload is network bound. * Examine {source, destination} pairs to identify possible bad hardware on individual hosts. #### Collective Slack Too Small One indicator that your workload is not optimized for compute/communication overlap is seeing small slack times for a subset of collectives. This can manifest as longer than expected `recv-done` traces in the trace viewer, or as collectives with zero or near-zero [slack time](https://cloud.google.com/tpu/docs/troubleshooting/troubleshoot-multislice#slack_time). If this is the case, look towards identifying bottlenecks in your workload that may be causing parts of your program to not overlap compute and network communication. #### High Network Bandwidth Demand If you are observing long `recv-done` op latencies within your model XProf, this could be an indication that the model is 'Bandwidth Bound' in those portions of the step function (is blocked by available network bandwidth in the system). You can generate a timeline of network usage for your model. If you see consistently high network usage throughout the step, or specific regions with large spikes, then your model may be bandwidth bound in those regions. Use the [Network Analysis](#network-analysis) to generate a timeline of network usage: To mitigate bandwidth bound models, you can: 1. Check the [Collective Slack](#collective-slack-too-small) of your model. Models with many collectives with low slack will have bandwidth bound regions. 2. Confirm that the network settings are optimized. 3. Examine your model structure and data sharding to see if there are ways to increase computation/communication overlap. 4. (Data parallel models) Confirm that you have sufficient batch size in each local replica to overlap with the communication. #### High Network Latency If the bandwidth is not saturated, you may want to generate the RPC latency timeline. If you see high RPC latencies is constantly or sporadically high, this means that there are some issues with the MXLA RPCs. Use the [Network Analysis](#network-analysis) to generate the RPC latency timeline, the following example shows that there is some sporadic 200ms long tail RPC latencies. #### Confirm Optimal Network Settings High RPC latencies on Cloud environment are often caused by suboptimal TCP configuration. Please confirm if all TCP parameters are configured properly within the container. If any of the TCP parameters are not correctly configured, consult Google Cloud ML Compute Services (CMCS) team on how to configure them properly. #### HLO Dump Please follow [these steps](https://openxla.org/xla/hlo_dumps) to dump HLO to the local filesystem on the TPU worker. You may need to upload the dump to GCS in order to share them with the XLA or Megascale team. ## Straggler **Note on Future Tooling:** Google is actively working on open-sourcing versions of diagnostic dashboards to provide a more streamlined experience for Cloud TPU customers to identify and diagnose stragglers. These will be available soon. Stragglers often manifest as outliers in key Megascale metrics. These key metrics include: - **kubernetes.io/container/multislice/network/dcn_transfer_latencies**: Measures the round trip time for a transfer. Measured at the sender. - **kubernetes.io/container/multislice/network/dcn_inbound_transfer_latencies**: Measures the one-way latency for a transfer. - **kubernetes.io/container/multislice/accelerator/compute_latencies**: Measures the time it takes to compute a reduction. Outliers are correlated with CPU and memory issues. These metrics are most useful when grouped by hosts as shown in the [dashboard examples](#using-the-dashboards) below. Once a host is identified, it can be diagnosed further for software or hardware issues. ### Using the Dashboards #### DCN (Data Center Network) Transfer Latencies The figure shows an example of an outlier in the network transfer latency metric. In this example, `test5-slice-job-1-0` is an outlier. #### DCN (Data Center Network) Inbound Transfer Latencies This figure shows an example of an outlier in the inbound transfer latency metric. In this example, `test2-slice-job-1-0` is the outlier. #### Host Compute Latencies The figure shows an example of an outlier in the compute latency metric. In this example, `slice-job-1-0` is the outlier. --- ### Megascale Performance Tuning (`docs/megascale/performance_tuning.md`) # Tuning Megascale XLA Flags MegascaleXLA (MXLA) coordinates communication across TPU slices over Data Center Networks (DCN). This guide provides guidelines and best practices for configuring and tuning MegascaleXLA flags for users running multi-slice TPU workloads. An effort has been made to set reasonable defaults for most cases. ## Collective Buffer Sizes For small collectives, it is generally more efficient to package small TPU send/receives and network send/receives into larger ones to reduce latency overhead. For larger collectives, it may be advantageous to use a greater number of smaller intermediate buffers. A number of tools exist which allow for altering this behavior. TPU DMA reads may be coalesced or split for some collectives by passing, ```bash --megascale_target_dma_size= ``` which will attempt to schedule DMA reads that are as close as possible to the target size. By default, a 8 MB default is used. It's only recommended to adjust this value if a profiling trace shows considerable delay between the time a DMA is initiated and the first network transfer for a particular collective is initiated. For small collectives, it is often beneficial to only use a subset of the available participants in a collective for performing reduction operations. This is known as "sparse reduction". The buffer size threshold for which sparse reduction is enabled is controlled by the flag, ```bash --megascale_sparse_reduction_threshold= ``` By default, this value is 256 KB. If a profiler shows considerable latency overhead for a reduction operation, it may be beneficial to adjust this value upwards so that fewer reducers are used, resulting in a smaller number of larger network transfers. In some circumstances it may also be necessary to avoid large network transfers, and instead use a greater number of smaller transfers. This can be adjusted with the flag: ```bash --megascale_max_reduction_shard_size= ``` which will divide the buffer up into a greater number of smaller shards than is otherwise needed by the collective. By default, this value is 8 MB. It should be the case that `megascale_max_reduction_shard_size` >> `megascale_sparse_reduction_threshold` since these have opposing effects. For one-to-one transfers (such as `collective-permute`), transfers on a host can be divided into chunks using: ```bash --megascale_chunk_size= ``` If not zero, this divides each one-to-one transfer on a host into multiple chunks of `megascale_chunk_size` bytes (default: 8 MB) to pipeline DMA and network transfers. All other collectives use `megascale_target_dma_size`, `megascale_sparse_reduction_threshold`, and `megascale_max_reduction_shard_size` to control sharding and chunking. --- ## Memory Management & Zero-Copy Transfers Megascale network transfers are designed to be zero-copy: host memory regions are pinned and premapped for DMA to and from the TPU. If premapped memory is exhausted at runtime, the system falls back to on-demand mapping or dynamic memory copies, severely impacting throughput. ### Pre-Mapped Memory Region ```bash --megascale_grpc_premap_memory_bytes=17179869184 # Default: 16 GB (16LL << 30) ``` Controls the size of the host pinned-memory region allocated at initialization for DMA transfers. ### Diagnostic Symptoms 1. **`MapDmaBuffer` in steady-state**: In the XProf Trace Viewer, search for `MapDmaBuffer`. These calls should only occur during startup. If they continue during training iterations, dynamic remapping is occurring. 2. **`Megascale: Memory Copy` traces**: If the Communication Transport track in XProf displays `Megascale: Memory Copy` events, memory buffers are being copied rather than transferred via zero-copy DMA. ### Resolution Increase `--megascale_grpc_premap_memory_bytes` (e.g., to 24 GB or 32 GB, depending on available host memory on your TPU VM instance) and restart the job. --- ### Pjrt Examples (`docs/pjrt/examples.md`) # PJRT Examples ## Example: JAX CUDA plugin 1. PJRT C API implementation through wrapper ([pjrt\_c\_api\_gpu.h](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_gpu.h)). 1. Set up the entry point for the package ([setup.py](https://github.com/google/jax/blob/main/jax_plugins/cuda/setup.py)). 1. Implement an initialize() method ([\_\_init\_\_.py](https://github.com/google/jax/blob/a10854786b6d1bc92a65dd314916b151640789af/plugins/cuda/__init__.py#L31-L51)). 1. Can be tested with any jax tests for CUDA. ## Frameworks Implementations Some references for using PJRT on the framework side, to interface with PJRT devices: - JAX + [jax-ml/jax](https://github.com/jax-ml/jax/blob/main/jax/_src/compiler.py#L248) interacts with PJRT APIs via the `xla_client` APIs - GoMLX + [gomlx/gopjrt](https://github.com/gomlx/gopjrt) + [gomlx/gomlx > backends/xla](https://github.com/gomlx/gomlx/tree/main/backends/xla/xla.go) - ZML + PJRT API wrapper [pjrt.zig](https://github.com/zml/zml/blob/master/pjrt/pjrt.zig) + Load PJRT Plugin [context.zig](https://github.com/zml/zml/blob/master/zml/context.zig#L30-L34) + Interacting with PJRT Buffers [buffer.zig](https://github.com/zml/zml/blob/master/zml/buffer.zig#L36) + Execute a module via PJRT [module.zig](https://github.com/zml/zml/blob/master/zml/module.zig#L863-L886) ## Hardware Implementations - Full integration plugins (PJRT+MLIR+XLA): + [XLA CPU Plugin](https://github.com/openxla/xla/tree/main/xla/pjrt/cpu/cpu_client.cc) + [XLA GPU Plugin](https://github.com/openxla/xla/tree/main/xla/pjrt/gpu/se_gpu_pjrt_client.cc) + [Intel XLA Plugin](https://github.com/intel/intel-extension-for-openxla) - Light integration plugins (PJRT+MLIR): + StableHLO Reference Interpreter plugin (MLIR-based, C++ plugin, to be linked after devlabs) + [Tenstorrent-XLA plugin](https://github.com/tenstorrent/tt-xla/blob/main/src/common/pjrt_implementation/api_bindings.cc) (MLIR-based, C plugin) --- ### Pjrt Pjrt Integration (`docs/pjrt/pjrt_integration.md`) # PJRT plugin integration This doc focuses on the recommendations about how to integrate with PJRT, and how to test PJRT integration with JAX. ## How to integrate with PJRT ### Step 1: Implement [PJRT C API interface](https://github.com/openxla/xla/blob/71a4e6e6e4e9f0f8b8f25c07a32ad489aff19239/xla/pjrt/c/pjrt_c_api.h) **Option A**: You can implement the PJRT C API directly. **Option B**: If you're able to build against C++ code in the [xla repo](https://github.com/openxla/xla) (via forking or bazel), you can also implement the PJRT C++ API and use the C→C++ wrapper: 1. Implement a C++ PJRT client inheriting from the [base PJRT client](https://github.com/openxla/xla/blob/main/xla/pjrt/pjrt_client.h) (and related PJRT classes). Here are some examples of C++ PJRT client: [pjrt\_stream\_executor\_client.h](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/pjrt_stream_executor_client.h), [tfrt\_cpu\_pjrt\_client.h](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/tfrt_cpu_pjrt_client.h). 1. Implement a few C API methods that are not part of C++ PJRT client: * [PJRT\_Client\_Create](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L344-L365). Below is some sample pseudo code (assuming `GetPluginPjRtClient` returns a C++ PJRT client implemented above): ``` #include "third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.h" namespace my_plugin { PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) { std::unique_ptr client = GetPluginPjRtClient(); args->client = pjrt::CreateWrapperClient(std::move(client)); return nullptr; } } // namespace my_plugin ``` Note [PJRT\_Client\_Create](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L344-L365) can take options passed from the framework. [Here](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_gpu_internal.cc#L48-L102) is an example of how a GPU client uses this feature. * [Optional] [PJRT\_TopologyDescription\_Create](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L1815-L1830). * [Optional] [PJRT\_Plugin\_Initialize](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L173-L180). This is a one-time plugin setup, which will be called by the framework before any other functions are called. * [Optional] [PJRT\_Plugin\_Attributes](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L182-L194). With the [wrapper](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_wrapper_impl.h), you do not need to implement the remaining C APIs. ### Step 2: Implement GetPjRtApi You need to implement a method `GetPjRtApi` which returns a `PJRT_Api*` containing function pointers to PJRT C API implementations. Below is an example assuming implementing through wrapper (similar to [pjrt\_c\_api\_cpu.cc](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api_cpu.cc)): ``` const PJRT_Api* GetPjrtApi() { static const PJRT_Api pjrt_api = pjrt::CreatePjrtApi(my_plugin::PJRT_Client_Create); return &pjrt_api; } ``` ### Step 3: Test C API implementations You can call [RegisterPjRtCApiTestFactory](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_test.h#L31C6-L31C33) to run a small set of tests for basic PJRT C API behaviors. ## How to use a PJRT plugin from JAX ### Step 1: Set up JAX You can either use JAX nightly ``` pip install --pre -U jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ pip install git+https://github.com/google/jax ``` or [build JAX from source](https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source). For now, you need to match the jaxlib version with the PJRT C API version. It's usually sufficient to use a jaxlib nightly version from the same day as the TF commit you're building your plugin against, e.g. ``` pip install --pre -U jaxlib==0.6.1.dev20250428 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ ``` You can also build a jaxlib from source at exactly the XLA commit you're building against ([instructions](https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source-with-a-modified-xla-repository)). We will start supporting ABI compatibility soon. ### Step 2: Use jax\_plugins namespace or set up entry\_point There are two options for your plugin to be discovered by JAX. 1. Using namespace packages ([ref](https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-naming-convention)). Define a globally unique module under the `jax_plugins` namespace package (i.e. just create a `jax_plugins` directory and define your module below it). Here is an example directory structure: ``` jax_plugins/ my_plugin/ __init__.py my_plugin.so ``` 2. Using package metadata ([ref](https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata)). If building a package via pyproject.toml or setup.py, advertise your plugin module name by including an entry-point under the `jax_plugins` group which points to your full module name. Here is an example via pyproject.toml or setup.py: ``` # use pyproject.toml [project.entry-points.'jax_plugins'] my_plugin = 'my_plugin' # use setup.py entry_points={ "jax_plugins": [ "my_plugin = my_plugin", ], } ``` Here are examples of how openxla-pjrt-plugin is implemented using Option 2: https://github.com/openxla/openxla-pjrt-plugin/pull/119, https://github.com/openxla/openxla-pjrt-plugin/pull/120. ### Step 3: Implement an initialize() method You need to implement an initialize() method in your python module to register the plugin, for example: ``` import os import jax._src.xla_bridge as xb def initialize(): path = os.path.join(os.path.dirname(__file__), 'my_plugin.so') xb.register_plugin('my_plugin', priority=500, library_path=path, options=None) ``` Please refer to [here](https://github.com/google/jax/blob/8f283bc9ed50d3828bd468ae57b1ee4df1527624/jax/_src/xla_bridge.py#L420) about how to use `xla_bridge.register_plugin`. It is currently a private method. A public API will be released in the future. You can run the line below to verify that the plugin is registered and raise an error if it can't be loaded. ``` jax.config.update("jax_platforms", "my_plugin") ``` JAX may have multiple backends/plugins. There are a few options to ensure your plugin is used as the default backend: * Option 1: run `jax.config.update("jax_platforms", "my_plugin")` in the beginning of the program. * Option 2: set ENV `JAX_PLATFORMS=my_plugin`. * Option 3: set a high enough priority when calling xb.register\_plugin (the default value is 400 which is higher than other existing backends). Note the backend with highest priority will be used only when `JAX_PLATFORMS=''`. The default value of `JAX_PLATFORMS` is `''` but sometimes it will get overwritten. ## How to test with JAX Some basic test cases to try: ``` # JAX 1+1 print(jax.numpy.add(1, 1)) # => 2 # jit print(jax.jit(lambda x: x * 2)(1.)) # => 2.0 # pmap arr = jax.numpy.arange(jax.device_count()) print(jax.pmap(lambda x: x + jax.lax.psum(x, 'i'), axis_name='i')(arr)) # single device: [0] # 4 devices: [6 7 8 9] ``` (We'll add instructions for running the jax unit tests against your plugin soon!) For more examples of PJRT plugins see [PJRT Examples](examples.md). --- ### Xla Pjrt C Docs Pjrt Integration Guide (`xla/pjrt/c/docs/pjrt_integration_guide.md`) # PJRT integration guide This file has been moved into the root [XLA documentation directory](https://github.com/openxla/xla/blob/main/docs/pjrt/pjrt_integration.md). --- --- METRICS --- - Files Extracted: 64 - Estimated Token Budget: ~193485 tokens - Canonical Reference: https://codewiki.google/github.com/openxla/xla