Control-Dependency Reification

All symbols and addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310). The HLO-side pass lives in neuronxcc/starfish/bin/hlo-opt; the MLIR simplifier and Penguin emitter live in neuronxcc/starfish/bin/hlo2penguin. Treat every address as version-pinned.

Abstract

A control dependency in XLA HLO is an ordering edge with no data flow: B must be scheduled after A even though B never reads A's value. In HloInstruction these edges live in a side table (control_predecessors() / control_successors()), not in the operand list. That is fatal for the Neuron backend, because the lowering path HLO → MHLO/StableHLO → Penguin Python IR → BIR crosses MLIR, and MLIR tracks only SSA value operands. Export to MHLO, an upstream canonicalization, and re-import would silently drop every side-table control edge — and with it, every ordering guarantee the frontend established.

PreserveControlDeps (pass #61, registered name preserve-control-deps, Run @ 0x1f5c110) solves this by reifying each control edge into a real SSA operand. For a successor S it builds a side-effecting AwsNeuronControlDep custom-call that takes the control-predecessors as leading operands and S's original operand(0) as the trailing operand, then rewires S.operand(0) to consume that call. The edge is now part of the data-dependence graph, so it survives MHLO import, MLIR canonicalization, fusion, and scheduling — exactly the passes that would otherwise lose it. The custom-call's side-effect bit ([cc+0x2B8]=1) keeps DCE and the scheduler from deleting or freely reordering it.

This page traces the full seam. Section 2 is the on-wire op and its operand convention. Section 3 is the three-phase Run algorithm on the HLO side. Section 4 is the MHLO/StableHLO handoff: NeuronControlDepTupleSimplifier (flatten-control-dep-tuple-operands, runOnOperation @ 0x20f6a40) un-bundles the mhlo.tuple the importer wraps the predecessor operands in, and MhloToPythonPrinter::printControlDeps (@ 0x20b8480) finally emits a Penguin .add_dep_edge( call. The recurring trap — answered up front in §2.3 — is that the HLO pass emits flat variadic operands; the "tuple" in the simplifier's name is an MHLO-import artifact, not something the HLO pass produces.

For reimplementation, the contract is:

• The wire op: AwsNeuronControlDep custom-call, empty opaque, api_version=1, side-effect flagged, result shape = forwarded operand shape.
• The operand convention: (p0, …, p_{k-1}, D) — control-predecessors first, the data value last; S.operand(0) rewired to the call.
• The filtering and module-stamp rules: drop kParameter predecessors; stamp module frontend attribute has_control_deps="1" iff ≥1 edge was reified.
• The MLIR read-back: recognize the call by call_target_name == "AwsNeuronControlDep", flatten any mhlo.tuple / mhlo.get_tuple_element wrapper, rebuild the call, and emit a Penguin dependency edge.

The Wire Op and Operand Convention

The on-wire op: AwsNeuronControlDep

The reified edge is an XLA custom-call whose target name is the literal AwsNeuronControlDep (hlo-opt rodata @ 0x26e82c, length 19; hlo2penguin @ 0x272822). The name is loaded into HloInstruction::CreateCustomCall with an explicit length of 0x13=19 — disasm 0x1f5c5ec: mov r8d, 13h immediately before the call site at 0x1f5c602. (CONFIRMED — disasm-anchored.)

The instruction is transparent: it returns the same shape it forwards (the data value's shape), so substituting it for S.operand(0) is type-preserving. The side-effect bit at HloCustomCallInstruction+0x2B8 is the mechanism that makes the reification stick — see §5.

GOTCHA — the side-effect flag is the whole point, and it is set inline. There is no named setter symbol; the body Cast<HloCustomCallInstruction>s the freshly-built call (@ 0x1f5c685) and writes the boolean directly: 0x1f5c68a: mov byte ptr [rax+2B8h], 1. A reimplementation that builds an un-flagged AwsNeuronControlDep call will see HloDCE delete it (no users care about an effect-free, single-use forwarder) and lose every control edge it was meant to pin. Offset +0x2B8 is CONFIRMED; the field name custom_call_has_side_effect_ is inferred (MED) from the single inline boolean write.

Operand convention

For a successor S whose filtered control-predecessors are P = {p0, …, p_{k-1}} (§3.1), the pass builds a flat variadic operand list with the predecessors first and the data value last:

AwsNeuronControlDep( p0, p1, …, p_{k-1}, S.operand(0) )
                     └────── k preds ──────┘  └─ data, index k (LAST) ─┘

Disasm evidence (all CONFIRMED):

• The k-predecessor pointers are memcpy'd into the operand SmallVector when the predecessor count ebx != 0; the count dword is written from ebx. The >6 path widens the inline storage via SmallVectorBase::grow_pod.
• The data value is appended last: mutable_operand(0) is read (@ 0x1f5c569) and pushed at index [size], giving total size k+1 (lea r12d,[rax+1] @ 0x1f5c5b7). So the data operand sits at index k.
• CreateCustomCall(shape = op0.shape, operands = span{…, count = k+1}, "AwsNeuronControlDep", "", 1) @ 0x1f5c602.
• The call is added to the same computation with an empty name, then S.ReplaceOperandWith(0, cc) @ 0x1f5c6a0 rewires the successor. The TF_CHECK_OK guard carries the literal successor->ReplaceOperandWith(0, controlDepCall) (str @ 0x393bd0, referenced @ 0x1f5c6af).

Net structural transform per successor group:

BEFORE:  S.operand(0) = D ;  control edge p_i → S         (side table, invisible to MLIR)
AFTER:   cc = AwsNeuronControlDep[side_effect=1](p0, …, p_{k-1}, D)
         S.operand(0) = cc                                 (edge now an SSA operand chain)

The "tuple" is an importer artifact, not an HLO construct

The HLO pass emits a flat Span<HloInstruction* const>; there is no CreateTuple call anywhere in Run. (CONFIRMED — the only constructor on the path is CreateCustomCall.) The word tuple in the downstream pass name flatten-control-dep-tuple-operands refers to the mhlo.tuple that the HLO→MHLO importer wraps the variadic predecessor bundle into (§4). The simplifier flattens that, not anything this pass built.

QUIRK — a reader who sees flatten-control-dep-tuple-operands and assumes PreserveControlDeps emits a kTuple operand will look for a CreateTuple that does not exist and mis-model the HLO operand list. The HLO list is flat; the tuple appears only after MHLO import, and only sometimes (the importer's tuple-wrapping is not statically provable to fire for every predecessor count). The simplifier handles both the tuple-wrapped and the already-flat shapes defensively.

Preservation Algorithm — PreserveControlDeps::Run

Run (@ 0x1f5c110, 7290 bytes, 282 basic blocks, 36 callees) is a stateless HloModulePass: the factory new(8)s an object holding only a vtable pointer (off_410110), with no tunable flags. It operates over all computations of the module — not just the entry computation. The body is three phases; each is anchored by a verbatim NeuronLogger DEBUG string tagged to the source file hilo/hlo_passes/PreserveControlDeps.cc (@ 0x2b8088).

NOTE — Run was extracted without Hex-Rays decompilation; the phases below are reconstructed from disassembly plus the verbatim rodata strings and callee sidecars. Operand order and the side-effect flag are CONFIRMED (disasm-anchored). The exact semantic of the phase-1 skip flag is MED, as noted in place.

Phase 1 — build the predecessor map

// loop @0x1f5c168..0x1f5c33f
function BuildPredMap(module):                          // [module@rdx] computations at +0x40..+0x48
    predMap = MapVector<HloInstruction*, SmallVector<HloInstruction*,6>>()   // operator[] @0x1f5b6d0
    for comp in module.computations():                  // stride 0x10
        for inst in comp.instructions():                // node->inst at +8
            if inst.flags[0x15] & 0x8:  continue        // test byte ptr[rdx+15h],8 @0x1f5c1df — skip (MED)
            preds = inst.rare[0x30].control_predecessors()   // kEmptyRare sentinel if null
            filtered = []
            for p in preds:
                if p.opcode_byte[0x14] == 0x4C:  continue    // cmp byte[r14+14h],4Ch @0x1f5c284 — kParameter
                filtered.push_back(p)
            predMap[inst] = filtered                    // insertion-ordered → deterministic worklist

• The opcode lives at HloInstruction+0x14 (low byte); 0x4C is kParameter. The compare cmp byte ptr [r14+14h], 4Ch ; 'L' @ 0x1f5c284 is CONFIRMED. Parameter predecessors are never reified — a control edge from a graph input is meaningless, since inputs are available before any compute begins (§5).
• [inst+0x15] & 8 is a fast-path skip (test byte ptr [rdx+15h], 8 @ 0x1f5c1df). It gates whether an instruction is considered at all (HIGH); the precise meaning is MED — most consistent with "carries no rare control metadata" or a dead/removed marker.
• The container is llvm::MapVector (operator[] @ 0x1f5b6d0, 1136 bytes), which is insertion-ordered. That makes the phase-2 worklist deterministic in computation/instruction order — important for reproducible builds.

Phase 2 — reify each map entry into a custom-call

// worklist loop @0x1f5c4e9..0x1f5c6c0
function ReifyEdges(predMap, comp):
    DEBUG("Worklist-size = " << predMap.size() << "]")          // str @0x27e813
    for (S, preds) in predMap:                                  // var_2A8 base, var_2A0 count, stride 0x48
        DEBUG("  Handling Pair:\n[\n  Successor [" S.ToString() ...   // strs @0x23bfde / @0x253051
              "  Predecessors: [" preds... "]")                       // str @0x25305d
        operands = []
        if preds.size() != 0:                                   // loc_1F5D7B0
            operands = memcpy(preds)                            // k entries; grow_pod if >6
        operands.push_back(S.mutable_operand(0))               // data value LAST (index k)
        cc = HloInstruction::CreateCustomCall(                 // @0x1f5c602
                 shape    = S.operand(0).shape(),
                 operands = operands,                          // {p0..p_{k-1}, D}
                 target   = "AwsNeuronControlDep",             // r8d=13h @0x1f5c5ec
                 opaque   = "",
                 api_ver  = 1)                                 // push 1 @0x1f5c5e6
        cc = comp.AddInstruction(cc, /*name*/"")
        Cast<HloCustomCallInstruction>(cc)                     // @0x1f5c685
            ->custom_call_has_side_effect_ = 1                 // mov byte[rax+2B8h],1 @0x1f5c68a
        TF_CHECK_OK( S.ReplaceOperandWith(0, cc) )             // @0x1f5c6a0; check str @0x393bd0
        DEBUG("ControlDep custom-call [" cc "]")               // str @0x256f7c

The DEBUG block is heavily duplicated in the disassembly — one getInstance / shouldLogToFile / shouldLogToConsole NeuronLogger triplet per << token, all gated, all source-tagged hilo/hlo_passes/PreserveControlDeps.cc with setSourceLine(0x31)=49 at the post-rewire log site. The severity is DEBUG (setCurLogLevel(1)). None of this fires in a non-verbose build, but the strings are the firmest anchors for the phase boundaries.

Phase 3 — stamp the module has_control_deps attribute

// @0x1f5daac..0x1f5dce0
function StampModule(module, anyReified):
    if not anyReified:  return false                       // fast exit @0x1f5dce5, result byte = 0
    FrontendAttributes fa                                  // ctor @0x1f5dac1
    fa.mutable_map()["has_control_deps"] = "1"             // key = 16-byte xmmword @0x406140 (movdqa @0x1f5db15); value 0x31
                                                           //   key NAME inferred (B29-1); load + value CONFIRMED
    module.frontend_attributes()[0xBA0].MergeFrom(fa)      // CopyFrom @0x1f5dcc3 / InternalSwap @0x1f5dcff
    return true                                            // changed

After all computations are processed, if any edge was reified the pass records a module-level frontend attribute whose 16-byte key is loaded from the rodata constant @ 0x406140 (movdqa xmm0, cs:xmmword_406140 @ 0x1f5db15, size 0x10); the value is the single character "1" (0x31, length 1). It is merged into HloModule frontend attributes at offset +0xBA0 via a tagged-pointer CopyFrom-vs-InternalSwap fast path.

CORRECTION (B29-1) — the key string is identified as "has_control_deps" from the source-derived backing analysis, but it is not isolable in the binary's string table: it is embedded as a 16-byte xmmword constant loaded via movdqa, so a string-table extractor never surfaces it as a standalone entry. The 16-byte movdqa load at 0x1f5db15 and its merge into module frontend attributes at +0xBA0 are CONFIRMED; the literal text "has_control_deps" is INFERRED (string length 16 matches xmmword, and it is the attribute key conventionally read by downstream fusion gates). Treat the key name as INFERRED, the mechanism as CONFIRMED. Downstream readers are MED — see §5.

The bool Run returns is changed = (≥1 cc inserted). An empty module, or one with no non-parameter control edges, takes the fast exit at 0x1f5dce5 and returns false ("no change").

HLO→MLIR Handoff — NeuronControlDepTupleSimplifier

Once HLO is exported and re-imported as MHLO (or StableHLO), the variadic AwsNeuronControlDep operand list may be wrapped: the importer can surface the multi-operand predecessor bundle as a single mhlo.tuple operand, or as chains of mhlo.get_tuple_element. NeuronControlDepTupleSimplifier (arg flatten-control-dep-tuple-operands @ 0x348e50) undoes that, restoring a flat operand list before the Penguin printer reads it.

The pass exists in two byte-twin forms differing only in mhlo:: ↔ stablehlo:: op TypeIDs:

All addresses CONFIRMED against hlo2penguin function_addresses.json (mangled symbols _ZN4hilo31NeuronControlDepTupleSimplifier… and the StableHLO-prefixed twin).

runOnOperation — collect the control-dep calls

function NeuronControlDepTupleSimplifier::runOnOperation():     // @0x20f6a40
    func = hilo::getMainFunction(module)                        // @0x21c1e70
    hits = SmallVector<Operation*>()
    func.walk<mhlo::CustomCallOp>(op):                          // forward walk; lambda @0x20f59c0
        if hilo::isControlDep(op):  hits.push_back(op)          // @0x21c0370
    replaceTuples(hits)                                         // @0x20f5d40

hilo::isControlDep(Operation*) (@ 0x21c0370, 399 bytes) is the shared control-dep recognizer: it returns true iff the op's call_target_name attribute equals "AwsNeuronControlDep" (str @ 0x425360). It is the single source of truth for "is this a control edge" — called by both tuple-simplifier walk lambdas and by the Penguin printers (printControlDeps, printControlDepCustomCall, and the operand/source printers). A reimplementation must use the exact target-name match; nothing else distinguishes the op. (CONFIRMED.)

replaceTuples — the actual flattening

function NeuronControlDepTupleSimplifier::replaceTuples(hits):  // @0x20f5d40
    for op in hits:                                             // each AwsNeuronControlDep mhlo::CustomCallOp
        flat = []
        for operand in op.operands():
            def = operand.getDefiningOp()                       // @0x20f5e71 / 0x20f6317 / 0x20f635d
            if isa<mhlo::TupleOp>(def):                         // TypeID cmp @0x20f5e96 / 0x20f639d
                for elem in def.operands():  flat.push_back(elem)   // inline the tuple's predecessors
            elif isa<mhlo::GetTupleElementOp>(def):             // TypeID cmp @0x20f633a
                flat.push_back(unwrap GTE source)               // OpResultImpl::getNextResultAtOffset @0x20f623a
            else:
                flat.push_back(operand)                         // already flat
        newOp = mhlo::CustomCallOp::build(                      // @0x20f60bf → builder @0x8fa9c60
                    builder, state,
                    TypeRange(op.result_types),
                    ValueRange(flat),                           // FLATTENED operands
                    { "call_target_name" = "AwsNeuronControlDep" })   // NamedAttr; key str @0x25a1fb
        op.replaceAllUsesWith(newOp)
        op.erase()                                              // Operation::erase @0x20f6179

The rebuild uses mhlo::CustomCallOp::build(OpBuilder&, OperationState&, TypeRange, ValueRange, ArrayRef<NamedAttribute>) (@ 0x8fa9c60), with Builder::getNamedAttr (@ 0x9ae5590), getStringAttr (@ 0x9ae6060), and RegisteredOperationName::lookup resolving the mhlo.custom_call TypeID (op name str @ 0x232721, "Building op " diag @ 0x286e24`). The net MHLO transform:

mhlo.custom_call{AwsNeuronControlDep}( mhlo.tuple(p0,…,p_{k-1}), D )
  →  mhlo.custom_call{AwsNeuronControlDep}( p0, …, p_{k-1}, D )

restoring exactly the flat (preds…, data) shape the HLO pass authored.

Terminal consumer — the Penguin emitter

The flattened AwsNeuronControlDep calls are not lowered to a Penguin op. Instead MhloToPythonPrinter::printControlDeps (@ 0x20b8480; StableHLO twin @ 0x2153360) recognizes them via isControlDep and prints Python dependency edges into neuronxcc.starfish.penguin — emission string .add_dep_edge( (str @ 0x262878), constructing neuronxcc.starfish.penguin.ir.Dependency edges. The round-trip in full:

HLO side-table edge  p_i → S
   │  PreserveControlDeps (#61)
   ▼
AwsNeuronControlDep[side_effect=1]( p0,…,p_{k-1}, D )     (flat SSA operands)
   │  HLO → MHLO import  (predecessors may be tuple-wrapped)
   ▼
mhlo.custom_call{AwsNeuronControlDep}( mhlo.tuple(p…), D )
   │  NeuronControlDepTupleSimplifier  (flatten-control-dep-tuple-operands)
   ▼
mhlo.custom_call{AwsNeuronControlDep}( p0,…,p_{k-1}, D )
   │  MhloToPythonPrinter::printControlDeps
   ▼
penguin.ir.Dependency : successor.add_dep_edge(predecessor)
   │  Penguin scheduling
   ▼
BIR ordering constraint

(Penguin emission: HIGH — anchored to the .add_dep_edge( string and the printControlDeps symbol; the precise Python argument order is not individually traced.)

Scheduling and Ordering Interaction

Why reification preserves order, mechanism by mechanism:

• Side-effect bit is the anti-DCE pin. [cc+0x2B8]=1 (§2.1) marks the call side-effecting, so HloDCE will not delete it (even though it is a single-use, effect-free-looking forwarder) and layout/schedule passes treat it as a node whose operands must stay live and whose position must be honored.
• Edges become data edges. Because the predecessors are real operands of a node feeding S.operand(0), the former control edges now live inside the data-dependence graph. Any topological scheduler — HLO, MLIR, or Penguin — honors them automatically; no separate control-edge tracking is needed downstream. This is the whole reason for the reification: it converts an MLIR-invisible side table into an MLIR-native operand chain.
• Parameter edges are pruned. A control edge p → S where p is a graph input is dropped in phase 1 (§3.1); inputs are available before any compute, so no ordering pin is needed and reifying one would create a spurious dependency on a parameter.
• Module attribute is a cheap signal. has_control_deps="1" (§3.3) lets later stages detect "this module reified control edges" without re-scanning instructions. The blob @ 0x406140 is rodata-adjacent to a "Don't fuse instr…" string, suggesting a fusion gate reads it, but the read site is not individually traced (MED).
• Placement is deliberate. #61 sits immediately before NeuronHloInstCombine (#62) and the fusion cluster, just after #60 UpcastAllToFP. Pinning the edges into operands before any peephole/fusion pass runs is what prevents those passes from dropping or reordering the edges. (Registry order from the --passes table — see Pass Registry.)

Reconstructed Structures and Signatures

// ── HLO side (hlo-opt) — stateless HloModulePass, no flags ─────────────────
class xla::hilo::PreserveControlDeps : public xla::HloModulePass {   // vtable 0x410100 (slots at +0x10)
  absl::string_view name() const { return {"preserve-control-deps", 21}; }   // 0x1f5b0d0
  StatusOr<bool> Run(HloModule*,
                     const absl::flat_hash_set<std::string_view>& exec_threads);   // 0x1f5c110
};
// factory: RegisterPreserveControlDeps()::lambda → operator new(8){ vptr = off_410110 }   // 0x1e70910

// HloInstruction field offsets used (CONFIRMED unless noted):
//   +0x14  uint16 opcode (low byte);  0x4C = kParameter
//   +0x15  flags byte (bit 0x8 = phase-1 skip; MED semantic)
//   +0x30  rare-metadata ptr → control_predecessors() vector  (kEmptyRare sentinel if null)
// HloCustomCallInstruction:
//   +0x2B8 byte custom_call_has_side_effect_  (set 1 on the reified op)   // offset CONFIRMED, name MED
// HloModule:
//   +0xBA0 FrontendAttributes  (recipient of has_control_deps="1")

xla::HloInstruction::CreateCustomCall(const Shape&, absl::Span<HloInstruction* const>,
        std::string_view target, std::string opaque, CustomCallApiVersion);   // called @0x1f5c602

// ── MLIR side (hlo2penguin) ────────────────────────────────────────────────
class hilo::NeuronControlDepTupleSimplifier
    : mlir::PassWrapper<…, mlir::OperationPass<mlir::ModuleOp>> {
  StringRef getArgument() { return "flatten-control-dep-tuple-operands"; }   // 0x20f5620
  void runOnOperation();                                                     // 0x20f6a40
  void replaceTuples(llvm::SmallVectorImpl<mlir::Operation*>&);              // 0x20f5d40
};
bool hilo::isControlDep(mlir::Operation*);   // call_target_name == "AwsNeuronControlDep"   // 0x21c0370
mhlo::CustomCallOp::build(OpBuilder&, OperationState&, TypeRange, ValueRange,
        ArrayRef<NamedAttribute>{ "call_target_name" = "AwsNeuronControlDep" });   // 0x8fa9c60

Function Map

Related Components

Cross-References

• Control-Dep Tuple-Flatten (MLIR) — 4.38; the MHLO/StableHLO flatten-control-dep-tuple-operands simplifier in depth
• Pass Registry (the --passes Table) — where preserve-control-deps sits at slot #61
• HLO/mhlo/stablehlo Ingestion Boundary — the xla::hilo:: (HLO) vs lowercase hilo:: (MLIR) split that puts these passes in different binaries
• Collectives → Custom-Call Forward Conversion — the sibling AwsNeuron* custom-call family this op belongs to
• Penguin Dependency Model — DependencyEdge & EdgeKind — Part 5; the DependencyEdge/EdgeKind model behind neuronxcc.starfish.penguin.ir.Dependency and the .add_dep_edge( calls the printer emits on the Penguin side
• BIR EdgeKind — Part 7; the BIR-level ordering edge the reified control dependency ultimately becomes