Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Penguin Pass Roster and Pipeline Driver

All symbols, counts, and strings on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310 Cython modules under neuronxcc/starfish/penguin/; cp311/cp312 are layout-identical, only the Cython hash suffix differs). Treat every count as version-pinned — re-run the ls/strings/nm recipes below against any other wheel.

Abstract

The Penguin middle-end is not a monolithic optimizer. It is a catalog of compiled pass modules — one Cython .so per pass — plus a small set of driver generators that sequence a subset of them into an ordered pipeline at run time. There is no addPass(...); addPass(...) registry hard-coded in one function the way LLVM's PassBuilder works. Instead, each per-target CodeGenFlow driver is a Cython generator that yields pass specifications (PassType objects, each a (pass_class, kwargs, required|optional) triple) into a PassConstructorBuilder, which folds them via the Python +/+= operators into a list a PassManager then runs. This page is the index for that whole machine: the roster (which passes exist), the driver (which order they run in), the runner (DotTransform/PassManager/IterativePassManager), and the gate (OptLevel).

The roster splits cleanly by target-dependence. penguin/transforms/ holds 71 target-independent IR passes (DCE, value numbering, LICM, the simplifier family, delinearization, predicate resolution). penguin/targets/transforms/ holds 131 target-oriented passes shared across the Sunda and Tonga backends (allocation/spill, layout, tiling, collective lowering, DMA, type legalization, Tritium fusion). Two thin per-target packages — targets/sunda/passes/ (6) and targets/tonga/passes/ (40) — add the backend-specific ISel/tiling/scheduler passes. Counts are ls-verified below and differ from earlier "~70 / ~140" estimates; the numbers here are exact module counts excluding __init__.

The driver core lives in one unstripped, fully-decompiled file: targets/sunda/SharedCodeGenFlow.{py→.so}, banner "Driver functions for Sunda CodeGen". It exposes exactly four generator functions — peephole_optimizations, legalize_pe_insts, nki_codegen, and the centerpiece codegen_optimization_post_tritium_fusion — and both the Sunda and Tonga CodeGenFlow drivers import and call them. The reimplementation frame the rest of this page builds on: roster (catalog) → PassType specs → PassConstructorBuilder (assemble) → PassManager/DotTransform (run, with run_or_rollback safety) → OptLevel (skip whole stages at lower -O).

For reimplementation, the contract is:

  • The roster taxonomy — the four module directories, their exact counts, and how to tell a real Pass class from a helper/util module that ships in the same directory.
  • The spec/build/run separationPassType (a spec, not a pass), PassConstructorBuilder (folds specs via nb_add/nb_inplace_add), PassManager/IterativePassManager (the runners), all subclassing the DotTransform visitor base.
  • The driver sequencing — the four SharedCodeGenFlow generators, with codegen_optimization_post_tritium_fusion given as annotated pseudocode in source-order, and the per-pass required=/optional= qualifier on each.
  • The gating modelOptLevel as an inverted staged-skip enum (skip_middle_endskip_loop_level_transformationsskip_allocators), applied at the flow level, plus per-pass Options.clOptBool('<flag>') gates.
Driver filetargets/sunda/SharedCodeGenFlow.cpython-310-…so — banner "Driver functions for Sunda CodeGen"
Driver generatorspeephole_optimizations, legalize_pe_insts, nki_codegen, codegen_optimization_post_tritium_fusion
Spec builderPassConstructor.cpython-310-…soPassConstructorBuilder, PassType, build_pass_constructor, build_pass_list
Runner baseDotTransform.cpython-310-…soDotTransform (visitor base) + PassManager / IterativePassManager / SuperPass / BaseAnalysis
Opt-level gateir/OptLevel.cpython-310-…so@unique @total_ordering enum, custom __gt__
Target-independent passespenguin/transforms/71 modules (excl. __init__)
Target-oriented passespenguin/targets/transforms/131 modules (excl. __init__, excl. autotune/)
Sunda / Tonga pass packstargets/sunda/passes/ 6 · targets/tonga/passes/ 40
IR level in / outtensoriser-IR (op level) → ISel → <Target>ISAInst-level tensoriser-IR → BIR (5.18 BirCodeGenLoop)

The Roster — what passes exist

The roster is enumerable by ls. Every pass is its own Cython module; there is no shipped .py. The four directories below partition the catalog. Counts exclude __init__.cpython-310-….so (the package marker) and, for targets/transforms/, the autotune/ subpackage which is counted separately.

GOTCHA — module count is not pass count. Several modules in transforms/ and targets/transforms/ are helper libraries (DataflowUtil, IPUtils, LoopTransformUtils, TensorOpUtils, InstTransformUtils, TensorContractUtils, InstBuilder, IntrinsicBuilder, DataflowUtils) that export no Pass/DotTransform subclass. A reimplementer enumerating "passes" by directory listing over-counts by ~8–10. To get the real pass set, filter to modules whose top-level class subclasses DotTransform or BaseAnalysis (every real pass does — see the runner section). The directory counts below are module counts, given exact because they are the verifiable anchor; the pass-class subset is a few smaller.

Verification recipe

# from <wheel>/neuronxcc/starfish/penguin/
ls transforms/*.so          | grep -v __init__ | wc -l    # -> 71
ls targets/transforms/*.so  | grep -v __init__ | wc -l    # -> 131  (autotune/ is a subdir, not counted)
ls targets/sunda/passes/*.so| grep -v __init__ | wc -l    # -> 6
ls targets/tonga/passes/*.so| grep -v __init__ | wc -l    # -> 40

CORRECTION (ROSTER-1) — earlier estimates put transforms/ at "~70" and targets/transforms/ at "~140". The ls-verified counts are 71 and 131 respectively (excluding __init__). The "~140" figure conflated targets/transforms/ with the per-target tonga/passes/ (40) and sunda/passes/ (6) packs; 131 + 40 + 6 ≈ 177 target-oriented modules total across all three directories, of which 131 are the shared target transforms. The tonga pack is 40, not 41 — the 41st in the earlier list was __init__.

transforms/ — target-independent (71 modules)

These operate on op-level tensoriser-IR and never name a target. Grouped by role; names are the literal module stems (<Name>.cpython-310-….so), CONFIRMED by ls.

GroupModulesConfidence
Alias / dependencyAliasDependencyElimination AliasDependencyInduction AliasDependencyReset AliasDependencyVerificationPass AliasDependencyVerifier TransformMustAlias LegalizeOpLevelAlias OptimizeAliasedCopyChainCONFIRMED
DCE / DSE / VNDeadCodeElimination DeadStoreElimination ValueNumbering MemcpyElimination ZeroSizeTensorEliminationCONFIRMED
Op-eliminationCastOpElimination ConcatOpElimination ReshapeOpElimination SliceOpElimination TransposeOpElimination PadEliminationCONFIRMED
SimplifySimplifier SuperSimplifier SimplifySlice TensorOpSimplifier GenericAccessSimplifier IPSimplifier SimplifyPredicatesCONFIRMED
Loop / affineLICM LoopFusion FlattenLoop PerfectLoopNest EliminateDivs ModDivDelinear Delinearization DelinearizationBase ConcatDelinearizer AffinePredicateResolutionCONFIRMED
PredicateResolvePredicates ResolveComplicatePredicates PredicateAffineSelect MaskPropagation InsertConflictResolutionOpsCONFIRMED
Tensor-opLowerTensorOp LateLowerTensorOp TensorOpTransform RecognizeOpIdiom InferIntrinsic MutateDataTypeCONFIRMED
AnalysisRangeAnalysis LoadStoreDependencyAnalysis IPLiveInterval IPSubgraphTensorAnalysis AnalyzeKernel AxesGroup DynamicInstEstimator DetailDynInstCONFIRMED
Codegen / IR / debugNkiCodegenPass SimulatorPass PrintFunction DataflowViewer LayoutViewer DoNothing Region PatternMatch CanonicalizeIR HoistAllGatherCONFIRMED
Helper (no Pass class)DataflowUtil InstTransformUtils IPUtils LoopTransformUtils TensorContractUtils TensorOpUtilsSTRONG

NOTE — CanonicalizeIR and HoistAllGather live in transforms/ (not targets/), and several names the earlier roster placed here (IntegerSetAnalysis, IslSimplifier, KernelMetrics, InferInitValue, InferSharedMemLoc, InferNonlocalTensors) are not present as transforms/ modules in this wheel — the Infer* passes live in targets/transforms/, and the Isl/IntegerSet machinery is folded into the AffineExpr algebra rather than shipped as a standalone pass module.

targets/transforms/ — target-oriented (131 modules)

Shared across Sunda and Tonga. The high-population groups (counts are within-group, the 131 total is ls-verified):

GroupRepresentative modulesConfidence
Alloc / spillAllocateBlocks AllocationDecision AllocatorState BufferState GlobalSpillCtx MaxLiveSpiller FastSpillGeneration SpillPSum SpillFreeKernel StackAllocator ModuloAllocation IntervalMinCut AnnotateNoSpill NeuronStackLiveInterval CachedLiveInterval TensorLifetimeAnalysis TensorLiveRangeCONFIRMED
LayoutLayoutPreprocessing LayoutRequirementAnalysis LayoutTilingPipeline LocalLayoutOpt CycleBasedLayoutCostModel NKIKernelLayout PAGLayoutAnalysis PAGLayoutGreedySearchAlgorithm PAGLayoutHelpers OperatorLayoutVisualizerCONFIRMED
Tiling / vectorizeFlattenAxesForTiling FactorizeBlkDims FactorizeFreeDims LoopSplitting PartitionVectorization PackParDim PartialLoopFusion PartialSimdFusion VectorizeLoop VectorizeMatMult Vectorizer SFKVectorizer PGAnalysis PGAnalysisForTiling PGAnalysisHelpers PGTilingHelpers BFComputeCutting PComputeCuttingCONFIRMED
Collective (SPMD/CC)AGOrderingAnalysis BucketizeCCOp CCOpFusion CCOpAxesGroupAnalysis CCOpAxesGroupDelinearizer CoalesceCCOp FineGrainedCCOpFusion LowerCCOpBlockAxis TileCCOps LegalizeCCOpLayout LowerToSendRecv SimpleAllReduceTiling HoistFSDPCollectives RewriteReplicationMatmul SPMDCodeGen SpmdDCE TensorParallel ShardingUtils LowerShardAxis InferShardAxis InsertImplicitShardAxis ReinsertShardAxis RemoveShardedPartitionAxesCONFIRMED
DMADMALegalizer DMALocalityOpt AssignDMAQoSLabels DataStreaming DataLocalityOpt DramToDramTranspose InsertOffloadedTransposesCONFIRMED
Lowering / legalizeTargetLowering LowerIntrinsics LowerBroadcast LowerTranspose LegalizeType LegalizeSundaAccess LegalizeTongaAccess LocalLegalizeType RelaxPredicates TransformConvOp MacroGeneration SoftwarePipelineCodeGenCONFIRMED
Tonga-flavoredTongaCpyElim TongaDelinear TongaInstComb TongaIslSimplifier TongaLICM TongaLiveInterval TongaLoopFusion TongaLoopInterchange TongaSimplifier TongaSimplifyPredicatesCONFIRMED
Tritium fusionTritiumFusion TritiumFusionBase (the "post_tritium_fusion" pre-stage; autotuner host — see § autotuner)CONFIRMED
Type / castDTypeMutator EnforceAluDTAcc DemoteLargeTensors ExpandBatchNormCONFIRMED
Verify / sim / latencyNeuronVerifier UnshardBIRKernelForSimulation InstructionLatencyModel KernelLatencyInfoCONFIRMED
InferInferInitValue InferNonlocalTensors InferPSumTensor InferSharedMemLocCONFIRMED
MiscDeConcat MatMultCombine Rematerialization RewriteWeights RoundtripTranspose SimplifyMacroPredicates SimplifyTensor SoftmaxDivisionDelay StaticTransposeLocalTensor SplitAPUnionSets FlattenAPIndices TensorContract TensorInitialization InlineNKIKernels InsertCoreBarrier CanonicalizeDAGCONFIRMED

Per-target packs

targets/sunda/passes/  (6, excl __init__):
    InferIntrinsicOnCC  LateLegalizeInst  LegalizePartitionReduce
    LegalizeSundaMacro  SundaISel  SundaSizeTiling

targets/tonga/passes/  (40, excl __init__):
    AutoCastFP32  BroadcastWeights  ColoringAllocator  CommuteConcat  DMAProfiler
    ExpandISAMacro  FlattenMacroLoop  GlobalLayoutOpt  InferTongaTensor
    InsertIOTransposes  InsertLocalTransposes  IPGlobalLayoutOpt  LegalizePartitionTile
    LegalizeTongaMacro  LinearizeFreeDim  LowerPartitionTile  ParAxesAnnotation
    ParLayoutOpt  PartitionLocalityOpt  Recompute  ReshapeWeights  ResolveAccessConflict
    Scheduler  SimplifyTongaTensor  SplitAccGrp  SplitFreeDim  StaticProfiler
    TilingProfiler  TongaAccessAnalysis  TongaBufferUsageAnalysis  TongaEliminateCasts
    TongaISel  TongaLayoutAnalysis  TongaPerfectLoopNest  TongaSizeTiling
    TongaValueNumbering  TransformLayout  UnrollReduceMacro  VectorizeDMA  WeightCoalescing

NOTE — targets/cayman/passes/ and targets/core_v4/passes/ ship only __init__ — no per-pass modules. Those backends reuse the Sunda/Tonga packs; their own packs are placeholders in this build.

The Penguin→BIR / →NKI emit drivers live in targets/codegen/: BirCodeGenLoop (banner "BirCodeGen - Generate Backend IR from tensoriser IR at the TongaISAInst level."), NkiCodegen, CodeGenBase, DCE, DumpGraphAndMetadata. The autotune engine is targets/transforms/autotune/: Autotuner, _Compiler, _PerformanceMetric, _Search, _TreeSearch, _TreeSearchVisualizer.


The Runner — DotTransform, PassManager, IterativePassManager

Every pass — target-independent or not — subclasses one base, DotTransform, defined in DotTransform.cpython-310-….so. This is the shared visitor: the runner walks the comp-unit (Module → Function → Block/BasicBlock → Stmt → Macro/While/ScopeRegion) and the pass overrides the per-node hooks. It is the same base BirCodeGenLoop subclasses to emit BIR — so "lower to BIR" is structurally just another DotTransform walk.

Purpose

DotTransform provides the drive loop, the IR walk, per-statement hooks, the rollback-on-failure safety net, and the changed bookkeeping that IterativePassManager uses for fixpoint. Analyses use a sibling base, BaseAnalysis, with the same drive/rollback harness but non-mutating (analyzeModule/analyzeFunction instead of transform*).

Method Map

Method groupSymbols (DotTransform.<name>, CONFIRMED via strings)Role
DriverunOnModule runOnFunction run_ timed_run_ runTransformstop-level entry; timed_run_ wraps run_ with timing
IR walktransformModule transformFunction transformBasicBlock transformBlock transformScopeRegion transformStmts transformWhile transformMacro transformrecursive descent; transform is the leaf node hook
Per-stmt hooksbeforeStmtTransform afterStmtTransformoverridden by passes (and by BirCodeGenLoop — STRONG cross-bin link)
Safetyrun_or_rollback always_rollback run_with_exception_handling rethrow_exception verifya pass that fails verify or raises TransformError is rolled back
Bookkeepingchanged print_changed enter_pass_scope pass_name createPass should_skip_tensor _noopchanged drives the fixpoint loop
Debug / dumpdump_files dump_function dump_git set_cur_debug_location is_debug_mode LogContextper-pass dumps gated on debug mode
FlagsclOptBool clOptStringthe CLI-option accessors passes self-gate on

The runners

// DotTransform.PassManager — the top-level runner (CONFIRMED method symbols)
class PassManager:
    transformModule(module):     // run each spec's pass over the whole comp-unit
        for spec in self.pass_list:        // spec is a PassType (see next section)
            if not spec.isenabled(): continue
            pass = spec.construct(target)   // createPass — instantiate class with kwargs
            pass.run_or_rollback(module)    // mutate, or revert if verify() fails / raises
    transformFunction(fn): ...              // per-function variant
    verify(module): ...                     // run the registered verifiers
    dumpJson(out) / fromJson(in): ...       // the whole manager is JSON-serialisable

// DotTransform.IterativePassManager — run a pass set to FIXED POINT (CONFIRMED)
class IterativePassManager:
    transformFunction(fn):
        do:
            changed = False
            for spec in self.pass_list:
                pass = spec.construct(target)
                pass.run_or_rollback(fn)
                changed |= pass.changed       // wraps IterativeTransform.iterate / iterate_fn
        while changed                          // re-run until nothing mutates

// DotTransform.SuperPass — a composite that bundles a sub-pipeline AS ONE pass
class SuperPass(DotTransform):                 // used to group a phase (dump/print_changed)
    ...

QUIRK — run_or_rollback makes passes transactional. Unlike an LLVM pass (which is assumed correct and whose failure aborts the pipeline), a Penguin pass runs inside a try/verify wrapper: if the pass raises or its post-verify fails, the IR is reverted to the pre-pass state and the pipeline continues. This is why optional perf passes can be aggressive — a miscompiled tile is rolled back, not fatal. A reimplementation that drops the rollback will turn recoverable optional-pass failures into hard compile errors.


The Spec Builder — PassConstructor and PassType

A driver does not call manager.addPass(SomePass()). It yields a specification object and lets a builder fold it in. This indirection is what lets the whole pipeline be OptLevel-gated and JSON-serialised before any pass instance exists.

PassType — a spec, not a pass

A PassType wraps a pass class plus its kwargs, a required=/optional= qualifier, and an optional target requirement. It is enable/disable-toggleable (isenabled, enable, disable, is_required). The class is instantiated only when the PassManager actually runs the spec.

PassConstructorBuilder — fold via + / +=

PassConstructorBuilder (in PassConstructor.cpython-310-….so) defines the numeric nb_add and nb_inplace_add slots (CONFIRMED via nm -D), so a driver appends a spec with Python operator syntax:

// inside a CodeGenFlow generator (Cython generator body)
builder = PassConstructorBuilder(target=target)        // .__init__ takes target
builder += PassType(LowerTranspose,  required=True)    // nb_inplace_add — guarded:
builder += PassType(LateNeuronInstComb, optional=True) //   "Pass must be a PassType"
...
// build_pass_constructor() finalizes; build_pass_list() -> inner pass_constructor closure
pipeline = builder.build_pass_constructor(...)

Guard strings CONFIRMED in the binary: "Pass must be a PassType" (rejects a bare class) and "PassConstructorBuilder does not have target but %s requires target" (a target=-tagged spec yielded into an untargeted builder fails loudly). The constructed pipeline is round-trippable: PassConstructor.dumpJson / .fromJson (inputJson/outputJson) serialise the entire spec list — so a pipeline can be cached, diffed, or introspected without instantiating a single pass.

NOTE — the += fold is why the driver functions are generators: each yield PassType(...) is consumed by a for spec in gen: builder += spec loop in the calling CodeGenFlow. The generator's source order therefore is the pipeline order — and because Cython interns each pass name as a __pyx_n_s_<Name> global fetched through a version-counter, sorting those lookups recovers the literal in-source order even from a stripped generator body.


The Driver — SharedCodeGenFlow's four generators

targets/sunda/SharedCodeGenFlow.cpython-310-….so is unstripped (full IDA decompile available) and defines exactly four module-level generator functions, CONFIRMED by their qualified-name and __pyx_scope_struct symbols:

peephole_optimizations               -> __pyx_scope_struct_1_peephole_optimizations
legalize_pe_insts                    -> __pyx_scope_struct__legalize_pe_insts
nki_codegen                          -> __pyx_scope_struct_2_nki_codegen
codegen_optimization_post_tritium_fusion
                                     -> __pyx_scope_struct_3_codegen_optimization_post_tritium_fusion

Both targets/sunda/CodeGenFlow.so and targets/tonga/CodeGenFlow.so from …targets.sunda import SharedCodeGenFlow and call these four. The Sunda driver exposes optimize / codegen_prepare / codegen_optimization and their _eager variants (define-by-run single-shot mode); Tonga exposes the three staged ones only (no _eager). All entry-function names CONFIRMED via strings.

peephole_optimizations() — local cleanup, all optional

function peephole_optimizations(builder):   // gen1; all yields optional=True
    yield PassType(FactorizeBlkDims,     optional=True)
    yield PassType(NeuronInstComb,       optional=True)   // instruction-combine
    yield PassType(NeuronValueNumbering, optional=True)
    yield PassType(NeuronInstComb,       optional=True)   // re-run after VN

The InstComb → VN → InstComb sandwich is the classic GVN-then-recombine shape: value numbering exposes redundancies that a second instruction-combine pass can fold. Pass-name tokens (FactorizeBlkDims, NeuronInstComb, NeuronValueNumbering) CONFIRMED in SharedCodeGenFlow.so's __pyx_n_s_ table.

legalize_pe_insts() — PE-engine instruction legalization

function legalize_pe_insts(builder):        // gen
    yield PassType(LowerTranspose,    required=True)   // transpose -> PE ops
    yield PassType(LowerBroadcast,    required=True)   // broadcast -> PE ops
    yield PassType(LateNeuronInstComb, optional=True)
    yield PassType(SplitAccGrp,       optional=True)   // split accumulation groups
    yield PassType(SpillPSum,         optional=True)   // PSUM-bank spill

The two required lowerings make every transpose/broadcast expressible on the PE (tensor) engine; the optional tail does late cleanup and PSUM spilling. All five names CONFIRMED in the binary.

nki_codegen() — emit NKI text

function nki_codegen(builder, nki_dl, nki_debug_mode):  // gen2
    yield PassType(NkiCodegenPass, required=True,
                   nki_dl=nki_dl,                 // debug-location kwarg
                   nki_debug_mode=nki_debug_mode) // debug-mode kwarg

A single required pass — NkiCodegenPass from penguin/transforms/ — with the debug-location (nki_dl) and debug-mode (nki_debug_mode) flags forwarded as kwargs. Both kwarg names CONFIRMED as __pyx_n_u_nki_dl / __pyx_n_u_nki_debug_mode unicode tokens.

codegen_optimization_post_tritium_fusion() — the main middle-end list

This is the canonical post-fusion optimization sequence, run after Tritium fusion and before the engine-legalize / NKI / BIR tail. Recovered by the version-counter name-trace (source order). The whole generator is itself gated by Options.clOptBool('enable_post_tritium_fusion_passes') (CONFIRMED __pyx_n_s_enable_post_tritium_fusion_passes).

function codegen_optimization_post_tritium_fusion(target, nki_dl, nki_debug_mode):  // gen3
    builder = PassConstructorBuilder(target=target)
    builder += nki_codegen(nki_dl, nki_debug_mode)          // chains 3c in

    builder += PassType(NeuronLICM,           optional=True)   // loop-invariant code motion
    builder += PassType(InferPSumTensor,      optional=True)   // mark PSUM-resident tensors
    if clOptBool('enable-local-legalize-type'):                // type legalization, local|global
        builder += PassType(LocalLegalizeType, required=True)
    else:
        builder += PassType(LegalizeType,      required=True)
    builder += PassType(WeightCoalescing,     optional=True)
    builder += PassType(LegalizeSundaAccess,  required=True)   // target=Sunda access legalize
    builder += PassType(RelaxPredicates,      required=True)
    builder += PassType(TensorInitialization, required=True)
    builder += PassType(NeuronSimplifyPredicates, optional=True)
    builder += PassType(ExpandISAMacro,       required=True)   // from tonga.passes
    builder += PassType(SimplifyNeuronTensor, optional=True)
    builder += PassType(DMALocalityOpt,       optional=True)
    builder += PassType(DataStreaming,        optional=True)
    builder += PassType(FastSpillGeneration,  optional=True)
    builder += PassType(SFKVectorizer,        optional=True)   // scalar/SIMD-free-kernel vectorize
    builder += PassType(FastSpillGeneration,  optional=True)   // re-run spill after vectorize
    if clOptBool('annotate-no-spill-hint'):
        builder += PassType(AnnotateNoSpill,  optional=True)
    builder += PassType(/* sunda.<pass> */,   required=True)   // attr on the `sunda` module (STRONG)
    builder += PassType(CoalesceCCOp,         required=True)   // collective-op coalescing
    if not clOptBool('disable-tiling-allreduce'):
        builder += PassType(SimpleAllReduceTiling, required=True)
    builder += PassType(InsertCoreBarrier,    optional=True)
    builder += PassType(DMAProfiler,          optional=True)   // from tonga.passes
    return builder.build_pass_constructor(...)

Every pass-name token above is CONFIRMED present in SharedCodeGenFlow.so's interned-name table (__pyx_n_s_NeuronLICM__pyx_n_s_DMAProfiler). The required/optional qualifier on each is the FINE gate of the opt-level model: required passes are correctness legalizations that always run; optional passes are the perf passes dropped at lower -O. The exact intra-list ordering is STRONG (from the version-counter trace, source order); the two FastSpillGeneration instances bracketing SFKVectorizer are the spill-vectorize-respill pattern, mirroring the InstComb sandwich in peephole_optimizations.

QUIRK — the list re-runs nki_codegen and FastSpillGeneration. A pass appearing twice in the spec list is not a bug: FastSpillGeneration runs once before vectorization and again after (vectorization creates new live ranges that may need spilling), and the leading nki_codegen sub-flow is chained in so the post-fusion list can itself drive NKI emission on the NKI path. A set-based reimplementation that de-dups specs by class will silently drop the second spill and miscompile spill-heavy kernels.


The Gate — OptLevel

ir/OptLevel.cpython-310-….so defines the optimization-level enum. It is @unique and @total_ordering with a custom __gt__ (all CONFIRMED as __pyx_n_s_unique / __pyx_n_s_total_ordering / OptLevel.__gt__). Docstring, verbatim: "The optimization level of an penguin function." (the less-optimization-toward-the-end ordering note is in the source comment).

This is an inverted enum: members later in the list run fewer optimizations. The CONFIRMED member tokens (consumed by NKICodeGenFlow.so):

default_level   — full optimization (the implicit top)
skip_middle_end
skip_loop_level_transformations
skip_allocators
prepare         — prepare-only (the bottom; least optimization)
                  (+ allow_none — sentinel for "no level set")

Two-level gating

Coarse (flow-level). The CodeGenFlow / NKICodeGenFlow drivers branch on if opt_level > OptLevel.skip_middle_end / … skip_loop_level_transformations / … skip_allocators and skip whole phases. CONFIRMED: OptLevel.__gt__ plus these three member tokens are consumed only by NKICodeGenFlow.so — opt-level decisions are made at the flow level, not buried per-pass. skip_allocators is also threaded as a kwarg through NKICodeGenFlow (CONFIRMED __pyx_n_s_skip_allocators).

// flow-level gating shape (STRONG — exact thresholds INFERRED)
function codegen_optimization(fn, opt_level):
    if opt_level > OptLevel.skip_middle_end:                  // run the middle-end at all?
        run TritiumFusion, fusion, simplify, …
    if opt_level > OptLevel.skip_loop_level_transformations:  // loop/tiling/layout phase?
        run LoopFusion, LayoutTilingPipeline, Vectorize*, ISel
        run SharedCodeGenFlow.codegen_optimization_post_tritium_fusion(...)
    if opt_level > OptLevel.skip_allocators:                  // register/buffer allocation?
        run StackAllocator | ColoringAllocator, AllocateBlocks
    run SharedCodeGenFlow.legalize_pe_insts(...)              // always (correctness)
    run SharedCodeGenFlow.peephole_optimizations(...)         // final cleanup

Fine (per-pass). Inside codegen_optimization_post_tritium_fusion, each spec's required=/optional= qualifier is the per-pass gate, plus the Options.clOptBool('<flag>') toggles. CONFIRMED flag strings: enable-local-legalize-type, annotate-no-spill-hint, disable-tiling-allreduce, enable-tritium-loopfusion, and the generator-level enable_post_tritium_fusion_passes.

GOTCHA — the level→pass cutoff table is not statically recoverable. The > comparisons live in CodeGenFlow.so (no IDA decompile for that module in this build). The staged-skip taxonomy above is firm (member names + __gt__ + the consuming flow are CONFIRMED), but the exact numeric OptLevel value each -O maps to, and which optional passes survive at each level, is INFERRED. The driver-CLI --optlevel "0".."3" (default "2", 3.10) maps to these enum members upstream; that map is not in this binary's recoverable text. Do not invent a cutoff table — gate on the named stage boundary, not a guessed integer.


Where the Autotuner Hooks In

The L-strand autotuner is not a pipeline stage. It is embedded in one middle-end pass: targets/transforms/TritiumFusion.cpython-310-….so (145 autotun* string hits — the only pass binary with heavy autotune integration; CONFIRMED strings | grep -ci autotun → 145). Mechanism:

  • AutotuneFusionPlanGenerator (CONFIRMED class) enumerates candidate matmul fusion plans: build, best_fusion_plan, best_fusion_plan_search, generate_fusion_plans_for_search_more_trip, force_tripcount, get_axes_info.
  • Plans are submitted via add_plans to Autotuner (autotune/Autotuner.so, docstring CONFIRMED verbatim: "Autotuner -- Automatically benchmark and tune decisions within penguin during compilation."), which benchmark-drives through _Compiler / _Search / _TreeSearch / _PerformanceMetric and returns the winner via apply_plan.
  • Gating: _autotune_enabled / _autotune_config; when disabled, a _BypassedAutotuner / add_no_change_plan short-circuits to a deterministic heuristic.

So at default opt the fusion decision is a deterministic heuristic; with -autotune it becomes a benchmarked search over a re-run of the same middle-end (each plan = a re-tiling). The autotuner sits inside codegen_optimization, at the TritiumFusion pass, before ISel — downstream passes are deterministic given the chosen plan.


Placement — entry and exit

HLO ──hlo2penguin (C-strand)──▶ Penguin Module / Function (tensoriser-IR, op level)
   │  Driver/Job layer (collectBackendPipeline) calls per-target CodeGenFlow as a backend Job
   ▼
CodeGenFlow.optimize():
   codegen_prepare       (canonicalize, infer, alias setup, DCE/DSE, simplify, dtype)
   codegen_optimization  (TritiumFusion[+autotune] → fusion → loop/tiling/layout →
                          vectorize → collective lower → ISel  ◀── op→ISAInst boundary
                          → SharedCodeGenFlow.codegen_optimization_post_tritium_fusion (§ driver)
                          → legalize_pe_insts → late lower → alloc → peephole → nki_codegen)
   ▼  IR now at <Target>ISAInst level (same Python object, mutated in place)
BirCodeGenLoop (P-strand) — a DotTransform subclass — walks it, emits BIR   [5.18]
   (NKI path: SharedCodeGenFlow.nki_codegen / targets/codegen/NkiCodegen emits NKI text instead)

There is no serialization between passes — it is one live Python object graph, mutated in place, with changed tracking and run_or_rollback reverting failed passes. ISel (SundaISel / TongaISel) is the IR-form boundary: op-level tensoriser-IR before, <Target>ISAInst-level after. The hand-off to BirCodeGenLoop is just-another-DotTransform-walk that happens to emit BIR rather than mutate.


Adversarial Self-Verification

The five strongest claims, re-challenged against the binary:

  1. Pass counts: transforms/ = 71, targets/transforms/ = 131, sunda = 6, tonga = 40. ls *.so | grep -v __init__ | wc -l in each directory. CONFIRMED, exact. This overturns the backing report's "~70 / ~140 / 41" — see CORRECTION (ROSTER-1). The 131 excludes the autotune/ subdir (6 modules) and __init__. Status: CONFIRMED (directory enumeration; pass-class subset is a few smaller because ~8–10 modules are helper libs, flagged inline).
  2. The four SharedCodeGenFlow driver generators exist and are generators. CONFIRMED: __pyx_scope_struct_{1,2,3,_} for peephole_optimizations/nki_codegen/codegen_optimization_post_tritium_fusion/legalize_pe_insts, plus __pyx_mdef_… method-defs and the banner "Driver functions for Sunda CodeGen". Status: CONFIRMED.
  3. codegen_optimization_post_tritium_fusion runs the §-driver list in that order. Every pass-name token (NeuronLICMDMAProfiler) is CONFIRMED in the interned __pyx_n_s_ table; the order is from the version-counter name-trace (source order) on the unstripped decompile — STRONG for intra-list ordering (the trace methodology is sound but the per-position index was not re-derived in this session; the membership and qualifiers are CONFIRMED). The generator-level enable_post_tritium_fusion_passes gate is CONFIRMED. Status: membership CONFIRMED, order STRONG.
  4. OptLevel is an inverted staged-skip enum with members prepare/skip_middle_end/skip_loop_level_transformations/skip_allocators/default_level/allow_none, @unique @total_ordering, custom __gt__. All tokens + docstring CONFIRMED via strings ir/OptLevel.so. The flow-level vs per-pass split is CONFIRMED (members consumed only by NKICodeGenFlow). The exact -O→member→pass cutoff table is INFERRED and flagged as such (the > comparisons are in IDA-skipped CodeGenFlow.so). Status: enum CONFIRMED, thresholds INFERRED.
  5. PassConstructorBuilder folds PassType specs via nb_add/nb_inplace_add; DotTransform/PassManager/IterativePassManager run them with run_or_rollback. CONFIRMED: nm -D shows nb_add/nb_inplace_add; guard strings "Pass must be a PassType" and "PassConstructorBuilder does not have target but %s requires target"; PassManager.transformModule/transformFunction/verify/dumpJson/fromJson + IterativePassManager.transformFunction + run_or_rollback/always_rollback/verify all present in DotTransform.so. Status: CONFIRMED.

Honest gaps: the codegen_prepare/codegen_optimization intra-phase ordering in § placement is STRONG (pass set CONFIRMED referenced; order from vocabulary + the SharedCodeGen internal order); the sunda.<pass> spec at position 18 of the §-driver list resolves to an attribute on the sunda module (name STRONG, not pinned); the Job → CodeGenFlow call site is in the compiled-opaque driver (INFERRED from 3.16/the pipeline).


ComponentRelationship
DotTransformShared pass-visitor base for every Penguin pass and BirCodeGenLoop
PassConstructorSpec builder — turns driver yields into a runnable, serialisable pipeline
OptLevelStaged-skip enum gating whole phases at the flow level
TritiumFusionThe autotuner host; the fusion stage the post-tritium list runs after
BirCodeGenLoopThe exit driver — a DotTransform subclass that emits BIR

Cross-References