wehub-resource-sync
72 lines
3.6 KiB
Markdown
72 lines
3.6 KiB
Markdown
# 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.
|