XLA Up Close: What It Optimizes, and What It Won't

XLA, the compiler behind JAX, TensorFlow, and PyTorch/XLA, optimizes array programs by freezing shapes, statically allocating buffers, and fusing operations against a global cost model, which makes it fast but rigid. A demonstration shows XLA eliminating add-zero, multiply-one, and double-transpose operations from a simple function, proving its algebraic-simplification strength. The same design decision that enables high performance also causes inflexibility, as XLA's strengths and limitations stem from its ahead-of-time compilation and static analysis approach.

XLA Up Close: What It Optimizes, and What It Won't Here is a small function and the entire program XLA compiled it into: 1 2 3 python def simp x : y = x + 0.0 1.0 return jnp.transpose jnp.transpose y 1 2 3 4 php ENTRY %main x: f32 3,3 - f32 3,3 { %x = f32 3,3 {1,0} parameter 0 ROOT %copy = f32 3,3 {1,0} copy %x } The add-zero, the multiply-one, and both transposes are gone. XLA proved the whole thing is the identity and emitted a single copy. This is not a toy result; it is the same algebraic-simplification machinery that lets XLA take a page of dense linear algebra and hand back a fused kernel that runs at the memory-bandwidth limit of a TPU. XLA is a genuine optimizing compiler, and a very good one. 1 fn:1 It is also rigid in ways that will bite you, and the interesting thing — the thing worth a long post — is that its strengths and its limitations are the same design decision . XLA freezes your shapes, sees the whole graph at once, allocates every buffer statically, and fuses against a global cost model you do not control. That is exactly why it is fast, and exactly why it is inflexible. This is a review and a critique of that bargain. I will show you what XLA actually does, using real HLO dumps, sourcing each claim to the OpenXLA source and docs, and then I will tell you where the bargain stops paying off. If you have not met the layer above XLA: it is the compiler behind JAX, behind TensorFlow, and behind PyTorch through PyTorch/XLA. 1 JAX traces your Python into a small typed IR and hands it to XLA; what XLA does with it from there is the subject here. The tracing half — how Python becomes that IR — is a topic of its own, and I leave it aside. Where XLA sits, and what it eats XLA is an ahead-of-time compiler for array programs. “Ahead of time” is relative to execution , not to your process: it compiles the first time a given shape signature appears, caches the result, and reuses it. 2 Its input is a graph of high-level tensor operations, and the modern portable front door to that graph is StableHLO , an MLIR dialect that JAX emits and that XLA ingests. The full path: 3 fn:3 1 2 3 4 jaxpr ──▶ StableHLO ──▶ HLO ──▶ optimization passes ──▶ optimized HLO ──▶ backend ├─ LLVM CPU ├─ LLVM/PTX + cuBLAS/cuDNN/Triton GPU └─ XLA:TPU closed backend The naming here is a genuine mess and worth untangling once. There is the classic HLO — a hand-written C++ IR with its own protobuf and text format that predates MLIR and is explicitly not based on it. 4 There is the MLIR world layered on later: MHLO , and then StableHLO , which is built on MHLO and adds serialization plus backward/forward compatibility guarantees, so that it can serve as a stable portability contract between frameworks and compilers. JAX lowers to StableHLO; XLA converts StableHLO into its internal HLO and optimizes that. 3 fn:3 5 fn:2 Even the acronym is unsettled — XLA’s own architecture page glosses HLO as “high level operations” while its terminology page says “High Level Optimizer.” 2 fn:7 5 fn:2 I will use “HLO” for the thing XLA optimizes and “StableHLO” for what it is handed, because that is how the dumps read. 4 fn:4 Here is what a jit -ed relu x @ W + b looks like when JAX hands it over, straight from jax.jit f .lower ... .as text : 1 2 3 4 5 6 7 8 9 10 11 12 func.func public @main %arg0: tensor<4x8xf32 , %arg1: tensor<8x16xf32 , %arg2: tensor<16xf32 - tensor<4x16xf32 { %0 = stablehlo.dot general %arg0, %arg1, contracting dims = 1 x 0 : tensor<4x8xf32 , tensor<8x16xf32 - tensor<4x16xf32 %1 = stablehlo.broadcast in dim %arg2, dims = 1 : tensor<16xf32 - tensor<1x16xf32 %2 = stablehlo.broadcast in dim %1, dims = 0, 1 : tensor<1x16xf32 - tensor<4x16xf32 %3 = stablehlo.add %0, %2 : tensor<4x16xf32 %cst = stablehlo.constant dense<0.000000e+00 : tensor<f32 %4 = stablehlo.broadcast in dim %cst, dims = : tensor<f32 - tensor<4x16xf32 %5 = stablehlo.maximum %3, %4 : tensor<4x16xf32 return %5 : tensor<4x16xf32 } Every value has SSA-style single-assignment semantics it is MLIR underneath , every tensor is typed, and broadcasting is explicit — the + b is a real broadcast in dim followed by an add , because XLA models it as an operation with a cost, not as sugar. The op vocabulary is small: XLA’s opcode enum holds roughly 150 opcodes and is statically asserted to stay under 256. 6 It covers contraction dot general , convolution , reduce , dynamic-slice , the collectives all-reduce , all-gather , and the escape hatch custom-call .That small, typed, functional vocabulary is what keeps the optimizer tractable. One caveat to hold: classic HLO shapes are static, but StableHLO can express 7 fn:5 bounded dynamism tensor<?x2xf32 , a distinction that matters a great deal in the critique. 3 fn:3 The optimizer, pass by pass Dump the pipeline XLA FLAGS=--xla dump to=... and XLA writes out every intermediate module — on a trivial function I counted 39 stages. The architecture doc frames it as three phases: target-independent optimization, target-dependent optimization, and target-specific codegen. 5 Cleaned of noise, the spine is: 1 2 simplification algsimp, CSE, DCE ─▶ layout assignment ─▶ fusion ─▶ copy insertion ─▶ scheduling ─▶ backend codegen LLVM IR ─▶ object / PTX Each stage is where one of XLA’s strengths — and one of its sharp edges — lives. Simplification The front of the pipeline is ordinary optimizing-compiler work: algebraic simplification, 8 common-subexpression elimination, and dead-code elimination. The 8 fn:10 simp example at the top is this stage doing its job — x+0 1 folds away, transpose transpose x cancels, and the whole computation collapses to a copy. This is real and it is good. But notice the ceiling. These are local, syntactic rewrites over the graph you handed it. XLA will cancel a double transpose; it will not discover that two differently-written subgraphs compute the same mathematical function, and it will not reassociate a reduction to expose parallelism it was not given. The algebra it knows is a fixed rule set in algebraic simplifier.cc , not a solver. Layout assignment This is the first thing that surprises people, because it is invisible in your source. Every array in optimized HLO carries a layout : the {1,0} in f32 4,16 {1,0} is the minor-to-major dimension order — here, row-major; {0,1} would be column-major. 9 XLA chooses these physical layouts to make consuming operations efficient, and on TPU the story is richer still, with tiled layouts because the hardware wants data in tiles, not flat rows. 9 fn:8 The catch is that layout is a constraint-propagation problem, and when constraints collide, XLA inserts copies. The layout assignment pass carries explicit machinery for exactly this — SetOperandLayout , CreateCopyWithNewLayout , CopyOperandIfLayoutsDiffer , and constraint propagation for custom calls. 10 Watch it happen when a computation feeds an operation that demands a particular layout — an eigendecomposition, which lowers to a LAPACK routine that insists on column-major input: 1 2 3 %eigh = f32 16,16 {0,1}, f32 16 {0}, s32 custom-call %multiply copy fusion , custom call target="lapack ssyevd ffi", operand layout constraints={f32 16,16 {0,1}} The data arrives row-major {1,0} and the routine wants column-major {0,1} , so XLA quietly materializes a transposing copy — the multiply copy fusion feeding the call. That copy never appears in your Python. It is a real cost, it is layout-driven, and the only way to see it is to read the HLO. A “mysterious” transpose in a profile is very often layout assignment reconciling a constraint you could not see. Fusion Fusion is where XLA earns its reputation, and where the contrast with a framework compiler like PyTorch’s TorchInductor is sharpest. The goal is the same one I covered in the torch.compile post /posts/how-torch-compile-works/ : keep intermediate values in registers instead of streaming them to and from memory. But XLA does it at the granularity of HLO ops, folds the result into explicit fusion nodes whose bodies are their own computations, and — critically — distinguishes fusion kinds . The FusionKind enum has exactly four: kLoop , kInput , kOutput , kCustom . 11 fn:11 Take tanh exp x 2 + 1 . XLA fuses the entire chain into one loop kernel: 1 2 3 4 5 6 7 8 9 php %fused computation param 0: f32 1024 - f32 1024 { %exp.0 = f32 1024 exponential %param 0 %mul.0 = f32 1024 multiply %exp.0, broadcast constant 2 %add.0 = f32 1024 add %mul.0, broadcast constant 1 ROOT %tanh.0 = f32 1024 tanh %add.0 } ENTRY %main x: f32 1024 - f32 1024 { ROOT %fusion = f32 1024 fusion %x , kind=kLoop, calls=%fused computation } Four elementwise ops, one kind=kLoop fusion, one kernel, one pass over memory. The base pass fuses “vertically,” producers into consumers. 12 And relu x @ W + b splits by kind: the matmul becomes a kind=kCustom fusion routed to a tuned kernel, while the bias-add and ReLU merge into a kind=kLoop fusion on top. On GPU the main fusion pass is priority fusion , which is explicit about being cost-model-driven: it “assigns a priority to each producer instruction based on the estimated performance benefit of fusing it into its consumers,” where priority = time unfused - time fused . 12 fn:12 That cost model is the strength and the problem. It is global and automatic, which is why XLA fuses well across a whole graph without you lifting a finger. It is also internal: the decision comes from the model in priority fusion , not from anything you can annotate. When the model under-fuses, or when the pattern you want cannot be expressed at HLO-op granularity — flash attention is the canonical example — you are out of levers inside XLA. That is precisely why Pallas exists. Buffer assignment Here is a defining property and a genuine strength: XLA does whole-program static memory allocation . Because it sees the entire graph and every shape is known, the buffer assignment pass computes buffer liveness and assigns every value a concrete offset in a preplanned allocation before the program runs; a single allocation may hold multiple values “with disjoint liveness” — buffer reuse, computed statically. 13 The dump is explicit: 1 2 3 4 5 6 7 8 BufferAssignment: allocation 0: size 1024, parameter 0, shape |f32 16,16 | allocation 1: size 64, output, maybe-live-out: value: eigh{1} size=64, offset=0 value: broadcast select fusion size=64, offset=0 <- same buffer, reused allocation 2: size 1092, preallocated-temp: value: multiply copy fusion size=1024, offset=64 value: eigh{} size=24, offset=0 allocation 1 is shared by two values whose lifetimes do not overlap; allocation 2 is a single temp arena holding several temporaries at fixed offsets. No malloc on the execution path, peak memory known at compile time. This pairs with donation , where an input buffer is aliased to an output so an update happens in place — configured through HloInputOutputAliasConfig and essential for keeping optimizer state and KV caches from doubling your memory. 13 For static-shaped programs this is close to optimal. It also depends entirely on static shapes; the moment a buffer’s size is unknown at compile time, the guarantee weakens. Hold that thought. Copy insertion and scheduling Two shorter passes close out the target-independent work. copy insertion is a legalization pass that “inserts copies kCopy instructions to eliminate several kinds of problems” — chiefly values that are simultaneously live while sharing a buffer, and enforcing the aliasing that donation promised. 14 Then the scheduler orders instructions; XLA’s latency hiding scheduler is “oriented to hiding latencies of operations that can run in parallel,” and it carries explicit knobs for overlapping collectives with compute all reduce overlap limit , all gather overlap limit , and friends — a detail that matters in the sharding post. 15 fn:15 Codegen and the backends After scheduling, XLA hands off to a backend. On CPU and GPU it lowers to LLVM IR — the dump even drops the .ll and .o files next to the HLO — and then to native code or, on GPU, PTX via the LLVM NVPTX backend. 5 fn:2 16 GPU additionally routes dense matmuls to cuBLAS and convolutions to cuDNN as custom calls cublas$gemm , cudnn$convForward , rewritten in by dedicated passes.And for “more advanced fusions which include matrix multiplication or softmax, XLA:GPU uses Triton as a code-generation layer,” through a real Triton emitter in the source tree — a direct through-line to the 16 fn:16 Triton compiler deep-dive /posts/triton-compiler-deep-dive/ . TPU goes to the closed-source XLA:TPU backend; a maintainer states plainly that “the TPU backend is non-OSS, so you can’t build it from the GH repo.” 17 fn:17 The whole StableHLO/HLO layering is one more instance of the MLIR dialect-stack pattern I wrote about in 18 fn:19 the MLIR tour /posts/mlir-dialect-stack-for-ml/ . This is not hand-waving; you can watch it on real hardware. Lowering a softmax on an A100-SXM4-80GB, XLA:GPU emits precisely a Triton fusion: 1 2 3 4 ROOT %triton softmax.2 = f32 2048,2048 {1,0} fusion %x.1 , kind=kCustom, calls=%triton softmax computation.2, backend config={"fusion backend config":{"kind":" triton", "block level fusion config":{"num warps":"4","output tiles": {"sizes": "1","2048" } }}} The "kind":" triton" is XLA choosing Triton, and the block level fusion config is XLA choosing the launch parameters for you — four warps, an output tile of 1x2048 . 17 A matmul, by contrast, lowers by default to a kCustom gemm fusion also Triton ; disable that path with --xla gpu enable triton gemm=false and the same matmul falls back to a cuBLAS custom-call, custom call target=" cublas$lt$matmul" .Which of the two you get is decided by autotuning, not by you. 16 fn:16 Autotuning One backend behavior deserves its own note because it dominates the compile-time story later. On GPU, XLA autotunes : for a given fusion it generates candidate implementations Triton configs, cuBLASLt, cuDNN , profiles them on the real device, and keeps the fastest — the autotuner core profiles all candidates and selects the best. 19 It is effective and it is expensive; OpenXLA’s own persisted-autotuning doc admits the tuning “can take a long time if there are many fusions,” which is the entire reason a cache exists. You can feel it directly: on the A100, compiling a single 4096x4096 matmul took about 19 fn:18 2 seconds , and an 8192x8192 one about 2.8 seconds — for one fused operation, almost all of it the autotuning search. The open-source autotuner is GPU-only; the equivalent for TPU lives in the closed backend. 18 fn:19 The critique Everything above is XLA at its best: a whole-program optimizer that fuses aggressively, allocates statically, and produces genuinely fast code. Now the other side. None of what follows says XLA is bad — it says the design has a shape, and the shape has edges you will hit. Static shapes are not a preference, they are the whole deal Every strength in the last section — static allocation, global fusion, layout planning, in-place donation — is downstream of one assumption: the shapes are known at compile time. When they are, XLA is excellent. When they are not, you are working against the grain of the entire system. StableHLO can represent bounded dynamism, but OpenXLA is candid that there is “limited framework support,” and specifically that “JAX does not currently trace operations which lead to data dependent dynamism.” 20 fn:21 And modern ML is full of shapes that are not static. LLM inference has a sequence length that grows token by token and a KV cache that grows with it; batched serving has ragged batches. These are the mainstream, not the exotic. XLA’s answer is bounded dynamic shapes plus padding and bucketing into a few static shapes compiled separately — JAX’s own issue tracker states that “JAX currently supports dynamic shape input through padding or bucket padding,” and the JAX scaling book notes that “prefills are padded to the longest sequence and we waste a lot of compute.” 20 fn:21 21 Truly data-dependent shapes it refuses outright. Ask for the indices of the nonzero elements of a vector inside jit : 1 2 3 jit lambda x: jnp.nonzero x jnp.array 0., 1., 0., 2. ConcretizationTypeError: Abstract tracer value encountered where a concrete value is expected: traced array with shape int32 JAX’s own nonzero docstring is explicit: “Because the size of the output of nonzero is data-dependent, the function is not compatible with JIT,” and it must be given a static size= to be usable under transformation. 22 This is the correct behavior given the design, and it is also a wall — the same tension I wrote about in the VLIW post /posts/why-vliw-architecture-is-popular-again/ : systems that win by assuming static, predictable structure pay for it precisely when the workload turns dynamic, and the workloads are turning dynamic. The compile-time tax is real, but be honest about where it comes from You will read that XLA “takes minutes to compile.” Sometimes it does, and it is worth being precise, because the naive version of the complaint is wrong. A 24-layer, 1024-wide dense tower — 360 HLO instructions after optimization — compiled on CPU in my measurements in about 0.03 seconds . XLA is not slow at compiling graphs. Move to the GPU, though, and a single 4096x4096 matmul took roughly 2 seconds to compile on an A100 — dozens of times the whole CPU tower, for one operation — because the backend was autotuning it. The minutes come from three places. First, autotuning : as above, XLA benchmarks kernel configs on real hardware at compile time, and that search dominates — which is why the persisted-autotuning cache exists. 19 Second, graph size : a full model with the backward pass and an optimizer step is not 360 instructions but tens of thousands, and several passes are super-linear; real long-compile reports exist independent of autotuning. Third — the one that actually hurts in deployment — 23 fn:28 per-shape specialization : every new shape signature recompiles, so a serving workload with varying sequence lengths spends real wall-clock recompiling, a long-standing complaint on the JAX tracker. The mitigations are a persistent on-disk compilation cache and compiling a repeated block instead of the whole model, and JAX documents both. 20 fn:21 They also tell you the shape of the problem: XLA wants to compile once, for one shape, and amortize. When your deployment cannot promise that, the model fights you. 24 fn:20 Fusion is automatic, which means you cannot argue with it The flip side of a global cost model is that it is the cost model. XLA’s fusion decision comes from the internal priority model, not from anything you can annotate. 12 For most code the automatic decision is fine to good. But if the heuristics decide not to fuse two operations you know should fuse, or to lay out a tensor in a way that forces a copy, your recourse is to restructure the source and hope the model reads it differently, or to leave XLA entirely via custom call . There is no user-facing “fuse these, not those.” For the last slice of performance on a kernel that matters, the absence of a manual override is a real ceiling — and it is why performance-critical teams reach for Pallas. custom call is an escape hatch and an optimization barrier When XLA cannot express or generate what you need, you drop to custom call — a node pointing at an external function. The StableHLO spec describes it exactly: it is an “escape hatch” whose “semantics … are implementation-defined and are not analyzed by the compiler.” 25 That is the whole problem in one sentence. Look again at the eigh example: XLA knows nothing about what happens inside lapack ssyevd ffi ; it cannot fuse across it, cannot reorder around it, and must satisfy its layout constraints with copies. JAX’s FFI docs make the runtime consequence concrete — an FFI call is “an opaque black box: JAX can’t look inside it,” and a sharded input may force XLA to “gather the whole array onto every device with an all-gather” and run the call redundantly.This is the same property I flagged for NCCL collectives in the torch.compile post: an opaque region that fragments optimization on either side. The more of your hot path lives in custom calls, the less of it XLA is actually optimizing. 25 fn:24 Opacity and the closed backend XLA is a black box in daily use. When performance is wrong, there is no gentle story — you dump HLO, learn to read fusion kinds and layout annotations, and reverse-engineer a cost model that is not documented as a contract. The tooling --xla dump to , the HLO-as-HTML renderer, the pass dumps is genuinely good once you know it exists, but the learning curve is steep and self-taught. And the highest-value backend, XLA:TPU, is closed: a maintainer notes that “the TPU dependent optimizations are mostly not open source, essentially just those parts that are shared with GPU and CPU.” 18 For a system this central to modern ML, the debuggability floor is higher than it should be. Portability of the IR is not portability of performance StableHLO’s pitch is a stable, portable IR — write once, run on any backend that speaks it — and as a portability contract it delivers. 3 But portable code is not portable performance. Fast on TPU means the right tiled layouts and XLA:TPU’s fusion choices; fast on GPU means cuBLAS, Triton, and GPU autotuning; fast on CPU means something else again. 9 fn:8 16 fn:16 The same StableHLO lowered to two backends gives you two 17 fn:17 correct programs, one or both of which may need backend-specific coaxing to be fast . StableHLO resolves the portability half and leaves the performance half where it was — the trilemma every heterogeneous system meets. Placement and topology live outside the IR The last one is the one this series has been circling, and the pipeline names it without meaning to: among the passes are a pre-spmd-partitioner , a sharding-removal , and an async-collective-replacer — a whole sub-machine for turning sharding annotations into partitioned, communicating programs. That machine is a separate one, bolted on the side, and the source makes the distinction sharp. An array’s type is its dtype and its shape. Its layout is physical and first-class. But where the array lives — which device, which slice of a mesh — is not part of the type: sharding is a nullable attribute hung off the instruction has sharding , set sharding , backed by the OpSharding proto, not a field of Shape . 26 A separate partitioner — GSPMD, and now the MLIR-based Shardy , which “incorporates the best of GSPMD and PartIR” — consumes those annotations and adds “the necessary data movement/formatting and collective operations.” 27 fn:26 You can literally watch the annotation get stripped by a pass. 28 fn:27 That is the seam. XLA models the physical layout of bytes within a device beautifully and models the placement of bytes across devices as an annotation that a side-machine consumes. The distinction is not academic: it is why XLA can fuse a matmul perfectly and still have no idea that one of its operands is about to require a cross-device transfer. The bargain, stated plainly XLA is one of the best static-shape, whole-program array compilers ever built, and I do not say that as a hedge before the criticism — I say it because the criticism only makes sense against that baseline. It fuses across an entire model, allocates every byte before the program runs, plans layouts down to tile order, and produces code that saturates real hardware. Every one of those wins comes from the same four commitments: freeze the shapes, see the whole graph, allocate statically, fuse with a global model. And every limitation — the recompiles, the rejected dynamic shapes, the un-overridable fusion, the custom call walls, the placement-as-annotation — is the bill for those same four commitments. You do not get XLA’s strengths and dodge its rigidities. They are one decision. That is worth internalizing before you reach for it. XLA shows how much performance is available once you assume a static graph on a single device. What it does not model — dynamism, and where the data lives — is not an oversight; it is the frontier. And the placement machine XLA strips out is where I would look next: sharding, like layout before it, is something you find out about at trace time rather than something the type system ever promised you. References Disclaimer: This article was generated using the Claude Opus 4.8 model. XLA README and Architecture. “The XLA compiler takes models from popular ML frameworks such as PyTorch, TensorFlow, and JAX, and optimizes them for high-performance execution.” Repo https://github.com/openxla/xla , Architecture https://openxla.org/xla/architecture ↩︎ fnref:1 ↩︎ fnref:1:1 2 JAX: Ahead-of-time lowering and compilation. .lower produces StableHLO; .compile produces the optimized executable; compiled functions are “specialized to a particular set of argument ‘types,’ such as arrays with a specific shape and element type.” Link https://docs.jax.dev/en/latest/aot.html ↩︎ fnref:7 ↩︎ fnref:7:1 2 StableHLO. Portability layer between frameworks and compilers; “based on the MHLO dialect” with serialization and versioning; supports bounded dynamic shapes. Repo https://github.com/openxla/stablehlo , Spec https://openxla.org/stablehlo/spec ↩︎ fnref:3 ↩︎ fnref:3:1 2 ↩︎ fnref:3:2 3 ↩︎ fnref:3:3 4 XLA Terminology — HLO. “HLO … is an internal graph representation IR for the XLA compiler … It is not based on MLIR, and has its own textual syntax and binary protobuf based representation.” Note: this page expands HLO as “High Level Optimizer,” while the architecture page uses “high level operations.” Link https://openxla.org/xla/terminology ↩︎ fnref:4 ↩︎ fnref:4:1 2 XLA Architecture — the compilation pipeline. Target-independent optimization CSE, fusion, buffer analysis , target-dependent optimization, and LLVM-based codegen for CPU/GPU. Link https://openxla.org/xla/architecture ↩︎ fnref:2 ↩︎ fnref:2:1 2 ↩︎ fnref:2:2 3 ↩︎ fnref:2:3 4The HLO opcode enum hlo opcode.h . kDot , kConvolution , kReduce , kDynamicSlice , kAllReduce , kAllGather , kCustomCall , … ; ~150 opcodes, statically asserted under 256. Source https://github.com/openxla/xla/blob/main/xla/hlo/ir/hlo opcode.h ↩︎ fnref:6 XLA Operation Semantics. Definitions of dot general , convolution , reduce , dynamic-slice , all-reduce / all-gather , and custom-call “Call a user-provided function within a computation” . Link https://openxla.org/xla/operation semantics ↩︎ fnref:5 Simplification passes. algebraic simplifier “mostly arithmetic simplifications” , hlo dce removes dead instructions/computations , hlo cse common-subexpression elimination . algsimp https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/algebraic simplifier.cc , dce https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/hlo dce.cc , cse https://github.com/openxla/xla/blob/main/xla/service/hlo cse.h ↩︎ fnref:10 ↩︎ fnref:10:1 2 XLA: Shapes and Layout. Minor-to-major layout order; “ {1,0} ” is row-major, “ {0,1} ” column-major; tiled layouts for accelerators. Link https://openxla.org/xla/shapes , Tiled layout https://openxla.org/xla/tiled layout ↩︎ fnref:8 ↩︎ fnref:8:1 2 ↩︎ fnref:8:2 3Assigns layouts “while satisfying all necessary invariants and minimizing cost,” with layout assignment.cc . SetOperandLayout , CreateCopyWithNewLayout , CopyOperandIfLayoutsDiffer , and custom-call constraint propagation. Source https://github.com/openxla/xla/blob/main/xla/service/layout assignment.cc ↩︎ fnref:9 FusionKind . enum class FusionKind { kLoop, kInput, kOutput, kCustom }; . Source https://github.com/openxla/xla/blob/main/xla/hlo/ir/hlo instruction.h ↩︎ fnref:11 Fusion passes. instruction fusion fuses “vertically, meaning producing instructions are fused into their consumers”; priority fusion XLA:GPU assigns priority from a cost model, “ priority = time unfused - time fused .” instruction fusion https://github.com/openxla/xla/blob/main/xla/service/instruction fusion.h , priority fusion https://github.com/openxla/xla/blob/main/xla/backends/gpu/transforms/priority fusion.h ↩︎ fnref:12 ↩︎ fnref:12:1 2 ↩︎ fnref:12:2 3Assigns LogicalBuffers to BufferAllocations at fixed offsets; “A single BufferAllocation may hold LogicalBuffers with disjoint liveness” reuse . Donation is configured via buffer assignment.cc and input/output aliasing. HloInputOutputAliasConfig . buffer assignment https://github.com/openxla/xla/blob/main/xla/service/buffer assignment.h , alias config https://github.com/openxla/xla/blob/main/xla/hlo/ir/hlo input output alias config.h ↩︎ fnref:13 ↩︎ fnref:13:1 2“A legalization HLO pass which inserts copies kCopy instructions to eliminate several kinds of problems in the HLO module.” copy insertion.cc . Source https://github.com/openxla/xla/blob/main/xla/service/copy insertion.cc ↩︎ fnref:14 “A scheduler oriented to hiding latencies of operations that can run in parallel with other operations,” with collective-overlap limits. latency hiding scheduler.cc . Source https://github.com/openxla/xla/blob/main/xla/service/latency hiding scheduler.h ↩︎ fnref:15 GPU libraries and rewriters. cuBLAS/cuDNN custom-call targets cublas$gemm , cudnn$convForward in cublas cudnn.h ; gemm rewriter / conv rewriter lower matmul/convolution to them. targets https://github.com/openxla/xla/blob/main/xla/service/gpu/cublas cudnn.h , GPU arch https://openxla.org/xla/gpu architecture ↩︎ fnref:16 ↩︎ fnref:16:1 2 ↩︎ fnref:16:2 3 ↩︎ fnref:16:3 4 XLA:GPU and Triton. “For more advanced fusions which include matrix multiplication or softmax, XLA:GPU uses Triton as a code-generation layer”; Triton emitter under xla/backends/gpu/codegen/triton/ . GPU arch https://openxla.org/xla/gpu architecture , Triton emitter https://github.com/openxla/xla/blob/main/xla/backends/gpu/codegen/triton/fusion.h ↩︎ fnref:17 ↩︎ fnref:17:1 2 ↩︎ fnref:17:2 3 XLA:TPU is closed-source. Maintainer: “the TPU backend is non-OSS, so you can’t build it from the GH repo”; and “the TPU dependent optimizations are mostly not open source.” 11599 https://github.com/openxla/xla/issues/11599 , 25613 https://github.com/openxla/xla/issues/25613 ↩︎ fnref:19 ↩︎ fnref:19:1 2 ↩︎ fnref:19:2 3 Autotuning. The autotuner profiles all candidate implementations and selects the fastest; OpenXLA notes tuning “can take a long time if there are many fusions,” motivating a persisted cache. autotuner https://github.com/openxla/xla/blob/main/xla/backends/autotuner/autotuner.cc , persisted autotuning https://openxla.org/xla/persisted autotuning ↩︎ fnref:18 ↩︎ fnref:18:1 2 ↩︎ fnref:18:2 3 Dynamic shapes are limited. StableHLO: “limited framework support … JAX does not currently trace operations which lead to data dependent dynamism.” JAX supports dynamic input “through padding or bucket padding”; variable-length arrays “trigger recompilation.” StableHLO dynamism https://openxla.org/stablehlo/dynamism , JAX 26265 https://github.com/jax-ml/jax/issues/26265 , JAX 2521 https://github.com/jax-ml/jax/issues/2521 ↩︎ fnref:21 ↩︎ fnref:21:1 2 ↩︎ fnref:21:2 3 KV-cache / prefill padding. “Prefills are padded to the longest sequence and we waste a lot of compute.” JAX scaling book, inference https://github.com/jax-ml/scaling-book/blob/main/inference.md ↩︎ fnref:23 “Because the size of the output of jax.numpy.nonzero . nonzero is data-dependent, the function is not compatible with JIT”; requires a statically-specified size= . Link https://docs.jax.dev/en/latest/ autosummary/jax.numpy.nonzero.html ↩︎ fnref:22 Real-world long-compile reports. E.g. XLA compilation orders of magnitude slower on newer GPU targets, and recurring compile-time spikes on the JAX tracker. openxla/xla 44517 https://github.com/openxla/xla/issues/44517 , jax 30185 https://github.com/jax-ml/jax/issues/30185 ↩︎ fnref:28 JAX: caching and persistent compilation cache. “JAX will store copies of compiled programs on disk”; jit “will get compiled, and the resulting XLA code will get cached,” reused on subsequent calls with the same shapes/dtypes. persistent cache https://docs.jax.dev/en/latest/persistent compilation cache.html , jit caching https://docs.jax.dev/en/latest/jit-compilation.html ↩︎ fnref:20 StableHLO spec: an escape hatch whose “semantics … are implementation-defined and are not analyzed by the compiler.” JAX FFI: “an opaque black box: JAX can’t look inside it.” custom call is opaque to the optimizer. StableHLO custom call https://openxla.org/stablehlo/spec custom call , JAX FFI https://docs.jax.dev/en/latest/ffi.html ↩︎ fnref:24 ↩︎ fnref:24:1 2 Sharding is an annotation, not part of the type. HloInstruction carries an optional, nullable sharding has sharding , set sharding separate from its Shape ; backed by the OpSharding proto. hlo instruction.h https://github.com/openxla/xla/blob/main/xla/hlo/ir/hlo instruction.h , hlo sharding.h https://github.com/openxla/xla/blob/main/xla/hlo/ir/hlo sharding.h ↩︎ fnref:25 SPMD partitioner GSPMD . SpmdPartitioner with SPMDCollectiveOpsCreator , “a set of functions that create the cross-partition collective ops”; GSPMD “consumes HLO with sharding annotations … and produces a sharded HLO … with proper collectives.” spmd partitioner https://github.com/openxla/xla/blob/main/xla/service/spmd/spmd partitioner.h , GSPMD paper https://arxiv.org/abs/2105.04663 ↩︎ fnref:26 Shardy. “An MLIR-based tensor partitioning system … Built from the collaboration of both the GSPMD and PartIR teams, it incorporates the best of both”; the partitioner adds “the necessary data movement/formatting and collective operations.” Repo https://github.com/openxla/shardy , Overview https://openxla.org/shardy/overview ↩︎ fnref:27 CC BY 4.0 https://creativecommons.org/licenses/by/4.0/ by the author.