{"slug": "a-tour-of-xla-where-mlir-lives-and-where-it-doesn-t", "title": "A Tour of XLA: Where MLIR Lives (and Where It Doesn't)", "summary": "XLA, the compiler under JAX, TensorFlow, and PyTorch/XLA, uses two intermediate representations: classic HLO (a hand-built C++ IR) and MLIR dialects such as StableHLO and CHLO, with a translation layer moving programs between them. The core optimizer works on classic HLO, while MLIR appears at the front door (framework targeting), the sharding system, and code generation, making XLA a mixed compiler rather than a pure MLIR one.", "body_md": "# A Tour of XLA: Where MLIR Lives (and Where It Doesn't)\n\nXLA is the compiler under JAX, TensorFlow, and PyTorch/XLA, turning tensor programs into fast code for CPUs, GPUs, and TPUs. This is a tour of its architecture through one particular lens — where it uses MLIR, and where it doesn’t — because XLA’s structure is mixed, and the mix is the part worth understanding.\n\nXLA predates MLIR by several years, and its optimizer core is still its own representation: classic HLO, a hand-built C++ IR with its own passes, serialization, and text format. That core is not MLIR. MLIR appears, instead, at particular layers around it — the front door the frameworks target, the sharding system, and a growing share of code generation. So the useful question about XLA is not whether it “is an MLIR compiler” but which parts are and which parts aren’t — and that turns out to be specific enough to map.\n\nThat map is the rest of this post, walked from the way in to the way out and grounded in the `openxla/xla`\n\nsource tree with paths cited. It is a companion to [the tour of MLIR](/posts/mlir-dialect-stack-for-ml/), which covered MLIR as a general IR construction kit; here we look at one large, mature consumer of it and note how much shows up, and where.\n\n## The two IRs\n\nThe single most important architectural fact about XLA is that it has **two** intermediate representations, and only one of them is MLIR.\n\nThe core is **classic HLO** — “High Level Operations.” It is a set of hand-written C++ classes (`HloModule`\n\n, `HloComputation`\n\n, `HloInstruction`\n\n, in `xla/hlo/ir/`\n\n), with its own protobuf serialization, its own pass manager, and a text format that is *not* MLIR’s. 1 This is what the optimizer works on, and it reads like this — note the\n\n`f32[4,16]{1,0}`\n\nshape-and-layout notation, the `fusion(...)`\n\nops, no `func.func`\n\n, no `tensor<>`\n\ntypes:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\nphp\n%fused_computation.1 (param_0.2: f32[4,16], param_1.4: f32[16]) -> f32[4,16] {\n  %add.1 = f32[4,16]{1,0} broadcast(%param_1.4), dimensions={1}\n  %add.0 = f32[4,16]{1,0} add(%param_0.2, %add.1)\n  ROOT %max.0 = f32[4,16]{1,0} maximum(%add.0, %constant.0)\n}\nENTRY %main {\n  %ynn_fusion = f32[4,16]{1,0} fusion(%x, %W), kind=kCustom, calls=%fused_computation\n}\n```\n\nThe other IR is a family of **MLIR dialects**. XLA defines some in-tree — `MHLO`\n\n(the MLIR-native rendering of HLO) and the `IFRT`\n\n/`VIFRT`\n\ndialects for its multi-device runtime interchange 2 — and depends on others, chiefly\n\n**StableHLO** and its higher-level sibling\n\n**CHLO**. A\n\n`jit`\n\n-ed `relu(x @ W + b)`\n\n, straight from JAX’s `lower()`\n\n, comes out as StableHLO, and it is unmistakably MLIR — `func.func`\n\n, SSA `%`\n\nvalues, `tensor<...>`\n\ntypes, dialect-prefixed ops:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\n9\nfunc.func public @main(%arg0: tensor<4x8xf32>, %arg1: tensor<8x16xf32>,\n                       %arg2: tensor<16xf32>) -> tensor<4x16xf32> {\n  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] :\n         (tensor<4x8xf32>, tensor<8x16xf32>) -> tensor<4x16xf32>\n  %1 = stablehlo.broadcast_in_dim %arg2, dims = [1] : (tensor<16xf32>) -> tensor<1x16xf32>\n  %3 = stablehlo.add %0, %1 : tensor<4x16xf32>\n  %5 = stablehlo.maximum %3, %cst : tensor<4x16xf32>\n  return %5 : tensor<4x16xf32>\n}\n```\n\nSame computation, two worlds. And critically, there is a translation layer, `xla/hlo/translate/`\n\n, whose whole job is to move programs across the seam: `stablehlo_to_hlo`\n\n, `hlo_to_mhlo`\n\n, `mhlo_to_hlo`\n\n. 3 So the life of a program is: a framework hands XLA\n\n**StableHLO**(MLIR), XLA translates it\n\n*down into classic HLO*, optimizes it\n\n*as classic HLO*, and then translates\n\n*back into MLIR*to generate code. MLIR bookends the classic-HLO core — and, as we will see, there is a second entry path that skips the core entirely.\n\n``` php\ngraph TD\n  FW[\"JAX / TF / PyTorch\"]\n  FW -->|\"lower()\"| SHLO[\"StableHLO / CHLO<br/>(MLIR dialects)\"]\n  FW -.->|\"Pallas kernel<br/>(escape hatch)\"| PAL[\"Pallas\"]\n  HEL[\"Helion<br/>(PyTorch kernel DSL)\"] -.->|\"compiles to\"| PAL\n  SHLO -->|\"xla/hlo/translate\"| HLO[\"classic HLO<br/>(C++ HloModule — NOT MLIR)\"]\n  HLO -->|\"optimizer:<br/>algsimp · fusion · layout · buffers\"| HLO2[\"optimized HLO\"]\n  HLO2 -->|\"pointwise/reduction<br/>emitters\"| EM[\"MLIR emitters<br/>arith · scf · vector → llvm\"]\n  HLO2 -->|\"GEMM · conv (GPU)\"| LIB[\"cuBLAS / cuDNN<br/>(library custom-calls)\"]\n  HLO2 -->|\"matmul/softmax<br/>fusions (GPU)\"| TRIT[\"Triton (tt dialect)\"]\n  HLO2 -->|\"TPU\"| MOS[\"Mosaic\"]\n  PAL -->|\"GPU\"| TRIT\n  PAL -->|\"GPU / TPU\"| MOS\n  EM --> OUT[\"LLVM IR → obj / PTX\"]\n  TRIT --> OUT\n  MOS -.->|\"TPU\"| LLO[\"LLO → libtpu\"]\n\n  classDef mlir fill:#e1f5ff,stroke:#4a6572,color:#1a1a1a;\n  classDef classic fill:#fff3cd,stroke:#7a6a00,color:#1a1a1a;\n  classDef lib fill:#f3e5f5,stroke:#6a4a6a,color:#1a1a1a;\n  class SHLO,PAL,EM,TRIT,MOS mlir;\n  class HLO,HLO2 classic;\n  class LIB,LLO,OUT lib;\n```\n\nThe blue is MLIR; the yellow, in the middle, is not; the purple is the vendor libraries and the final targets. The rest of this post walks the blue, roughly left to right.\n\n## The front door: CHLO, MHLO, StableHLO\n\nThe way into XLA is MLIR, and it is a small stack of related dialects under `xla/mlir_hlo/`\n\n. 4 From high to low:\n\n**CHLO**(“client HLO”) is the most permissive — implicit broadcasting, a few composite ops, the conveniences a frontend wants — and lowers to MHLO (\n\n`map_chlo_to_hlo_op.h`\n\n). **MHLO** is the MLIR-native form of HLO’s operation set.\n\n**StableHLO** is MHLO with a stability contract — versioning, backward/forward compatibility — so it can be the portable interchange\n\n*between*frameworks and compilers. This is what\n\n`jit(...).lower()`\n\nemits and what a framework commits to.\n\n[5](#fn:5)Here is a detail worth noting about the interchange itself: **StableHLO’s syntax is MLIR’s.** The StableHLO specification says so in as many words — *“StableHLO syntax is heavily inspired by MLIR, which is not necessarily the most ergonomic alternative, but is arguably the best fit for StableHLO’s goal of creating more interoperability between ML frameworks and ML compilers.”* 6 The op grammar is MLIR’s op grammar (\n\n`OpName ::= '\"' 'stablehlo' '.' OpMnemonic '\"'`\n\n), and the function-signature form is, per the spec, “deliberately part of StableHLO syntax for compatibility with MLIR.” You can see it in an op that carries a region — a reduction-style op whose body is a nested block, written exactly as MLIR writes regions and blocks:\n\n```\n1\n2\n3\n4\n5\n6\n7\n%result = \"stablehlo.select_and_scatter\"(%operand, %source, %init_value) ({\n  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):\n    %0 = \"stablehlo.compare\"(%arg0, %arg1) {\n      comparison_direction = #stablehlo<comparison_direction GE>\n    } : (tensor<i32>, tensor<i32>) -> tensor<i1>\n    \"stablehlo.return\"(%0) : (tensor<i1>) -> ()\n}) : ...\n```\n\nThe `^bb0`\n\nblock, the region in `({ ... })`\n\n, the `#stablehlo<...>`\n\nattribute, the trailing type signature — none of that is XLA-specific; it is MLIR’s generic op syntax. So StableHLO uses not just MLIR as a library but MLIR’s textual and structural conventions as its on-disk interchange format. This is the front door the frameworks emit, and the one XLA translates through to reach its classic-HLO core.\n\n## The sharding layer: Shardy\n\nOff to the side of the optimizer sits an entire MLIR-based subsystem for distribution: **Shardy**, dialect `sdy`\n\n, integrated under `xla/service/spmd/shardy/`\n\n. 7 It is the newer, MLIR-native partitioner superseding GSPMD, propagating shardings and inserting collectives. What is worth showing here is that it is a\n\n*real dialect*, with its own attribute syntax for meshes and shardings:\n\n```\n1\n2\n3\n// a named device mesh, and a value sharded across its axes\nsdy.mesh @mesh_xyz = <[\"x\"=2, \"y\"=2, \"z\"=2]>\n%arg0 : tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_xyz, [{\"x\"}]>}\n```\n\nThat `#sdy.sharding<@mesh, [{\"x\"}, {\"y\"}]>`\n\nattribute is exactly the kind of thing MLIR’s extensibility is *for*: a domain that needs a first-class, verifiable representation of something the built-in dialects do not model — here, how a tensor’s dimensions map onto a device mesh. Architecturally, the giveaway of how Shardy is wired in is that its directory is full of *round-trip* pipelines — `stablehlo_round_trip`\n\n, `sdy_round_trip`\n\n— because sharding information has to be carried back and forth between the `sdy`\n\nMLIR world and the StableHLO/HLO the rest of XLA speaks. 7 XLA reached for MLIR here precisely because sharding annotations are gnarly enough to deserve a dedicated dialect with real verification, and it pays a round-tripping tax to bridge back to the classic core.\n\n## The way out: emitters, cuBLAS/cuDNN, Triton, and Mosaic\n\nCode generation is where MLIR does the most work in modern XLA, and where the “vote for MLIR” is most visible. But the GPU way out is not one path — it is a fan-out, and only some of the branches are MLIR.\n\n**The emitters framework (MLIR).** The old code generator was a hand-written `IrEmitter`\n\nthat walked HLO and called the LLVM `IRBuilder`\n\ndirectly. It is being replaced by an **emitters framework** built on MLIR, now used by *both* the CPU and GPU backends (`xla/codegen/emitters/`\n\n, `xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc`\n\n, `xla/backends/gpu/codegen/emitters/`\n\n). 8 The path is textbook progressive lowering: a fusion’s HLO is converted to MLIR by\n\n`elemental_hlo_to_mlir`\n\n, producing a mix of standard dialects — the emitters pull in `arith`\n\n, `math`\n\n, `scf`\n\n, `affine`\n\n, `vector`\n\n, `tensor`\n\n, `gpu`\n\n, and `mhlo`\n\n— and that is lowered, through the emitters’ own transform passes (`vectorize_loads_stores`\n\n, then `lower_to_llvm_gpu`\n\n/ `lower_to_llvm_common`\n\n), down to the **and out to LLVM IR.**\n\n`llvm`\n\ndialectThis is exactly the\n\n[8](#fn:8)`linalg`\n\n/`scf`\n\n/`vector`\n\n→ `llvm`\n\nstory from the MLIR tour, happening inside XLA.**cuBLAS and cuDNN (not MLIR — vendor libraries).** For the operations that vendors have already tuned to the metal, XLA does not generate code at all; it calls the library. Dedicated rewriter passes turn matmuls and convolutions into `custom-call`\n\ns targeting the NVIDIA libraries — the targets are spelled out in `xla/service/gpu/cublas_cudnn.h`\n\n: `\"__cublas$gemm\"`\n\n, `\"__cublas$lt$matmul\"`\n\n(and F8/MX variants), `\"__cudnn$convForward\"`\n\n, `\"__cudnn$convBiasActivationForward\"`\n\n, and more. 9 Disable the Triton GEMM path and a matmul falls back to exactly one of these:\n\n```\n1\n2\n%custom-call = (f32[4096,4096]{1,0}, s8[4194304]{0}) custom-call(%a, %b),\n               custom_call_target=\"__cublas$lt$matmul\"\n```\n\nThis is a real and important branch of the “way out”: a large fraction of a model’s FLOPs, on GPU, never touch MLIR or LLVM at all — they are a `custom-call`\n\ninto a closed vendor library. MLIR owns the *fusions around* the GEMMs, not usually the GEMMs themselves.\n\n**Triton (MLIR).** For matmul-and-softmax-shaped fusions on GPU, XLA emits the `tt`\n\n(Triton) dialect (`xla/backends/gpu/codegen/triton/`\n\n), which the Triton compiler — itself MLIR — takes to LLVM and PTX. 10 The kernel XLA emits is ordinary Triton IR:\n\n```\n1\n2\n3\n4\n5\n6\n7\n8\ntt.func @triton_fn(%arg0: !tt.ptr<f32>, %arg1: !tt.ptr<f32>) {\n  %r  = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32>\n  %p  = tt.addptr %base, %off : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32>\n  %x  = tt.load %p : tensor<1024x!tt.ptr<f32>>\n  ...\n  tt.store %q, %y : tensor<1024x!tt.ptr<f32>>\n  tt.return\n}\n```\n\n**Mosaic (MLIR).** On TPU, the codegen dialect is **Mosaic**, another MLIR dialect (with its own TPU memory spaces, vector layout, and systolic-array ops). So there is not one MLIR codegen path but three — the general emitters, Triton, and Mosaic — chosen by target and fusion shape, plus the non-MLIR cuBLAS/cuDNN library path alongside them.\n\n## The other door: Pallas\n\nEverything so far is the main road: framework → StableHLO → classic-HLO optimizer → codegen. There is a second entrance that skips the optimizer almost entirely. **Pallas** is JAX’s kernel language, and a Pallas kernel does not go through HLO fusion and layout assignment — it is compiled straight to a device kernel and dropped into the surrounding program as a `custom-call`\n\n. 11 The architectural point here is that Pallas is a\n\n*second producer of the MLIR codegen dialects*, reaching them without the classic-HLO core:\n\n- On GPU,\n`JAX → Pallas → Triton (tt) → LLVM → PTX`\n\n. (In recent JAX, Pallas-Triton even has its own compilation pipeline rather than delegating PTX generation to XLA.)[11](#fn:11) - On TPU,\n`JAX → Pallas → Mosaic → LLO → libtpu`\n\n.\n\nFrom XLA’s perspective the whole kernel is one opaque `custom-call`\n\n— `@__gpu$xla.gpu.triton`\n\nfor the Triton path — with its launch grid attached; XLA schedules around it but does not look inside. So the dotted arrow in the diagram is not a detail: Pallas is a deliberate hole punched through the classic-HLO core so that a hand-written kernel can reach the same MLIR codegen backends (Triton, Mosaic) that XLA’s own emitters use. The same MLIR dialects, two producers.\n\nAnd the door has become a doorway others walk through. PyTorch, working with Google, shipped a TPU backend for **Helion** — its high-level kernel-authoring DSL — that “compiles Helion kernels to Pallas.” 12 So a PyTorch author writing a single Helion kernel reaches XLA’s Mosaic TPU backend\n\n*through*JAX’s Pallas —\n\n`Helion → Pallas → Mosaic → TPU`\n\n— with the explicit aim of keeping “the same set of kernels across TPU and GPU.” The escape hatch is now a convergence point: two frameworks and two kernel DSLs reaching one shared set of MLIR codegen backends. Triton and Mosaic are, in other words, not only XLA’s internal codegen path but a common target reached from several directions.To collect the blue regions in one place:\n\n| Where | MLIR piece | Role |\n|---|---|---|\nFront door | CHLO / MHLO / StableHLO | portable interchange; syntax is MLIR’s\n|\nSharding | Shardy (`sdy` ) | dialect for mesh/sharding; round-trips to HLO\n|\nCPU/GPU codegen | emitters (`arith` /`scf` /`vector` /`llvm` ) | HLO → MLIR → LLVM IR\n|\nGPU GEMM/conv | none — cuBLAS/cuDNN | library `custom-call` s\n|\nGPU fusions | Triton (`tt` ) | matmul/softmax → PTX\n|\nTPU / Pallas | Mosaic | kernels → TPU\n|\nRuntime interchange | IFRT / VIFRT | multi-device array IR\n|\nThe optimizer core | none — classic C++ HLO | algsimp, fusion, layout, buffers\n|\n\n## What XLA depends on\n\nRead as a dependency graph, XLA is downstream of a lot of shared infrastructure. Its Bazel `third_party/`\n\npulls in the whole `llvm-project`\n\nmonorepo (where MLIR itself lives), plus the `openxla`\n\nsatellites and Triton, each of which is *itself* MLIR-based.[13](#fn:12)\n\n```\ngraph TD\n  LLVM[\"llvm-project<br/>(LLVM + MLIR core)\"]\n  XLA[\"XLA<br/>classic-HLO core + MLIR edges\"]\n  SHLO[\"stablehlo repo\"]\n  SHARDY[\"shardy repo\"]\n  TRITON[\"triton repo\"]\n\n  LLVM --> XLA\n  LLVM --> SHLO\n  LLVM --> SHARDY\n  LLVM --> TRITON\n  SHLO --> XLA\n  SHARDY --> XLA\n  TRITON --> XLA\n\n  classDef core fill:#e1f5ff,stroke:#4a6572,color:#1a1a1a;\n  classDef leaf fill:#f3e5f5,stroke:#6a4a6a,color:#1a1a1a;\n  class LLVM core;\n  class XLA,SHLO,SHARDY,TRITON leaf;\n```\n\nMLIR is not a dependency XLA could take or leave — it arrives through `llvm-project`\n\nand again through every satellite (StableHLO, Shardy, and Triton are all MLIR-based). The one part of XLA that owes nothing to that ecosystem is the classic-HLO core, which is exactly why it is the part that has not been rewritten: it works, it is mature, and MLIR arrived after it was already load-bearing.\n\n## Why the hybrid, and what it costs\n\nThe honest description of XLA today is a mature pre-MLIR optimizer with MLIR bolted to its inputs, outputs, and sharding, plus a Pallas side-door to the MLIR backends, migrating inward at the rate a heavily-used production compiler can tolerate. That is what incremental adoption looks like when a rewrite is off the table. It has costs worth stating plainly.\n\n**Two IRs means translation, and translation means round-trips.** The `xla/hlo/translate/`\n\nlayer and Shardy’s `*_round_trip`\n\npipelines exist solely to shuttle programs across the classic-HLO/MLIR seam. Every round-trip is code to maintain and a place for information to be lost or subtly changed.\n\n**The optimizer cannot use MLIR’s machinery.** The pattern rewriter, the dialect-conversion driver, PDL/DRR declarative rewrites — the reusable transformation infrastructure the MLIR tour praised — none of it is available to the classic-HLO passes, because they operate on C++ HLO objects, not MLIR ops. XLA maintains its own pass manager and rewriting for HLO in parallel with the MLIR one it uses in the emitters. Two of everything.\n\n**“Built on MLIR” means different things for different parts.** The interchange and the codegen are MLIR; the optimizer is not. That distinction predicts what you can and cannot do: you can retarget the front door with StableHLO and write kernels that reach the same MLIR backends through Pallas, but you cannot express an HLO optimization as an MLIR pattern.\n\n## Where it stops\n\nWhat MLIR provides XLA is concrete and easy to list: a portable, versioned front door whose syntax the ecosystem shares, a dialect for the hard problem of sharding, a reusable progressive-lowering path to LLVM, and a set of codegen backends (Triton, Mosaic) that a hand-written Pallas kernel can reach too. It is what lets XLA retarget, and share codegen rather than hand-write it. But it is worth ending where the MLIR tour did. MLIR is plumbing: it gives XLA the machinery to *represent* StableHLO, `sdy`\n\nshardings, and TPU memory spaces, and to lower them. It does not decide that a tensor should live on a particular device, or in VMEM rather than HBM, and it does not put any of that in the *type* of the value. Those remain, in XLA as everywhere, annotations and dialect attributes checked by passes — the `sdy`\n\nsharding on an op, the memory space on a Mosaic `Ref`\n\n— rather than properties the type system enforces where you wrote the program. That gap is the one I keep coming back to.\n\nThe map, in one line: StableHLO in, Shardy to the side, emitters and Triton and Mosaic (and cuBLAS/cuDNN) out, Pallas and Helion as side-doors to the same backends — and a mature C++ HLO optimizer still in the middle, where most of the compiling happens. That is the shape of XLA today: a pre-MLIR optimizer at the center, MLIR at the interfaces around it — inbound interchange, sharding, and code generation — with the classic HLO core the one place MLIR has not reached. How much MLIR is in XLA, in the end, depends on which layer you are standing in.\n\n## References\n\n*Disclaimer: This article was generated using the Claude Opus 4.8 model.*\n\n**Classic HLO — the C++ core.**`HloModule`\n\n/`HloComputation`\n\n/`HloInstruction`\n\nand the optimizer that runs on them; not MLIR. ([xla/hlo/ir](https://github.com/openxla/xla/tree/main/xla/hlo/ir),[XLA architecture](https://openxla.org/xla/architecture))[↩︎](#fnref:1)[↩︎](#fnref:1:1)2**In-tree MLIR dialects: MHLO, IFRT, VIFRT.** MHLO is the MLIR rendering of HLO; IFRT/VIFRT are the multi-device runtime interchange dialects. ([xla/mlir_hlo](https://github.com/openxla/xla/tree/main/xla/mlir_hlo),[xla/python/ifrt](https://github.com/openxla/xla/tree/main/xla/python/ifrt))[↩︎](#fnref:2)[↩︎](#fnref:2:1)2**The translation layer.**`xla/hlo/translate/`\n\nbridges StableHLO ↔ MHLO ↔ classic HLO (`stablehlo_to_hlo`\n\n,`hlo_to_mhlo`\n\n,`mhlo_to_hlo`\n\n). ([xla/hlo/translate](https://github.com/openxla/xla/tree/main/xla/hlo/translate))[↩︎](#fnref:3)**CHLO / MHLO under** The MLIR-HLO dialect family and its passes (`xla/mlir_hlo`\n\n.`map_chlo_to_hlo_op.h`\n\n,`mhlo`\n\ntransforms). ([xla/mlir_hlo](https://github.com/openxla/xla/tree/main/xla/mlir_hlo))[↩︎](#fnref:4)**StableHLO.** The portable, versioned MLIR interchange dialect; what JAX`lower()`\n\nemits and XLA imports. ([openxla.org/stablehlo](https://openxla.org/stablehlo))[↩︎](#fnref:5)[↩︎](#fnref:5:1)2**StableHLO syntax is inspired by MLIR.**“StableHLO syntax is heavily inspired by MLIR … arguably the best fit for StableHLO’s goal of creating more interoperability between ML frameworks and ML compilers”; the op grammar and function-signature form are MLIR’s. ([StableHLO spec, Operations](https://github.com/openxla/stablehlo/blob/main/docs/spec.md#operations))[↩︎](#fnref:6)[↩︎](#fnref:6:1)2**Shardy (** MLIR-based partitioner integrated under`sdy`\n\n).`xla/service/spmd/shardy/`\n\n, with StableHLO/HLO round-trip pipelines;`#sdy.sharding<@mesh, [...]>`\n\nattributes and`sdy.mesh`\n\n. ([Repo](https://github.com/openxla/shardy),[sdy dialect](https://github.com/openxla/shardy/blob/main/docs/sdy_dialect.md),[in XLA](https://github.com/openxla/xla/tree/main/xla/service/spmd/shardy))[↩︎](#fnref:7)[↩︎](#fnref:7:1)2[↩︎](#fnref:7:2)3**The emitters framework.** MLIR-based CPU/GPU codegen replacing the hand-written`IrEmitter`\n\n:`elemental_hlo_to_mlir`\n\n→`arith`\n\n/`scf`\n\n/`vector`\n\n/`affine`\n\n/`gpu`\n\n/`mhlo`\n\n, lowered via`lower_to_llvm_*`\n\nto the`llvm`\n\ndialect. ([xla/codegen/emitters](https://github.com/openxla/xla/tree/main/xla/codegen/emitters))[↩︎](#fnref:8)[↩︎](#fnref:8:1)2[↩︎](#fnref:8:2)3**cuBLAS/cuDNN as** Target names in`custom-call`\n\ns.`cublas_cudnn.h`\n\n(`__cublas$gemm`\n\n,`__cublas$lt$matmul`\n\n,`__cudnn$convForward`\n\n, …);`gemm_rewriter`\n\n/`conv_rewriter`\n\nlower matmul/convolution to them. ([cublas_cudnn.h](https://github.com/openxla/xla/blob/main/xla/service/gpu/cublas_cudnn.h),[GPU architecture](https://openxla.org/xla/gpu_architecture))[↩︎](#fnref:9)[↩︎](#fnref:9:1)2**XLA:GPU and Triton.** The GPU backend emits the Triton (`tt`\n\n) dialect for matmul/softmax fusions. ([xla/backends/gpu/codegen/triton](https://github.com/openxla/xla/tree/main/xla/backends/gpu/codegen/triton),[gpu architecture](https://openxla.org/xla/gpu_architecture))[↩︎](#fnref:10)[↩︎](#fnref:10:1)2**Pallas — the escape hatch.** JAX kernels compiled via Triton (GPU) or Mosaic (TPU) and embedded as a`custom-call`\n\n, bypassing HLO fusion/layout. ([Pallas design](https://docs.jax.dev/en/latest/pallas/design/design.html),[Mosaic/Triton lowering](https://docs.jax.dev/en/latest/pallas/index.html))[↩︎](#fnref:11)[↩︎](#fnref:11:1)2[↩︎](#fnref:11:2)3**Helion on TPU.** PyTorch’s high-level kernel DSL; its TPU backend “compiles Helion kernels to Pallas,” aiming to keep the same kernels across TPU and GPU. ([PyTorch blog](https://pytorch.org/blog/helion-on-tpu-towards-hardware-heterogeneous-kernel-authoring/))[↩︎](#fnref:13)**External dependencies.** XLA’s`third_party/`\n\npulls`llvm-project`\n\n(LLVM + MLIR), and the`stablehlo`\n\n,`shardy`\n\n, and`triton`\n\nrepos — each itself MLIR-based. ([xla/third_party](https://github.com/openxla/xla/tree/main/third_party))[↩︎](#fnref:12)\n\n[CC BY 4.0](https://creativecommons.org/licenses/by/4.0/)by the author.", "url": "https://wpnews.pro/news/a-tour-of-xla-where-mlir-lives-and-where-it-doesn-t", "canonical_source": "https://hiraditya.github.io/posts/a-tour-of-xla-where-mlir-lives/", "published_at": "2026-07-24 03:00:00+00:00", "updated_at": "2026-07-24 15:32:06.842244+00:00", "lang": "en", "topics": ["machine-learning"], "entities": ["XLA", "JAX", "TensorFlow", "PyTorch/XLA", "MLIR", "StableHLO", "CHLO", "openxla/xla"], "alternates": {"html": "https://wpnews.pro/news/a-tour-of-xla-where-mlir-lives-and-where-it-doesn-t", "markdown": "https://wpnews.pro/news/a-tour-of-xla-where-mlir-lives-and-where-it-doesn-t.md", "text": "https://wpnews.pro/news/a-tour-of-xla-where-mlir-lives-and-where-it-doesn-t.txt", "jsonld": "https://wpnews.pro/news/a-tour-of-xla-where-mlir-lives-and-where-it-doesn-t.jsonld"}}