HLO Misc & Cleanup Sweep

All addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310, neuronxcc/starfish/bin/hlo-opt). Other versions will differ.

Abstract

After the heavyweight collective, layout, and fusion passes have run, the Neuron HLO pipeline closes with a long tail of ~28 small passes — the plumbing tier. None of them is individually large enough to earn its own page, but collectively they do the unglamorous work that makes the IR backend-legal: they reshape the entry tuple, fold weight parameters and RNG to constants, peel optimization barriers, strip trivial single-device collectives, lower a handful of unsupported idioms (variadic argmax-reduce, offloaded dropout, scatter-partition), legalize device assignment on every collective, and flatten predicated conditionals to select. This page is the consolidated reference for that closer.

Two registration kinds dominate. Most passes are plain xla::HloPass-derived, whose virtual Run(HloModule*, …) (vptr+0x18) is the real entry. A second group are xla::OpExpanderPass subclasses: they share one Run at 0x29f0bb0 and supply two per-pass virtuals — InstructionMatchesPattern (vptr+0x28, the matcher) and ExpandInstruction (the rewriter). For an OpExpander pass the table below cites the matcher address with the ExpandInstruction body in parentheses, since the matcher is what the registry row points at but the Expand body is where the algorithm lives.

The pass-order numbers (#NN) are the registration index in RegisterHiloHloPasses and serve as a strong proxy for default-pipeline order: output/IO shaping runs early (#14–#16, #54, #59), algebraic peepholes and canonicalizers in the middle (#17, #24, #66, #67, #90, #109, #110, #55), and metadata/legalization/debug late (#20–#22, #52, #69, #70). Every Run/matcher address below comes from the per-function disassembly and demangles to the class it is attributed to; the Confidence column therefore grades the rewrite logic (how strong the callee and string evidence for the algorithm is), not the address.

For reimplementation, the contract is:

• The OpExpander dispatch split: one shared Run, per-pass matcher + ExpandInstruction. Drive the pass off the matcher predicate, not the opcode alone.
• The five passes with real algorithms — InlineWeights, ReplaceRng, DeviceAssignmentLegalization, DynamicSliceTranspose, RewriteScatterPartition, plus HiloConditionalToSelect — given as annotated pseudocode below.
• The two open gaps: PartitionEmbedingOps sharding mechanics and RewriteModuleDtype's from/to dtype pair are not string-anchored.

1. The 28-pass table

Run / matcher addr is the registry entry point. For OpExpander rows the ExpandInstruction body addr follows in parentheses. Every class identity comes from its demangled symbol; the Confidence column grades the rewrite logic.

NOTE — HiloConditionalToSelect (0x1eb0dd0) is the Neuron variant with worker xla::DoConditionalToSelect (xlaL21…, file-local). The stock xla::ConditionalToSelect::Run at 0x2956eb0 is also linked into the binary but is not in the Neuron registry — do not conflate. Likewise MetadataNaming (#21) is plain-XLA op-naming, distinct from aws_neuron_default_metadata (#65, an OpExpander that attaches default OpMetadata).

2. OpExpander dispatch

Eight passes in the closer are OpExpanderPass subclasses. They never override Run; the base Run at 0x29f0bb0 walks the module post-order and, for each instruction, calls the subclass virtuals.

xla::OpExpanderPass::Run (0x29f0bb0)          ── shared driver, all OpExpander passes
  for each computation, post-order each inst:
    if this->InstructionMatchesPattern(inst):  ── vptr+0x28, per-pass matcher
        new = this->ExpandInstruction(inst)     ── per-pass rewriter (the algorithm)
        computation->ReplaceInstruction(inst, new)

GOTCHA — the registry row for an OpExpander pass points at InstructionMatchesPattern, not Run. A reimplementer who treats that address as the pass driver will model the matcher predicate as the whole pass and miss the ExpandInstruction rewrite. The matcher returns a bool; the rewrite lives in the sibling ExpandInstruction body cited in §1.

OpExpander members of the closer: DynamicSliceTranspose (#55), ResolveSelfComparison (#67), ConstantSliceClampSimplifier (#90), RewriteScatterPartition (#93), ConstantSliceConvertCanonicalizer (#109), DynamicSliceReshapeCanonicalizer (#110). The ExpandInstruction addresses (0x1ea0570, 0x2002d00, 0x1f75b70, 0x20093c0, 0x1f77b10, 0x1f81270) are each a non-virtual function reached from the shared Run, and each carries its own demangled symbol.

3. Algorithm notes — the non-trivial passes

3.1 InlineWeights — Run 0x1ee8c60 (10176 B, 494 bb)

The most substantive pass in the closer. It compile-time-folds weight parameters into HLO constants by reading an out-of-band model directory. The disassembly references metadata.json, model_files, nlohmann, npy2literal, RemoveParameter, CreateConstant, and the NCC_INL diagnostic family.

function InlineWeights_Run(module, dir):              // 0x1ee8c60
    cfg     = nlohmann::json::parse(ifstream(dir + "/metadata.json"))   // basic_ifstream open/close
    weights = cfg["model_files"] / cfg["weights"]      // map<param-name, npy-path>; strtol on indices
    for each entry-parameter p with a weights[p] entry:
        lit = hilo::npy2literal(weights[p])            // _ZN4hilo11npy2literalE…
        require ShapeUtil::Compatible(lit.shape(), p->shape())   // else NCC_INL001 "mismatching layout for parameter"
        require numpy_dtype(lit) == p->element_type()            // else NCC_INL002 "numpy file type mismatch"
        p->ReplaceAllUsesWith( HloInstruction::CreateConstant(lit) )
        computation->RemoveParameter(p->parameter_number())      // checks 0 <= no < param_instructions_.size()
    // re-derive entry ComputationLayout now that parameters are gone
    module->set_entry_computation_layout( ComputeProgramShape() )   // entry_computation_layout_.has_value() CHECK

Diagnostics emitted: NCC_INL001, NCC_INL002, NCC_MOD003; hard errors "Corrupted input dir, model file does not exist!" and "Computation … has no parameter number …", plus the JSON type-error path ("type must be string, but is "). The npy loader is hilo::npy2literal; the JSON lib is the vendored nlohmann::json_abi_v3_11_3. The CreateConstant literals this pass plants — together with ReplaceRng's host-sampled literals — are what the later BIR const-materialization stage (Part 12.8) lowers into device constant buffers.

NOTE — RemoveParameter shifts every later parameter's parameter_number(). The loop must re-resolve parameter indices (or iterate highest-index-first) or it will delete the wrong parameter. The bounds check 0 <= no < param_instructions_.size() is the only guard the binary keeps.

3.2 ReplaceRng — Run 0x1f63280 (2821 B, 140 bb)

Replaces HLO rng ops with statically sampled constant tensors, so the backend never executes an RNG op. The disassembly shows a std::mt19937 engine (mersenne_twister_engine with the 1812433253 initialization multiplier), uniform_int_distribution, a std::random_device seed, CreateConstant, and CreateTuple.

function ReplaceRng_Run(module):                      // 0x1f63280
    rd  = std::random_device()                        // _M_getval / _M_fini
    gen = std::mt19937(rd())                           // mersenne_twister_engine<…,1812433253>
    for each inst in post-order:
        if inst->opcode() != kRng: continue
        shape = inst->shape()
        if shape.IsTuple():                            // (state, value) rng
            elems = []
            for i in range(shape.tuple_shapes_size()):
                lit_i = sample_literal(shape.tuple_shapes(i), gen)   // uniform_int_distribution into Piece::buffer
                elems.push( CreateConstant(lit_i) )
            new = CreateTuple(elems)
        else:
            new = CreateConstant( sample_literal(shape, gen) )
        parent->ReplaceInstruction(inst, new)

GOTCHA — the seed comes from std::random_device, so two compiles of the same module produce different folded RNG constants — compilation is non-deterministic at this pass unless an upstream seed override is wired in, and no such override was located in this binary. A reimplementer reproducing bit-for-bit golden outputs must pin the seed.

3.3 HiloConditionalToSelect — Run 0x1eb0dd0 (494 B) → DoConditionalToSelect 0x1eb1160 (2347 B)

Flattens a predicated 2-branch kConditional into a per-element select, removing control flow for the data-parallel TPB engines. Run walks the call graph (CallGraph::VisitNodes, "Running conditional-to-select pass"); the worker does the surgery. Its callees include HasSideEffect, CallInliner, MakeSelectHlo, and CreateTuple.

function DoConditionalToSelect(cond):                 // 0x1eb1160
    for each branch b in cond->branch_computations():
        if b->HasSideEffect():
            return  // "Not transforming conditional; branches have side effects:"
    CallInliner::Inline(cond->true_computation())      // splice both branches inline
    CallInliner::Inline(cond->false_computation())
    pred = cond->operand(0)
    results = []
    for each leaf element of cond->shape():
        p   = convert_to_pred(pred, leaf)              // CreateConvert / ChangeElementType to PRED
        sel = MakeSelectHlo(p, true_val(leaf), false_val(leaf))
        results.push(sel)
    new = (cond->shape().IsTuple()) ? CreateTuple(results) : results[0]
    parent->ReplaceInstruction(cond, new)

3.4 DynamicSliceTranspose — matcher 0x1ea02d0, Expand 0x1ea0570

OpExpander. Matcher: transpose(dynamic-slice(...)). The Expand pushes the transpose below the dynamic-slice so the slice operates on the original (pre-transpose) layout. Its callees are InversePermutation, PermuteDimensions, CreateDynamicSlice, and CreateTranspose.

function DynamicSliceTranspose_Expand(transpose):     // 0x1ea0570
    ds    = transpose->operand(0)                      // the dynamic-slice
    perm  = transpose->dimensions()                    // transpose permutation
    inv   = InversePermutation(perm)
    // permute the slice sizes and start indices into the operand's frame
    sizes = ShapeUtil::PermuteDimensions(inv, ds->dynamic_slice_sizes())
    starts = permute(ds->index_operands(), inv)
    new_ds = CreateDynamicSlice(ds->operand(0), starts, sizes)
    return CreateTranspose(new_ds, perm)               // transpose now applies to the smaller, sliced tensor

3.5 TransformVariadicReduce — Run 0x1f71950 (149 B) → maxArgMax 0x1f6c890 (20653 B)

Run is a 4-callee post-order dispatcher calling xla::maxArgMax on each instruction. maxArgMax is the 20 KB worker: nested match::AllOf patterns recognize a variadic kReduce that jointly computes a max value and its arg-max index (a to_apply() doing compare+select+index-update over a 2-tuple). It lowers this single variadic reduce into native ops: MakeReduceHlo(kMax) for the value, and an iota/compare/select arg-max chain (MakeCompareHlo with a comparison_direction(), MakeSelectHlo, MakeBroadcastHlo, LiteralUtil::MaxValue sentinel), re-tuples via CreateTuple, and ReplaceInstruction(inst, newTuple). The exact comparator topology — specifically the tie-break direction when two elements share the maximum — was not resolved [UNRESOLVED].

3.6 PartitionEmbedingOps — Run 0x1f59390 (3158 B, 171 bb)

A large body with no distinctive strings; callees are generic shape/instruction builders. By name and pipeline position (immediately after the output-shaping passes, before EliminateRedundantCompare) it partitions embedding/gather ops across the tensor-parallel degree, sharding the embedding-table operand and stitching the per-shard gather results. The exact gather→per-shard-gather construction is not string-anchored — the reading rests on the class name, body size, and pipeline position alone, and the rewrite mechanics remain open [INFERRED].

3.7 ReplaceMinimumConstant — Run 0x1f62990

Despite the name, this pass does not rewrite a minimum opcode. It rewrites a scalar constant — the −inf/lowest sentinel that feeds a min/clamp/dynamic-update-slice mask-fill — to the dtype's largest-negative finite value, so the masked positions survive downstream arithmetic that would otherwise propagate NaN/inf. The per-dtype finite floor comes from the callee hlo_utils::GetMinNonInfValue, wrapped by CreateConstant and applied by ReplaceInstruction. The matcher gate is shouldReplaceConstant.

GOTCHA — the "min" in ReplaceMinimumConstant names the sentinel value it replaces, not the opcode it matches. The instruction that gets rewritten is the scalar constant operand of a min/clamp/dynamic-update-slice, retyped to a finite floor; the min op itself is untouched.

3.8 DeviceAssignmentLegalization — Run 0x1f7c570 (6591 B, 350 bb)

Rewrites every collective's logical replica groups into a physical CollectiveDeviceList derived from the module's device assignment. Its callees and strings include IsCollective, GetGlobalDeviceReplicaGroups, GetParticipatingDevicesGroupsForSourceTargetPairs, CollectiveDeviceList, CreateAllGather, CreateCollectivePermute, and the "Unhandled Collective" FATAL path.

function DeviceAssignmentLegalization_Run(module):    // 0x1f7c570
    for each inst:
        if !neuron::IsCollective(inst): continue       // AG/AR/RS/AllToAll/CollectivePermute
        switch inst->opcode():
          case kCollectivePermute:
            devs = GetParticipatingDevicesGroupsForSourceTargetPairs(inst)
            new  = CreateCollectivePermute(inst->shape(), inst->operand(0),
                                           CollectiveDeviceList(devs), inst->channel_id())
          case kAllGather / kAllReduce / kReduceScatter / kAllToAll:
            CHECK(inst->operands().size() == 1)         // single-operand path
            devs = GetGlobalDeviceReplicaGroups(inst)
            new  = CreateAllGather/AllReduce/…(inst->shape(), inst->operand(0),
                                               CollectiveDeviceList(devs),
                                               inst->channel_id(), inst->to_apply())
          default:
            FATAL("Unhandled Collective instruction:")  // every collective opcode must be handled
        parent->ReplaceInstruction(inst, new)           // VLOG "Collective before/after Device Assignment Legalization:"

3.9 RewriteScatterPartition — matcher 0x2009240, Expand 0x20093c0 (2614 B)

OpExpander. The matcher delegates to anonymous-namespace pattern_one (0x2008f90) / pattern_two (0x20090e0), which walk an operand chain via mutable_operand and test opcode constants for a scatter-style partition built from collectives + reshapes around a loop: all-gather (4), all-reduce (7), collective-permute (0x1D), GTE (0x3A), reshape (0x5B), while (0x79). The Expand rebuilds it as a single CreateAllGather+CreateReshape, recomputing the replica-group list and casting through HloAllGatherInstruction. Its callees and strings include CreateAllGather, CreateReshape, HloAllGatherInstruction, ReplaceOperandWithDifferentShape, and the VLOG strings "Num groups", "Num per group", "reshapes tuple elemnt" (sic — the typo is in the binary).

function RewriteScatterPartition_Expand(inst):        // 0x20093c0
    match = pattern_one(inst) || pattern_two(inst)     // walk mutable_operand chain
    if !match: FATAL("Instruction does not match any known pattern")
    groups = recompute_replica_groups(inst)            // RepeatedField<long>::Reserve; "Num groups N"/"Num per group N"
    ag = CreateAllGather(target_shape, inst->operand(0), groups, channel_id)  // VLOG "Before/After all-gather"
    ag = static_cast<HloAllGatherInstruction*>(ag)
    rs = CreateReshape(final_shape, ag)                // "reshapes tuple elemnt"
    inst->ReplaceOperandWithDifferentShape(idx, rs)    // "After operand"
    return rs

4. Evidence summary

5. Limits of this reading

• PartitionEmbedingOps (§3.6): sharding mechanics are name-and-position only, with no string anchor.
• RewriteModuleDtype: the from/to dtype pair lives in a constructor literal with no accompanying string, so the specific pair is unknown.
• InjectNumericalErrors: the target-source flag or environment key was not located — only the "No targets specified." diagnostic.
• TransformVariadicReduce: tie-break direction of the arg-max comparator.
• ReplaceRng's internal Piece::buffer write is reconstructed from the literal-build call pattern rather than traced instruction by instruction.
• DeviceAssignmentLegalization's per-opcode switch is derived from the callee set; the exact branch ordering is reconstructed.
• The OpExpander matchers' opcode-value guards were read from callees rather than from decompiled control flow.

Related Components

Cross-References

• Pass Registry — Part 4.1: registry rows (Run/matcher addr, vtable, factory) for every pass tabulated here
• Collectives → Custom-Call — the collective-conversion pass family that DeviceAssignmentLegalization (#69) and the CC-removal passes (#32/#33) legalize and prune
• DUS/DS Simplifier & DynamicSlice Mover — the dynamic-slice pass family this page's OpExpander canonicalizers (#55/#90/#109/#110) belong to