XLA / tf2xla / JIT / AOT

The TF↔XLA bridge. Lives under tensorflow/compiler/. Note: the XLA compiler itself (HLO IR, optimization passes, codegen) was extracted to its own repository at openxla/xla in 2023. What remains here is the bridge: code that lowers TF graphs to HLO and integrates with the TF runtime.

Purpose

• Lower TF subgraphs to HLO (HloModule), the IR consumed by XLA.
• Decide which subgraphs to compile (auto-clustering at runtime, or explicit jit_compile=True).
• Provide an AOT path (tfcompile) that compiles a graph to a static C++ library at build time.
• Surface XLA-compiled functions as ops (XlaCompile, XlaRun, _XlaCompile/_XlaRun) inside the TF runtime.

Directory layout

tensorflow/compiler/
├── tf2xla/                # TF → HLO bridge
│   ├── kernels/           # XLA implementations of TF ops (XlaOpKernel subclasses)
│   ├── ops/               # XLA-specific op registrations
│   ├── lib/
│   ├── xla_compiler.{h,cc}
│   └── xla_op_kernel.{h,cc}
├── jit/                   # Auto-clustering and JIT runtime
│   ├── compilability_check_util.{h,cc}
│   ├── encapsulate_subgraphs_pass.{h,cc}
│   ├── encapsulate_xla_computations_pass.{h,cc}
│   ├── mark_for_compilation_pass.{h,cc}
│   ├── xla_compile_op.cc
│   ├── xla_compile_on_demand_op.cc
│   ├── xla_launch_util.{h,cc}
│   └── ...
├── aot/                   # tfcompile AOT compiler
│   ├── tfcompile_main.cc
│   ├── codegen.{h,cc}
│   ├── compile.{h,cc}
│   └── ...
├── mlir/                  # MLIR-based passes (see compilers/mlir.md)
├── plugin/                # Plugin device/compiler hooks
├── tests/                 # Lit/MLIR + Python compiler tests
└── tf2tensorrt/           # TensorRT (see compilers/tensorrt.md)

tf2xla — TF→HLO

tf2xla lowers a TF subgraph to an xla::HloModule. The unit of lowering is an XlaOpKernel — a subclass of OpKernel that, instead of computing on Tensors, builds HLO into an XlaBuilder.

class AddOp : public XlaOpKernel {
 public:
  explicit AddOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
  void Compile(XlaOpKernelContext* ctx) override {
    xla::XlaOp x = ctx->Input(0);
    xla::XlaOp y = ctx->Input(1);
    ctx->SetOutput(0, xla::Add(x, y));
  }
};
REGISTER_XLA_OP(Name("Add"), AddOp);

Each TF op that supports XLA has one of these in tensorflow/compiler/tf2xla/kernels/. The XLA op kernel set is separate from the regular op kernel set in tensorflow/core/kernels/ — the same op (Add, MatMul, Conv2D, …) has both.

XlaCompiler (tensorflow/compiler/tf2xla/xla_compiler.cc) walks a Graph, calls each op's Compile method, and assembles an HloModule.

JIT — auto-clustering

Auto-clustering decides at runtime which subgraphs of a Graph should be compiled. The pass mark_for_compilation_pass.cc (tensorflow/compiler/jit/) labels nodes with _XlaCluster attributes; encapsulate_subgraphs_pass.cc then replaces each cluster with an _XlaCompile/_XlaRun pair.

Key bits:

• compilability_check_util.{h,cc} — decides whether each op is XLA-compatible (some ops have no XLA implementation; some can't be JITed in graphs that read external state).
• xla_compile_op.cc — runs at execution time, produces an XlaExecutable, caches it.
• xla_launch_util.{h,cc} — handles input/output marshalling between TF tensors and XLA buffers.

Auto-clustering can be opt-in (set tf.config.optimizer.set_jit(True)) or driven by @tf.function(jit_compile=True).

jit_compile=True

The recommended way to use XLA:

@tf.function(jit_compile=True)
def model(x):
    return tf.nn.relu(tf.linalg.matmul(x, w) + b)

This bypasses auto-clustering and forces the whole function to compile. The compilation happens once per input signature; the resulting XlaExecutable is cached.

AOT — tfcompile

tensorflow/compiler/aot/ contains the tfcompile binary, which AOT-compiles a GraphDef into:

• A C++ class with Run(), result(N), arg(N) methods.
• A .o file linkable into a binary.
• A small header for the user.

tfcompile requires the graph's input shapes to be known at compile time; it produces a self-contained binary that doesn't link the TF runtime. This is the path that powers, e.g., on-device inference in some Google products.

The Bazel macro tf_library (tensorflow/compiler/aot/tfcompile.bzl) wraps invocation of tfcompile for build files.

Integration points

• Auto-cluster output: produces nodes with op type _XlaCompile / _XlaRun. These live in the TF graph and are executed by the regular TF executor.
• Runtime: XlaCompileOnDemand runs once and caches.
• Profiler: TraceMe events in xla_compile_op.cc show up in TensorBoard.
• MLIR (compilers/mlir): the new path lowers TF→MLIR (tf dialect) → mhlo/stablehlo and then into the same xla::HloModule. Eventually most paths will go through MLIR.

Where the actual XLA compiler is

Optimization passes, codegen for CPU/GPU/TPU, the runtime that consumes xla::Executable, and StableHLO definitions live in openxla/xla. This repo references that code via the third_party/xla Bazel target.

Entry points for modification

• New XLA-supported op → add an XlaOpKernel subclass in tensorflow/compiler/tf2xla/kernels/ and a REGISTER_XLA_OP.
• Auto-clustering policy → tensorflow/compiler/jit/mark_for_compilation_pass.cc and compilability_check_util.cc.
• AOT codegen → tensorflow/compiler/aot/codegen.cc.

Related

• compilers/mlir — the MLIR pipeline that increasingly replaces tf2xla.
• systems/tpu — TPU execution is XLA-only.
• features/tf-function-and-autograph — jit_compile=True lives here.