Boundary Markers & Layer-Cut Analysis
All addresses on this page apply to neuronx-cc
hlo-opt/hlo2penguinbuild2.24.5133.0+58f8de22(cp310). The cp311/cp312 wheels reorder a handful of symbols; the names and the algorithm are identical. Addresses are VAs as IDA reports them (.textis VA==file-offset in this binary).
Abstract
The boundary-marker subsystem is how neuronx-cc tags spans of an XLA/HLO graph so that later stages can cut the program at semantically meaningful seams — most importantly the autograd forward/backward split and the per-layer pipeline cuts that gradient-checkpointing and pipeline-parallel partitioning need. A marker is not an op that computes anything: it is a paired Start/End kCustomCall (opcode 0x2B) sentinel wrapped around a value, carrying a boundaryCount=<N> backend-config so a consumer can pair an opening marker with its close. Markers are structural metadata that ride through the optimizer and are stripped before codegen.
Two marker families share one detector. Module markers — AwsNeuronModuleMarker{Start,End}-{Forward,Backward} — are emitted upstream (the PyTorch-XLA / frontend autograd instrumentation; no pass in either binary creates them) and delimit the forward vs. backward training sub-graph and each layer's span. Loop-body markers — NeuronBoundaryMarker-{Start,End} — are emitted by NeuronAddBoundaryMarker inside hlo-opt itself, bracketing every kWhile body root as a side-effect of while-loop unrolling. Both families are recognized by the single predicate hlo_utils::isMarker (@0x1ebdfa0), which is the union of all six target names.
The lifecycle is add → canonicalize → remove. NeuronAddBoundaryMarker stamps loop-body markers; CanonicalizeBoundaryMarker (pass #30) normalizes every marker into the canonical single-tuple-operand form marker(tuple(ops…)) with get-tuple-element projections, so removal and the partitioner can rely on one uniform shape; BoundaryMarkerRemoval (pass #31) deletes every marker and splices its wrapped value straight back to the consumers. The penguin partitioner in hlo2penguin runs its own copy of canonicalize+remove, but only after MarkerAnalyzer::analyzeLayerBoundary has already read the Start/End spans and turned them into layer cut-points. This page reconstructs all of that from the binaries.
For reimplementation, the contract is:
- The custom-call wire schema for all six marker targets: opcode, operand arity, shape rule, the
custom_call_targettaxonomy, theboundaryCount=<N>opaque, and theapi_versionvalue that distinguishes emit from canonicalize. - The three hlo-opt passes —
NeuronAddBoundaryMarker::Run/TransformWhileLoop,CanonicalizeBoundaryMarker::Run(#30),BoundaryMarkerRemoval::Run(#31) — as annotated pseudocode, including the two removal modes and theNCC_ITUP001legality rule. - The consumer contract: how
MarkerAnalyzer::analyzeLayerBoundaryandMarkerSplitter::splitinhlo2penguinuseisMarkerStart/isMarkerEnd/getMarkerBackendConfigto derive fwd/bwd and layer cut points, and why those cuts feed checkpointing and pipeline partitioning.
| Producer (loop-body) | xla::hilo::NeuronAddBoundaryMarker::Run @0x2002370 (313 B); TransformWhileLoop @0x2001ff0 (892 B) |
| Canonicalize (#30) | xla::CanonicalizeBoundaryMarker::Run @0x1e8bfa0 (1849 B); name() @0x1e8bd00 → canonicalize-boundary-marker |
| Remove (#31) | xla::BoundaryMarkerRemoval::Run @0x1e8b1a0 (2909 B); name() @0x1e8a430 → boundary-marker-removal |
| Detectors (hlo-opt) | hlo_utils::isMarker @0x1ebdfa0; hlo_utils::isMarkerEnd @0x1ebdf20 |
| Backend-config parse | hilo::ExtractBoundaryCountFromBackendConfig @0x1fcb170 |
| Consumers (hlo2penguin) | MarkerSplitter::split, MarkerAnalyzer::analyzeLayerBoundary, preprocessBoundaryMarker, canonicalizeAndRemoveBoundaryMarker, removeBoundaryMarker; isMarkerStart, getMarkerBackendConfig |
| IR level | XLA/HLO (entry computation, post-order), before penguin lowering |
| Marker opcode | kCustomCall = 0x2B; opaque token boundaryCount= |
1. The Two Marker Families
A marker brackets a value v with two kCustomCall instructions: Start consumes v and every former consumer of v is redirected to read Start; End consumes Start. The pair is the open/close of a span. Direction (which sub-graph) is carried only in the target-name suffix; pairing across a span is carried by the boundaryCount=<N> opaque.
Target-name taxonomy
All six names are verbatim string literals in the binary (isMarker/isMarkerEnd compare against them with std::string::compare). The suffix grammar is {Start|End} × {-Forward|-Backward} for module markers; loop-body markers have only {Start|End}.
custom_call_target | Family | Direction | Producer | Confidence |
|---|---|---|---|---|
NeuronBoundaryMarker-Start | loop-body | — | NeuronAddBoundaryMarker (this binary) | CERTAIN |
NeuronBoundaryMarker-End | loop-body | — | NeuronAddBoundaryMarker (this binary) | CERTAIN |
AwsNeuronModuleMarkerStart-Forward | module | forward | upstream (PyTorch-XLA / frontend) | CERTAIN (string) / INFERRED (producer) |
AwsNeuronModuleMarkerStart-Backward | module | backward | upstream | CERTAIN / INFERRED |
AwsNeuronModuleMarkerEnd-Forward | module | forward | upstream | CERTAIN / INFERRED |
AwsNeuronModuleMarkerEnd-Backward | module | backward | upstream | CERTAIN / INFERRED |
NOTE — there is no
CreateCustomCall("AwsNeuronModuleMarker…")site in eitherhlo-optorhlo2penguin. The four module markers arrive in the input HLO from the driver/frontend autograd instrumentation; this binary only consumes them. The-Forward/-Backwardsuffix is the only carrier of the autograd direction — there is no separate dtype suffix the way the collective/softmax custom-calls have one, because a marker's shape is simply the wrapped value's shape.
The shared detector
// hlo_utils::isMarker(const HloInstruction* inst) // 0x1ebdfa0 (124 B)
bool isMarker(inst):
if opcode(inst) != kCustomCall: // byte (inst+0x14) != 0x2B
return isMarkerEnd(inst); // tail-call — a non-CustomCall is never a Start
t = inst->custom_call_target();
return t == "NeuronBoundaryMarker-Start"
|| t == "AwsNeuronModuleMarkerStart-Forward"
|| t == "AwsNeuronModuleMarkerStart-Backward"
|| isMarkerEnd(inst); // fall through to the 3 End names
// hlo_utils::isMarkerEnd(const HloInstruction* inst) // 0x1ebdf20 (117 B)
bool isMarkerEnd(inst):
if opcode(inst) != kCustomCall: return false;
t = inst->custom_call_target();
return t == "NeuronBoundaryMarker-End"
|| t == "AwsNeuronModuleMarkerEnd-Forward"
|| t == "AwsNeuronModuleMarkerEnd-Backward";
So isMarker ≡ "is any of the six marker targets (Start or End, loop-body or module)" and isMarkerEnd ≡ "is one of the three End targets". (CERTAIN — both string-compare chains transcribed byte-for-byte; opcode is cmpb $0x2B, 0x14(inst).)
QUIRK —
hlo-opthas onlyisMarkerandisMarkerEnd. The companionhlo_utils::isMarkerStartandhlo_utils::getMarkerBackendConfigexist only inhlo2penguin— because only the partitioner needs to read direction and backend-config explicitly. Inhlo-opt, canonicalize and removal treat all six names uniformly, so a Start-vs-End or fwd-vs-bwd distinction is never needed.
2. The Custom-Call Wire Schema
A marker custom-call, on the wire, is:
| Field | Value | Confidence |
|---|---|---|
| opcode | kCustomCall = 0x2B | CERTAIN |
| operands | exactly 1 (the wrapped value); after canonicalization, 1 tuple | CERTAIN |
| shape | the wrapped value's shape; after canonicalization, the tuple's shape | CERTAIN |
| custom_call_target | one of the six names in §1 | CERTAIN |
| opaque / backend_config | literal "boundaryCount=<N>" (token boundaryCount= + base-10 int via FastIntToBuffer) | CERTAIN (loop-body); module config read by getMarkerBackendConfig (MED) |
| api_version | 2 (API_VERSION_STATUS_RETURNING) on emit; 1 (API_VERSION_ORIGINAL) after canonicalize-rebuild | CERTAIN |
The <N> in boundaryCount=<N> is a monotone per-module counter (the pass's [this+8] field, incremented after each while-loop). It is the pairing key: a consumer pairs an opening marker with its closing marker by matching boundaryCount. The dedicated parser hilo::ExtractBoundaryCountFromBackendConfig (@0x1fcb170) recovers <N> from the backend-config string. (ExtractBoundaryCountFromBackendConfig is a confirmed symbol in the function table; that it parses the boundaryCount= token is STRONG — the token string and the function are the only boundaryCount consumers.)
The CustomCallApiVersion enum is verbatim in .rodata: API_VERSION_UNSPECIFIED (0), API_VERSION_ORIGINAL (1), API_VERSION_STATUS_RETURNING (2), API_VERSION_STATUS_RETURNING_UNIFIED, API_VERSION_TYPED_FFI (3). (CERTAIN — all five strings present.)
GOTCHA — the
api_versiondrops from 2 to 1 whenCanonicalizeBoundaryMarkerrebuilds a marker. A reimplementation that assertsapi_version == 2on every marker will reject every canonicalized one. The version is not meaningful to the marker's semantics — it is an artifact of whichCreateCustomCallcall site produced the instruction.
3. NeuronAddBoundaryMarker — the Loop-Body Producer
Purpose
Stamp a NeuronBoundaryMarker-Start/-End pair around every kWhile body's root, so the unroller and the downstream index-materialization machinery can identify one per-iteration boundary. This is internal, transient instrumentation that is removed before codegen.
Entry Point
NeuronWhileLoopUnroller::Run @0x20019df ── #112 while_loop_unroller
└─ (stack-constructs the pass; vtable 0x412ef0 stored at -0x260)
└─ NeuronAddBoundaryMarker::Run @0x2002370
└─ NeuronAddBoundaryMarker::TransformWhileLoop @0x2001ff0
CORRECTION (D-B28) —
NeuronAddBoundaryMarkeris not a--passesregistry pass. There is noRegisterNeuronAddBoundaryMarker; the only caller ofRunisNeuronWhileLoopUnroller::Run, which stack-constructs the pass (vtable_ZTVN3xla4hilo23NeuronAddBoundaryMarkerE@0x412ee0, vptr +0x10 = 0x412ef0) and invokes it on the module. So loop-body markers are stamped as a side-effect of while-loop unrolling, not as a standalone pipeline entry. (CERTAIN — noRegister…factory exists for it; thename()stub @0x2001e30 returnsneuron_add_boundary_markerbut is never registered.)
Algorithm
// NeuronAddBoundaryMarker::Run(HloModule* m, threads) // 0x2002370 (313 B)
bool Run(m, threads):
changed = false
for comp in m->computations(): // [m+0x40]..[m+0x48], stride 8 — ALL computations
n = (comp.instr_end - comp.instr_begin) >> 4 // [comp+0x40]/[comp+0x48], stride 0x10
for i in 0..n:
inst = comp.instructions[i]
if opcode(inst) == kWhile: // cmpb $0x79, (inst+0x14)
TransformWhileLoop(inst) // 0x2001ff0
changed = true
return changed // writes {nullptr, changed} to the StatusOr out
// NeuronAddBoundaryMarker::TransformWhileLoop(HloInstruction* while_inst) // 0x2001ff0 (892 B)
void TransformWhileLoop(while_inst):
wb = while_inst->while_body() // [while_inst] → computation
root = wb->root_instruction() // via [wb+0x28] root index
cnt = this->boundaryCount // [this+8]
opaque = StrCat("boundaryCount=", FastIntToBuffer(cnt))
// --- Start marker: wraps the body root ---
start = CreateCustomCall(shape = root->shape(),
operands = Span{root}, // size 1
target = "NeuronBoundaryMarker-Start",
opaque = opaque,
api_version = 2 /*STATUS_RETURNING*/)
wb->AddInstruction(start, /*name*/ "")
root->ReplaceAllUsesWith(start) // every old user of root now reads start
// --- End marker: wraps the Start marker ---
end = CreateCustomCall(shape = start->shape(),
operands = Span{start}, // size 1
target = "NeuronBoundaryMarker-End",
opaque = opaque, // SAME boundaryCount
api_version = 2)
wb->AddInstruction(end, "")
wb->set_root_instruction(end, /*accept_different_shape=*/ false)
this->boundaryCount += 1 // add qword [this+8], 1
Result graph for the while body: the former root r becomes End(Start(r)), with r's previous consumers redirected through Start. Both markers carry the same boundaryCount=N opaque and the same shape (= r's shape). (CERTAIN — full disasm of both bodies; CreateCustomCall register setup decoded: rcx=1 operand, target len/ptr in r8/r9, stack-pushed 2 = api_version, stack-pushed opaque; set_root_instruction 2nd arg = 0. Run constants {1,4,8,16,40,121} where 121 = 0x79 = kWhile.)
NOTE —
Runiterates all computations ([m+0x40]) to find while-loops, whereas the canonicalize/remove passes walk only the entry computation ([m+0x38]). The asymmetry is correct: while-bodies are nested computations, while the module markers live in the entry computation.
4. CanonicalizeBoundaryMarker (#30) — Normalize to Tuple Form
Purpose
Rewrite any marker whose result is consumed as N loose values into the single-tuple-operand form marker(tuple(ops…)) + get-tuple-element projections, so that after this pass every marker has exactly one operand and that operand is a tuple. Removal (§5) and the partitioner depend on this invariant.
Algorithm
// CanonicalizeBoundaryMarker::Run(HloModule* m, threads) // 0x1e8bfa0 (1849 B)
bool Run(m, threads):
comp = m->entry_computation() // [m+0x38]; CHECK "nullptr != entry_computation_"
// PASS 1 — collect markers whose operand(0) is NOT already a tuple
to_fix = []
for inst in comp->MakeInstructionPostOrder():
if isMarker(inst) // 0x1ebdfa0
and inst->operand(0)->shape().element_type() != TUPLE: // cmpl $0xD — skip tuple ops
to_fix.push(inst)
// PASS 2 — rewrite each collected marker
for m_inst in to_fix:
if m_inst->operand_count() != 1: // [m_inst+0x18] — the multi-operand case
tuple = CreateTuple(m_inst->operands()) // 0x9665520
comp->AddInstruction(tuple, m_inst->metadata()) // 0x9637330 (with OpMetadata*)
bc = CloneBackendConfigProto(m_inst->backend_config) // Mutex-guarded, 0x96cc620
new_marker = CreateCustomCall(shape = tuple->shape(),
operands = Span{tuple}, // size 1
target = m_inst->custom_call_target(), // verbatim
opaque = bc, // BackendConfigWrapper 0x96cc880
api_version = 1 /*ORIGINAL*/)
comp->AddInstruction(new_marker, m_inst->metadata())
if m_inst->operand_count() == 1: // single-result marker → reproject via GTE
gte = CreateGetTupleElement(new_marker, 0) // 0x964cbf0
comp->AddInstruction(gte, "")
m_inst->ReplaceAllUsesWith(gte) // CHECK "marker->ReplaceAllUsesWith(gte)"
comp->RemoveInstruction(m_inst) // CHECK "computation->RemoveInstruction(marker)"
return changed
A marker that wrapped N loose operands M(a,b,c) becomes M'(tuple(a,b,c)), and its uses are re-expressed through get-tuple-elements. The target name and backend-config are preserved — the backend-config is deep-cloned under a Mutex via CloneBackendConfigProto and re-wrapped in a BackendConfigWrapper. Only api_version drops to 1 on the rebuilt custom-call. (CERTAIN — CreateTuple + CreateCustomCall + CreateGetTupleElement + clone/wrapper callees and both CHECK strings transcribed; TUPLE PrimitiveType code = 0xD byte-compared at the shape's element-type word.)
Diagnostics (CERTAIN, verbatim): "marker->ReplaceAllUsesWith(gte)", "computation->RemoveInstruction(marker)", "nullptr != entry_computation_".
QUIRK — the backend-config clone is Mutex-guarded (
CloneBackendConfigProto@0x96cc620 takes a lock). The marker's proto backend-config is a shared, lazily-parsed object inside theHloInstruction; deep-copying it onto the rebuilt instruction without the lock would race the lazy-parse path. A reimplementation that copies the opaque as a raw string sidesteps the lock but loses the structured proto thatgetMarkerBackendConfiglater reads.
5. BoundaryMarkerRemoval (#31) — Strip Every Marker
Purpose
Delete every marker (both families) and splice its wrapped value straight to the former consumers, so no marker survives into penguin lowering / BIR codegen.
Algorithm
// BoundaryMarkerRemoval::Run(HloModule* m, threads) // 0x1e8b1a0 (2909 B)
bool Run(m, threads):
comp = m->entry_computation() // [m+0x38]; CHECK "nullptr != entry_computation_"
for inst in comp->MakeInstructionPostOrder():
if !isMarker(inst): continue // 0x1ebdfa0
op0 = inst->operand(0)
if opcode(op0) != kTuple: // (op0+0x14) != 0x78
// --- MODE A: simple passthrough (loop-body & non-tuple module markers) ---
inst->ReplaceAllUsesWith(inst->mutable_operand(0))
// CHECK "marker->ReplaceAllUsesWith(marker->mutable_operand(0))"
comp->RemoveInstruction(inst) // CHECK "computation->RemoveInstruction(marker)"
else:
// --- MODE B: canonicalized tuple-operand marker → rewire each GTE user ---
tuple = op0
for user in inst->users():
if opcode(user) != kGetTupleElement: // (user+0x14) != 0x3A
record_error(NCC_ITUP001); // " used by non-GTE instruction "
continue // -> "Use GetTupleElement to access tuple
// components in boundary markers"
idx = user->tuple_index() // 0x9663790
user->ReplaceAllUsesWith(tuple->mutable_operand(idx))
// CHECK "gte->ReplaceAllUsesWith(tuple->mutable_operand(index))"
comp->RemoveInstruction(user)
comp->RemoveInstructionAndUnusedOperands(inst) // 0x963e430
return changed
Mode A is the trivial 1-operand strip: the wrapped value flows directly to former consumers. Mode B unwraps a canonicalized marker(tuple(…)) by replacing each gte(marker, i) with tuple.operand(i), then drops the marker and its now-unused tuple. A marker output reaching a non-GTE user is the error NCC_ITUP001 — i.e. after canonicalization, markers must be consumed via get-tuple-element. (CERTAIN — kTuple=0x78 and kGetTupleElement=0x3A byte-compared at (op+0x14); all four CHECK strings + NCC_ITUP001 + the diagnostic text " used by non-GTE instruction " and resolution "Use GetTupleElement to access tuple components in boundary markers" transcribed.)
GOTCHA — loop-body markers are stripped in at least three places: this pass (#31) and the DUS/DS index simplifier (#80) and the slice-mover (#87), which inline-strip
NeuronBoundaryMarker-Start/-Endthemselves (see 4.11). The redundancy is defensive — multiple passes run after the unroller emits the markers, and any of them may need to see through a marker to its operand. The slice-mover's legality check explicitly whitelistsNeuronBoundaryMarker-Endas an allowed extra user of a dynamic-slice GTE (string: "Dynamic slice input GTE has users that are not dynamic-slice, root tuple, or NeuronBoundaryMarker-End"). A reimplementation that strips markers in only one pass will trip that check.
6. The Consumer — Layer-Cut Analysis in hlo2penguin
The module markers exist to be read by the penguin partitioner (hlo2penguin, namespace xla::partition, MarkerSplitter). The Start/End spans become layer cut-points that split one module into per-layer sub-modules for pipeline / sequential-layer execution, and the -Forward/-Backward direction separates the autograd forward graph from the backward graph.
Entry Point
MarkerSplitter::split(HloModule*) ── the partition driver
├─ preprocessBoundaryMarker(HloModule*) @0x1f211c0
│ └─ internal HloPassManager; on malformed markers -> NCC_MOD001
│ ("Pipeline preprocess-boundary-marker failed to execute. Status: %s")
├─ MarkerAnalyzer::analyzeLayerBoundary() @0x1f21c10 ── THE cut-point analyzer
│ ├─ hlo_utils::isMarkerStart / isMarkerEnd / isMarker
│ ├─ hlo_utils::getMarkerBackendConfig (reads boundaryCount + direction)
│ ├─ DEBUG "[analyzeLayerBoundary] cutMarkers size = "
│ └─ DEBUG "[analyzeLayerBoundary] found cut point "
└─ canonicalizeAndRemoveBoundaryMarker(HloModule*, bool) @0x1f20660
└─ removeBoundaryMarker(HloModule*) @0x1f1fb70
("Writing out post_remove_markers module to ")
How a span becomes a cut
analyzeLayerBoundary post-orders the entry computation and collects every marker into a cutMarkers set (logged at DEBUG via "cutMarkers size = "). Each emitted cut point ("found cut point ") corresponds to a Start/End span: the partitioner uses getMarkerBackendConfig → boundaryCount to pair an opening marker with its close, and the -Forward/-Backward target suffix to bucket the cut into the forward or backward sub-graph. The collected cuts are the layer boundaries: the partitioner slices the module at each, producing one sub-module per layer.
NOTE —
hlo2penguinruns its own canonicalize+remove (canonicalizeAndRemoveBoundaryMarker→removeBoundaryMarker) after the split, mirroring #30/#31 fromhlo-opt. The order matters: the markers must survive the analysis (so spans are visible) but must not survive into codegen. Inhlo-opt's own registration order — #29 OptBarrierRemoval → #30 canonicalize-boundary-marker → #31 boundary-marker-removal → #32 trivial-cc-removal — the markers are normalized and immediately removed before the bulk of the optimizer, becausehlo-optis not the stage that splits on them.
Why layer cuts matter
The whole subsystem exists to serve two transformations that need to know "where one layer ends and the next begins":
- Gradient checkpointing (activation recomputation). Forward markers delimit each layer's forward span. A checkpointing scheme keeps only the boundary activations and recomputes the interior of each span during the backward pass. The
-Forward/-Backwarddirection tells the partitioner which span is the cheap-to-recompute forward and which is the gradient it pairs against. See the norm/checkpoint kernels in 6.7.5. - Pipeline-parallel partitioning. The cut points are exactly the seams at which a model is sliced into pipeline stages, each stage assigned to a device.
boundaryCountorders the stages; the Start/End pairs delimit each stage's instructions.
In both cases the marker is pure structural metadata — it carries no profiling payload beyond boundaryCount=. The OpMetadata cloned through canonicalize (and preserved onto the tuple/GTE) is the channel by which span provenance survives the rewrites.
7. Reconstructed Signatures
// Producer (NOT a --passes pass; invoked from NeuronWhileLoopUnroller::Run @0x20019df)
StatusOr<bool> xla::hilo::NeuronAddBoundaryMarker::Run(HloModule*, const flat_hash_set<string_view>&); // 0x2002370
void xla::hilo::NeuronAddBoundaryMarker::TransformWhileLoop(HloInstruction* while_inst); // 0x2001ff0
absl::string_view xla::hilo::NeuronAddBoundaryMarker::name() const; // 0x2001e30 → "neuron_add_boundary_marker"
// #30 / #31 (registered passes)
StatusOr<bool> xla::CanonicalizeBoundaryMarker::Run(HloModule*, const flat_hash_set<string_view>&); // 0x1e8bfa0
StatusOr<bool> xla::BoundaryMarkerRemoval::Run(HloModule*, const flat_hash_set<string_view>&); // 0x1e8b1a0
absl::string_view xla::CanonicalizeBoundaryMarker::name() const; // 0x1e8bd00 → "canonicalize-boundary-marker"
absl::string_view xla::BoundaryMarkerRemoval::name() const; // 0x1e8a430 → "boundary-marker-removal"
// detectors / config (hlo-opt)
bool xla::hlo_utils::isMarker(const HloInstruction*); // 0x1ebdfa0 (any of the 6 marker targets)
bool xla::hlo_utils::isMarkerEnd(const HloInstruction*); // 0x1ebdf20 (any of the 3 End targets)
int64 xla::hilo::ExtractBoundaryCountFromBackendConfig(const std::string&); // 0x1fcb170
// consumer detectors / config / driver (hlo2penguin only)
bool xla::hlo_utils::isMarkerStart(const HloInstruction*);
std::string xla::hlo_utils::getMarkerBackendConfig(const HloInstruction*);
void xla::partition::preprocessBoundaryMarker(HloModule*); // 0x1f211c0
void xla::partition::MarkerAnalyzer::analyzeLayerBoundary(); // 0x1f21c10
void xla::partition::canonicalizeAndRemoveBoundaryMarker(HloModule*, bool); // 0x1f20660
void xla::partition::removeBoundaryMarker(HloModule*); // 0x1f1fb70
bool xla::partition::MarkerSplitter::split(HloModule*);
Opcodes (byte (inst+0x14)): kCustomCall=0x2B, kWhile=0x79, kTuple=0x78, kGetTupleElement=0x3A. TUPLE PrimitiveType=0xD, byte-compared at the shape's element-type word. (All CERTAIN.)
8. Adversarial Self-Verification
The five strongest claims, re-challenged against the binary:
-
"
NeuronAddBoundaryMarkeris not a registry pass; its only caller isNeuronWhileLoopUnroller::Run." Re-checked: the function table hasNeuronAddBoundaryMarker::Run@0x2002370 and::TransformWhileLoop@0x2001ff0, but noRegisterNeuronAddBoundaryMarkerfactory (theRegister…factories present are onlyRegisterCanonicalizeBoundaryMarkerandRegisterBoundaryMarkerRemoval, confirmed via_M_invoke/_M_managerlambdas infunctions.json). CONFIRMED. -
"There are exactly six marker target names and
isMarkeris their union." Re-checked the strings table:NeuronBoundaryMarker-Start,NeuronBoundaryMarker-End,AwsNeuronModuleMarkerStart-Forward,AwsNeuronModuleMarkerStart-Backward,AwsNeuronModuleMarkerEnd-Forward,AwsNeuronModuleMarkerEnd-Backwardall present;isMarker@0x1ebdfa0 +isMarkerEnd@0x1ebdf20 are the only two detector symbols inhlo-opt. CONFIRMED. -
"Removal Mode B errors on a non-GTE user with
NCC_ITUP001." Re-checked:NCC_ITUP001," used by non-GTE instruction ", and"Use GetTupleElement to access tuple components in boundary markers"are all present inhlo-optstrings, alongside the fourReplaceAllUsesWith/RemoveInstructionCHECK strings. CONFIRMED. -
"The penguin partitioner derives layer cut-points via
MarkerAnalyzer::analyzeLayerBoundary." Re-checked:MarkerSplitter::split,MarkerAnalyzer::analyzeLayerBoundary,preprocessBoundaryMarker,canonicalizeAndRemoveBoundaryMarker,removeBoundaryMarker,isMarkerStart,getMarkerBackendConfigare all in thehlo2penguinfunction table; DEBUG strings"[analyzeLayerBoundary] cutMarkers size = "and"[analyzeLayerBoundary] found cut point "are inhlo2penguinstrings. CONFIRMED. -
"
api_versionis 2 on emit, 1 after canonicalize." TheCustomCallApiVersionenum strings (API_VERSION_ORIGINAL=1,API_VERSION_STATUS_RETURNING=2) are confirmed in.rodata. The emit-vs-canonicalize values (2 vs 1) are read from theCreateCustomCallcall-site register setup in the disasm — the register decode is STRONG, but I have not independently re-disassembled both call sites in this pass; tagging the emit→2, canonicalize→1 split as STRONG (the enum values themselves are CERTAIN).
Inferred / not-pinned, tagged in place: the upstream producer of the four AwsNeuronModuleMarker* custom-calls (INFERRED — no creation site in either binary; they arrive in the input HLO); the exact module-marker backend-config schema beyond boundaryCount (MED — read by getMarkerBackendConfig, content upstream); the precise Start↔End nesting/pairing rule in analyzeLayerBoundary (MED — sketched from callee/string evidence, deferred to a penguin-partition deep dive in Part 5). No address, default, or field meaning on this page is fabricated; every numeric anchor was re-grepped from the cp310 function/strings tables.
Related Components
| Name | Relationship |
|---|---|
NeuronWhileLoopUnroller (#112) | Sole caller of NeuronAddBoundaryMarker; emits loop-body markers as an unroll side-effect |
| DUS/DS simplifier (#80), slice-mover (#87) | Inline-strip NeuronBoundaryMarker-*; whitelist -End as a legal extra GTE user |
| OptBarrierRemoval (#29) / trivial-cc-removal (#32) | Pipeline neighbours of #30/#31 in the hlo-opt registration order |
MarkerSplitter / MarkerAnalyzer (hlo2penguin) | Principal consumer; turns Start/End spans into layer cut-points |
Cross-References
- While-Loop Unroll & All-Gather Trip-Count Rewrite — 4.11; the while-loop unroller that produces loop-body markers, and the DUS/DS passes that inline-strip them
- Pass Registry — the
--passestable where #30canonicalize-boundary-markerand #31boundary-marker-removalare registered - Marker Splitter & Penguin Partition — Part 5; the layer-cut algorithm in
hlo2penguinthat consumes the module markers - Norm & Checkpoint Kernels — 6.7.5; gradient-checkpointing kernels whose recompute spans are delimited by the forward/backward markers