The Compile Pipeline at a Glance

All symbols and addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310 Cython modules under neuronxcc/driver/). Other wheels differ; treat every address as version-pinned.

Abstract

A single neuronx_cc invocation is a process pipeline: the driver parses the compile arguments, builds an ordered list of Job names, and runs each Job in sequence — most of them by forking a standalone tool ELF and waiting on it. This page is the orientation map. It gives both views the rest of the book needs — the conceptual IR descent (HLO → Penguin → BIR → per-engine ISA → NEFF) and the actual job schedule the driver emits — and it marks where the naive reading of the two is wrong.

The job pipeline is assembled in one place: CompileCommand.buildPipeline (0x619d0, CompileCommand.py:1146–1349) concatenates three collector sub-lists — collectFrontendPipeline + collectWalrusPipeline + collectBackendPipeline — into a flat name list, which JobRegistry.makePipeline (0xbb80) turns into Job instances, and runPipeline (0x7a3a0) executes sequentially. The string list those three collectors return is the canonical compile order. There is no hidden scheduler; the pipeline shape is exactly what those three functions concatenate.

Two facts shape everything downstream and are easy to get wrong. First, the MLIR front half (hlo2penguin, emitted as the HLOToTensorizer job) runs first — before the in-process Frontend job and before the golden reference evaluator — so the "optimize HLO, then lower to Penguin" mental order does not match the process order. Second, WalrusDriver is a single Job whose ~18 backend stages (lsa, coloring, pre_sched, …) are forwarded as --walrus-passes to one walrus_driver process, not separate jobs.

The driver job pipeline

buildPipeline returns the concatenation of three collectors. The table is the default flow (no special flags, --optlevel 2); the Gate column says when each job is present. Job-name strings, tools, and gating py-lines are read from the decompiled CompileCommand / Pipeline / JobRegistry Cython modules.

GOTCHA — WalrusDriver is one Job, not a chain. collectWalrusPipeline (0x29e00, py 1655–1725) returns the single-element list ["WalrusDriver"]. The ~18 stage names it appears to contain — lsa, coloring, unroll, lower_ac, pre_sched, partitioner, bir_splitter, coloring_allocator / linear_scan_allocator, dma_optimization, anti_dependency_analyzer, vn_splitter, post_sched, tensorcopy_accel, dep_reduction, print_metrics, and finally birverifier / bir_sim — are accumulated into self.walrus_passes and handed to the one walrus_driver process as --walrus-passes. They are backend pass names, not driver Jobs. See Part 8.

GOTCHA — the Frontend job does not fork hlo-opt. Every job that forks a tool names that tool in its own module's string pool — hlo2penguin in jobs/HLOToTensorizer.so, walrus_driver in jobs/WalrusDriver.so, xla_infergoldens in jobs/XLAInferGoldens.so, hlo-neff-wrapper in jobs/NeffWrapper.so. The literal hlo-opt occurs in exactly one file in the whole wheel: the ELF neuronxcc/starfish/bin/hlo-opt itself (grep -rl --binary-files=text hlo-opt neuronxcc/). jobs/Frontend.so names no tool, carries no Popen/Job.shellCommand reference, and imports neuronxcc.starfish.penguin.Penguin with the entry names runXLAFrontend / runPenguin — it runs the HLO/penguin frontend in-process. hlo-opt is a standalone pass-driver ELF wired to no Job; the hlo-opt pass registry documents passes reachable through that ELF, not through F2.

GOTCHA — --optlevel does not add or remove top-level jobs. It is a string ("0".."3", default "2") that tunes the --walrus-passes set and the enable_internal_* flags the collectors read. Only indirectly — when it flips enable_internal_new_backend or selects the modular / bir-linker flow — does it change which jobs appear. The Backend sub-list (B0–B2) is suppressed entirely in the modular, bir-linker, and legacy-standalone flows; in the common modular path walrus_driver's output is consumed directly and Backend/Kelper never run.

The IR descent

Underneath the job schedule, the program descends six IR levels. The mapping to jobs is not one-to-one — the heavy HLO→Penguin lowering and HLO optimization are split across the forked HLOToTensorizer (hlo2penguin) job and the in-process Frontend job, and the exact division of HLO-level optimization between the two is not fully resolved from the binaries (both reach runXLAFrontend / runPenguin entry names). What is firm is the level sequence and the code that owns each.

  framework graph ──libneuronpjrt──►  XLA HLO  (HloModule protobuf)
                                          │
        ┌────────── hlo2penguin (F0, forked) + Frontend (F2, in-process) ──────┐
        │  HLO / MHLO / StableHLO / CHLO  →  Penguin Python IR (.py emission)   │
        │  collectives→custom-call, layout, quantize legalize, NeuronOpFusion,  │
        │  TensorizerLegalization, PenguinizeFunctions, MhloToPythonPrinter     │
        └──────────────────────────────────────────────────────────────────────┘
                                          │
   Penguin IR  (penguin.* Cython galaxy)  ◄────── NKI @nki.jit traces in here directly
        │  middle-end: layout, tiling, fusion, scheduling; ISL glue over stock islpy
        │
        │   ┌─ beta3 (live):    BirCodeGenLoop.cpython  →  bir::Instruction directly
        ▼   └─ beta2 (dormant on re-trace): KLR → KlirToBirCodegen (in libwalrus)
   BIR  (libBIR.so; bir::Module→Function→BasicBlock→Instruction; 110-opcode enum)
        │
        │   WalrusDriver job → walrus_driver (708 KB shim) → libwalrus.so (65 MB)
        ▼
   Per-engine ISA  (PE / Pool / Activation / DVE / SP — 64-byte bundles)
        │   NeffPackager (in libwalrus) / Kelper / NeffWrapper
        ▼
   NEFF on disk

The six levels, the dialect each speaks, and the part that owns the full detail:

NOTE — two routes into BIR. The live path in this build is beta3: BirCodeGenLoop.cpython lowers Penguin IR straight to bir::Instruction. The older beta2 route through KLR and KlirToBirCodegen (in libwalrus) is retained but dormant for internal kernel re-tracing here. A reimplementer needs only beta3 for current behavior; beta2 is documented for completeness in Part 7. The two converge on the same bir:: node.

Execution & failure model

runPipeline (0x7a3a0, py 1513–1572) initializes GlobalState, chdirs into the work directory (logging "Intermediate files stored in %s"), and drives Pipeline.runSingleInput (0xe030), which runs each Job under a timer and GlobalState scope ("Starting job %s" / "Finished job %s"). Each native Job is a subprocess.Popen(...).communicate() over its tool ELF. The whole subcommand runs inside a multiprocessing.Process that CommandDriver.run (0x1c510) joins; the child's exitcode is routed through handleError (0x30470) to one of three fatal diagnostics — [F134] (abnormal termination), [F137] ("forcibly killed", a negative exitcode, typically SIGKILL/OOM), or [F139] (a second abnormal-termination variant). The exact exitcode→diagnostic thresholds are a derived branch and remain [INFERRED]; the ownership in CommandDriver is read directly.

Two traps in the common mental model

GOTCHA — the job order is not "optimize then lower". collectFrontendPipeline (py 1636–1648) emits HLOToTensorizer (hlo2penguin) as F0, the golden evaluator as F1, and the in-process Frontend job as F2. Whatever HLO-level optimization Frontend performs is scheduled after the tensorizer job, not before it.

GOTCHA — Watchpoint, SaveTemps, and BIRVerifier are not collector jobs. Watchpoint appears only as a name-equality special-case in buildPipeline's post-makePipeline loop (_Pyx_PyUnicode_Equals(name, "Watchpoint") at 0x619d0); none of the three collectors emit it, and the other two are debug/aux jobs wired by separate paths. Building the pipeline from the three collectors will not produce any of them.

Cross-References

• Binary Inventory & the .so Map — the tools and libraries each job runs, with sizes and roles.
• Methodology & the Confidence Model — how the job order and the IR descent are evidenced, and what the confidence tags mean.
• 3.3 CompileCommand pipeline & the canonical job order — buildPipeline / runPipeline in full (planned).
• 3.4 JobRegistry & the sub-tool process model — name→class resolution and the forked-tool model (planned).
• Worked Example A — a matmul end-to-end — this descent followed for one operator (planned).