HLO Pre-Passes

Symbol names, VAs, and the build-id below apply to libtpu.so from the libtpu-0.0.40-cp314 wheel (build-id 89edbbe81c5b328a958fe628a9f2207d). Other versions differ; treat every VA as version-pinned.

Abstract

The HLO pre-passes are the ordered set of xla::HloPassInterface passes that run on libtpu's incoming XLA program before it leaves the HloInstruction world for MLIR. They scrub free-form frontend graphs into a normalized shape the rest of the compiler can assume: scheduling annotations are legalized, linear-algebra and RNG custom-calls are decomposed into primitive arithmetic, dynamic shapes are statically padded, sub-byte and high-precision dtypes are bracketed, sharding is propagated and the program is SPMD-partitioned per device, layout is assigned, and a final pass materializes TPU frontend_attribute strings into typed backend_config so the MLIR import sees nothing TPU-private. Every pass here consumes HLO and emits HLO; the handoff to MLIR is the next phase, owned by overview.md.

This page is the full ordered enumeration of that pre-pass set — the table that compile-phases.md delegates to. It fixes, for every pass: its C++ class, the builder function that adds it, the pipeline-builder phase it sits in (1–6, matching the compile-phases.md numbering), what HLO invariant it consumes, what HLO invariant it guarantees on output, whether it is TPU-specific or open-source XLA, and its byte-anchored entry symbol where one was recovered. The five builder functions live in the (anonymous namespace) of deepsea_compiler_hlo_passes.cc and are reached from DeepseaCompilerBase::RunHloPasses (0x1093a420), the body of the public C-ABI export TpuCompiler_RunHloPasses and of the separate-compilation entry xla::CompilePhase1HloOptimizations (0xf84ee00).

This is the enumeration page. It does not re-derive the top-level phase ordering or each phase's entry symbol — that spine lives on compile-phases.md and is linked, not duplicated. It does not re-derive per-pass transformation algorithms (the actual graph rewrites) — those have their own pages: algebraic-simplifier.md for TpuAlgebraicSimplifier, sharding-propagation.md / auto-sharding-spmd.md for the sharding flows. The 372-entry HloPassInterface RTTI catalog with name() strings lives on hlo-pass-registry.md.

For reimplementation, the pre-pass contract is:

• One container type. xla::HloPassPipeline is the only HLO pass container — there is no TPU-private pipeline class. Every TPU pass derives directly from xla::HloPassInterface (so RunImpl(HloModule*, execution_threads) returns absl::StatusOr<bool>) and is added with AddPass<T>(...).
• Six phases, five builders. The ordering decomposes into six pipeline-builder phases. Phases 1–5 are HLO→HLO; phase 6 is the final pre-MLIR HLO domain. Five decompiled builder functions add the passes: PreOptimizationPipeline, AddAutoShardingAndRelatedPasses, AddTpuPartitioningPasses, HloOptimizeThroughLayoutAssignment, PostOptimizationPipeline.
• The pre-passes legitimately introduce ops. Expanders emit new ops; the partitioner emits collectives. The acceptance gate (TpuHloSupportChecker) therefore runs in phase 4 after expansion and partitioning, not at the front — checking earlier would reject programs that are actually compilable.
• Fixed-point loops. Several passes are wrapped in xla::HloPassFix<P> and re-run to convergence; HloDCE and HloCSE are re-run between most stages whether or not wrapped.
• Invariant checkers re-run after every pass. MaybeAddInvariantCheckers (0x10944600) adds HloVerifier, LegalizeSchedulingAnnotations (as checker), and HloCycleDetection via the separate AddInvariantChecker<T> API at the head of every nested pipeline — they re-validate after each pass, not once.

Pipeline-Level Architecture

RunHloPasses (0x1093a420) is the monolithic driver. It builds and runs a sequence of xla::HloPassPipeline containers; the finer numbered stages map onto the five builder functions confirmed in the symbol table by their exact signatures:

RunHloPasses (0x1093a420)  ==  CompilePhase1HloOptimizations (0xf84ee00)
  ├─ PreOptimizationPipeline ........................ PHASE 1  input scrub
  │
  ├─ CreateHloPipeline (0x1093efe0)  opens the main pipeline
  │    ├─ AddAutoShardingAndRelatedPasses (0x10939c40) PHASE 2  sharding
  │    ├─ AddTpuPartitioningPasses (0x1278a440) ...... PHASE 3  SPMD prep
  │    ├─ HloOptimizeThroughLayoutAssignment (0x1094ada0) PHASE 4 pre-layout
  │    │    └─ LayoutAssignment  ........... (see layout-assignment.md)
  │    └─ PostOptimizationPipeline (0x1093fd40) ...... PHASE 5  post-layout
  │                                                    PHASE 6  pre-MLIR HLO
  └─ MLIR conversion  ..................... (see overview.md)

Two structural facts shape the whole set:

• xla::HloPassFix<P> wraps a pass in a fixed-point loop. Vtable evidence in the binary shows it instantiated for HloPassPipeline, HloDCE, AllReduceReassociate, ReduceWindowRewriter, ReduceScatterReassociate, WhileLoopConstantSinking, AllReduceReduceScatterReorder, jellyfish::TpuReduceWindowRewriter, and jellyfish::MosaicFusion. The fixed-point mechanism is CONFIRMED. Its crash-on-non-convergence behaviour is reportedly gated by --xla_tpu_crash_if_hlo_pass_fix_did_not_converge; that flag string was not located in the sampled strings table. [Confidence: LOW on the crash flag; CONFIRMED on HloPassFix.]

• Invariant checkers re-add at every nested pipeline head. MaybeAddInvariantCheckers (0x10944600) uses AddInvariantChecker<T> (not AddPass<T>), so each checker re-runs after every pass:

NOTE — decompile cross-check. All five builder signatures were confirmed verbatim in *_functions.json: PreOptimizationPipeline (lambdas $_0..$_8, predicate $_7 confirmed as the HloDomainIsolator argument), PostOptimizationPipeline (lambdas $_0..$_12), HloOptimizeThroughLayoutAssignment, AddAutoShardingAndRelatedPasses, AddTpuPartitioningPasses (takes PartitioningPipelineConfig), plus AddGatherScatterExpanderPasses, MaybeAddInvariantCheckers, and CreateHloPipeline. The four TPU vtables (TpuHloSupportChecker, jellyfish::TpuCallInliner, jellyfish::TpuInt2AutoUpDownCaster, jellyfish::TpuRngBitGeneratorTupleDecomposer) are present in *_names.json. [Confidence: CONFIRMED.]

The Ordered Pre-Pass Table

The master table below is the discovery-order enumeration recovered by walking each builder's AddPass<T> instantiations. The Phase column matches the pipeline-builder numbering on compile-phases.md: phase 1 = PreOptimizationPipeline, phase 2 = sharding, phase 3 = SPMD prep, phase 4 = HloOptimizeThroughLayoutAssignment (through layout assignment), phases 5–6 = PostOptimizationPipeline (post-layout refinement, then the final pre-MLIR HLO domain). Numbering is discovery order, not source contiguity; loop/fixed-point passes appear once but re-run.

The Src column: T = TPU-specific (xla::jellyfish::, xla::tpu::, or xla::Tpu* prefix); O = open-source XLA pass (some called with a TPU Target).

Phase 1 — PreOptimizationPipeline (input scrub)

O* = open-source class, TPU-specialized through a Target argument.

Phase 2 — AddAutoShardingAndRelatedPasses (sharding)

Phase 3 — AddTpuPartitioningPasses (SPMD prep)

Phase 4 — HloOptimizeThroughLayoutAssignment (through layout)

Layout assignment (Phase2PreLayoutAssignment / TpuLayoutAssignment) runs at the tail of phase 4 and is owned by layout-assignment.md.

Phases 5–6 — PostOptimizationPipeline (post-layout refinement, then pre-MLIR)

GOTCHA — the acceptance gate (#50) is in phase 4, not phase 1. A reimplementation that wants to reject unsupported ops early will be tempted to front-load TpuHloSupportChecker. libtpu does not: it runs after the expanders (phase 1) and the SPMD partitioner (phase 3) because those passes legitimately introduce new ops (expander outputs, partitioned collectives) that must themselves pass the check. Checking before expansion would reject compilable programs. The checker never mutates — it walks every HloComputation and validates each result Shape with ShapeUtil::ValidateShapeWithOptionalLayout, returning an error Status on the first unsupported shape. Entry RunImpl @ 0x11071480. [Confidence: CONFIRMED on RunImpl VA + vtable; HIGH that name() returns "tpu-hlo-support-checker" (the literal did not surface in the sampled strings table).]

TPU-Specific vs Open-Source Split

Of the 97 pipeline-mentioned passes, 33 carry an explicit TPU prefix (xla::jellyfish::, xla::tpu::, or xla::Tpu*); the remainder are open-source XLA passes, some parameterized with a TPU Target. The TPU-specific ones cluster by role:

• Custom-call decomposition (expanders): TpuCholeskyExpander, TpuQrExpander, TpuEighExpander, TpuTriangularSolveExpander, TpuGatherExpander, TpuScatterExpander, jellyfish::TpuRngBitGeneratorExpander (+ TupleDecomposer), jellyfish::RaggedAllToAllExpander, jellyfish::Infeed/OutfeedDecomposer.
• Shape / canonicalization rewriters: jellyfish::TpuBroadcastRewriter, jellyfish::TpuDegenerateDimensionRewriter, jellyfish::TpuReduceRewriter (flag + no-flag), HloPassFix<jellyfish::TpuReduceWindowRewriter>, xla::TpuGatherScatterFlattener, xla::TpuGatherSplit, jellyfish::TpuConvolutionTypeCanonicalizer.
• Algebraic / simplifier: jellyfish::TpuAlgebraicSimplifier (superset of the open-source simplifier — see algebraic-simplifier.md), PostFusionTpuSubgraphSimplifier, TpuTrivialFusionRemover, TpuTrivialInstructionUnfuser.
• Call-graph / acceptance: jellyfish::TpuCallInliner (must-fuse aware; name() → "tpu-call-inliner" / -must-fuse / -inner-must-fuse / -non-must-fuse), TpuHloSupportChecker.
• Sharding / partitioning: TpuAutoSharding (wraps xla::AutoSharding), jellyfish::TpuSpmdPartitioner, TpuPartitionAssignment, TpuSpmdConcatRewriter.
• Precision / dtype: jellyfish::TpuInt2AutoUpDownCaster, XPrecisionRewriter (x6/x9/x128 dot precision), TpuHloPrecisionTracer.
• Custom-call / fusion prep: jellyfish::MosaicFusion (Pallas/Mosaic kernels), UserGuidedFusionIdAssigner, WrapFusionOutputForDebug, TpuCustomCallScopedVmemAdjuster (confirmed AddPass in PostOptimizationPipeline::$_12).

A further set of TPU passes is present in the binary's RTTI (vtable-only) but was not observed inside any of the five decompiled builders — they are added by configuration-flag-gated branches deeper inside RunHloPasses. These include the SparseCore offload family (SparseCoreComputeOffloader, OffloadGatherToSparseCore, etc.), the async-collective family (TpuAsyncCollectiveCreator, AsyncCollectiveMerger, AsyncOpScheduler), the fusion family (TpuInstructionFusion, TpuMultiOutputFusion), FlashAttention, Gmm (group-matmul), and the int4 path TpuInt4Rewriter. [Confidence: CONFIRMED these classes exist (vtable RTTI); LOW on their pipeline position and gating.]

Custom-Call Targets Consumed at the HLO Boundary

The pre-passes consume / produce HLO kCustomCall ops with these target strings (recovered from the string section). They form the surface between the JAX/PyTorch frontend and the TPU compiler:

__cudnn$convForward and similar GPU targets are not consumed by any TPU pass — they are rejected by TpuHloSupportChecker (#50).

What Is Not Recovered

• Exact ordering of the vtable-only TPU passes. The SparseCore, async-collective, fusion, FlashAttention, and int4 families appear in RTTI but were not observed in the five decompiled builders; they live in flag-gated branches deeper inside RunHloPasses not walked in this pass. [Confidence: LOW on position.]
• Per-TpuVersion divergence. The decompilation reflects one pipeline that branches on Target / TpuCompilationEnvironment flags; which passes are skipped for, e.g., TPU v3 vs v6e was not isolated. [Confidence: LOW on per-gen differences.]
• Sharding-flow selector logic. Which of manual / auto / Shardy runs is decided by flags plus frontend-attribute detection inside the builders; the precise dispatch branch was not isolated (the auto_sharding flag family is CONFIRMED present in the string table, but the exact selector is LOW). See auto-sharding-spmd.md.
• Per-pass rewrite algorithms. This page recovers names + order + I/O invariants, not the internal graph rewrites. The simplifier algorithm has its own page (algebraic-simplifier.md); the rest are pending.
• The --xla_tpu_crash_if_hlo_pass_fix_did_not_converge flag string was not found in the sampled strings table. [Confidence: LOW.]

Cross-References

• compile-phases.md — the top-level ordered phase list and each phase's entry symbol (this page is the table it delegates to).
• overview.md — compiler orientation and the HLO → … → LLO IR-layer stack; owns the MLIR handoff that follows phase 6.
• hlo-ingestion.md — how the StableHLO/MHLO frontend program becomes the HLO module these pre-passes consume.
• hlo-pass-registry.md — the 372-entry HloPassInterface class catalog and name() strings.
• algebraic-simplifier.md — the TpuAlgebraicSimplifier rewrite algorithm (pass #60).
• sharding-propagation.md / auto-sharding-spmd.md — the phase-2/3 sharding and SPMD flows (passes #34–48).
• optimization-barrier.md — OptimizationBarrier handling across the pre-pass set.
• layout-assignment.md — TpuLayoutAssignment at the tail of phase 4.
• back to index