Lowering On-Chip Collectives — lower_local_collectives / insert_ptcom_flat
All symbols and addresses on this page apply to
neuronx_cc2.24.5133.0+58f8de22 (cp310 wheel; cp310/11/12 share the C++ core with a fixed per-build address delta — re-confirm any raw address against the target wheel). The two passes live inneuronxcc/starfish/lib/libwalrus.so(build-id92b4d331a42d7e80bb839e03218d2b9b0c23c346, stripped; all names resolve from.dynsym, there is no.symtab). For.text(0x62d660+) and.rodata(0x1c72000+) the virtual address equals the file offset;.datacarries a+0x400000delta. The IR-side handles (InstGPSIMDSB2SB::isCompatible,getEffectiveBytesPerPartition, theMaxBytesPerPartitionceiling) live inlibBIR.so(build-ida9b1ea38…). The pre-dumped.csidecars for these two passes are PLT re-export thunks (_0x5e9xxx); the real bodies disassemble from.textat0x1615xxx/0x16axxx— the two-VA-frame artifact. Treat every address as version-pinned. See Build & Version Provenance.
A backend collective — an AllReduce, a SendRecv, a ReduceScatter — arrives at the
libwalrus backend as a single opaque opcode: bir::InstCollectiveCompute, InstructionType
48, with a CollectiveKind sub-field that multiplexes all eleven collective variants
through one instruction. By the time the backend pipeline reaches order 80, two earlier
multi-core passes have already run: lnc_splitter (order 8) cloned the symbolic module into
one bir::Module per physical core (NeuronCoreId 0..lnc-1), and ExpandReplication
(order 7) made the SPMD replication explicit. lower_local_collectives is the pass that takes
that vector of per-core modules and decomposes the on-chip collectives — the ones whose data
movement stays inside the LNC group of cores on a single die (tensor-parallel / FSDP traffic
within a node) — into concrete engine ops: GPSIMD cross-core SBUF→SBUF copies, core-to-core
DMA, CoreBarrier fences. The cross-node collectives (ICI / EFA fabric) it leaves untouched.
The headline is a single unsigned compare. lower_local_collectives dispatches on
CollectiveKind: kinds 0/1/2 (SendRecv / SendRecvCCE / AllReduce) get a real local
lowering; kind ≥ 3 (the ReduceScatter / AllGather / AllToAll / Permute family) is
fence-only — the op is left intact for a downstream / runtime path. Inside the kind-0
SendRecv lowering sits a second, finer choice: each swap is carried either by the GPSIMD
engine's SB2SB copy or by DMA, selected by a byte-budget gate. Its sibling pass
insert_ptcom_flat then drops point-to-point completion tokens (InstTensorCompletion,
"PTCOM") so a consumer core can wait on a single producer write rather than on a whole barrier.
See also GPSIMD Engine — the Pool-Alias Cross-Core SB2SB Mover
for the InstGPSIMDSB2SB (IT 33) mechanism and the libBIR isCompatible predicate this pass
gates on, and Collective / GPSIMD / CustomOp Encoding
for the bir::CollectiveKind ordinals and the IT-48 wire encoding. The cross-node collective
lowering (the kind-≥3 remote path) is documented in the planned Part 13 distribution chapter.
The symbol roster
Both passes follow the dual-overload shape shared by every multi-core backend pass
(LncSplitter, LncVerifier): a per-Module& overload that is a no-op / abort shim, and a
vector<unique_ptr<Module>>& overload that is the real worker. The BackendPass dispatcher
drives the vector overload because the work is inherently cross-module / cross-core — the
per-Module& shim is never run for these passes. The same shim/worker split appears on
LncSplitter and LncVerifier.
LowerLocalCollectives (factory register_generator_lower_local_collectives__ @0x162a6a0):
| Symbol | Address | Role |
|---|---|---|
run(bir::Module&) | 0x1615680 | per-module shim (stores 0x2, no-op) |
run(vector<unique_ptr<Module>>&) | 0x1624850 | real worker (~0x5e50 bytes) |
getRemoteCores(unsigned long self) | 0x1617940 | {0..lnc-1}\{self} |
getScatterDimension(PhysicalAccessPattern const&) | 0x1615d40 | tail-jmp alias → getFirstSplittableNonPartitionDimension 0x1615c70 |
extractAndGroupCCOps(...) | 0x1618440 | clusters IT-48 ops into SPMD groups |
lowerSendRecv(...) | 0x161cc70 | CollectiveKind 0 |
lowerSendRecvCCE(...) | 0x161f130 | CollectiveKind 1 |
lowerAllReduce(...) | 0x1621a80 | CollectiveKind 2 |
createRemoteAP(...) | 0x161cad0 | clone AP onto a peer core |
createPartialAP(...) | 0x1620fb0 | per-core sub-tensor AP (the scatter slice) |
getMemoryLocation(...) | 0x161ab30 | remote-memloc minter (vnc_remote_addr_map) |
insertCoreBarrier(...) | 0x161a150 | emit InstCoreBarrier (IT 87) |
updateDeps(...) | 0x1619620 | dependency fix-up before deletion |
sendRecvToGPSIMDMaxBytesPerPartition | .bss 0x3e050a0 | cl::opt byte-budget value |
InsertPTCOMFlat (factory register_generator_insert_ptcom_flat__ @0x16ae030):
| Symbol | Address | Role |
|---|---|---|
run(bir::Module&) | 0x16a5710 | per-module shim (stores 0x2, no-op) |
run(vector<unique_ptr<Module>>&) | 0x16ad1c0 | real worker |
identifyPTCOMBlocks(...) | 0x16aa040 | per-memloc convergence block (postdom) |
identifyLastWrites(...) | 0x16a9280 | per-(ML,BB) last writer |
extendLastWrites(...) | 0x16abd80 | propagate last-write across cores |
insertPTCOMs(...) | 0x16acc50 | emit InstTensorCompletion (IT 94) |
findMemLocOnCore0(...) | 0x16a7540 | flat core-0-name lookup |
findCBOnCore0(...) | 0x16a8010 | match a CoreBarrier on core 0 |
meetsPTCOMRequirements(...) | 0x16a5840 | free fn — non-loop-SCC gate |
enablePTCOM | .bss 0x3df6a40 | pass-enable cl::opt |
enableRemoteSemaphoreDMA | .bss 0x3e05160 | extra-barrier gate (kind 0) |
What "local" means — the on-chip vs remote split
InstCollectiveCompute carries its CollectiveKind at Inst+0xf8. The enum is:
0 SendRecv 1 SendRecvCCE 2 AllReduce 3 ReduceScatter 4 AllGather
5 AllToAll 6 AllToAllV 7 Permute 8 PermuteReduce 9 PermuteImplicit
10 PermuteReduceImplicit
A single IT-48 opcode multiplexes all of them; the kind field is the discriminator. The worker prologue first gates on the LNC width and then dispatches on the kind.
The LNC gate. The run prologue at 0x1624876 is cmp DWORD[Module->PassOptions+0x1A4], 1; je return, where PassOptions+0x1A4 is lnc_size, the number of cores. When lnc == 1 there
is no cross-core exchange to lower, so the worker returns cleanly — the pass is a no-op on
single-core configs.
The kind dispatch (worker dispatch 0x1629337):
// LowerLocalCollectives::run(vector<Module>&) @0x1624850 (worker dispatch @0x1629337)
// 'cc0' is group[0]; kind = cc0[+0xf8] loaded into eax.
unsigned kind = *(uint32_t *)(cc0 + 0xf8);
cmp(kind, 1);
ja >= 3 -> default_remote; // 0x162999f (unsigned 'ja' past 1 => kind >= 3)
je == 1 -> lowerSendRecvCCE; // 0x1629949
cmp(kind, 2);
je == 2 -> lowerAllReduce; // 0x1629859
test_eq_zero -> lowerSendRecv; // 0x1629676 (kind == 0)
The control flow is an unsigned ja past 1 (which folds every kind ≥ 3 into one branch
because the comparison is unsigned), plus explicit == 1 / == 2 / == 0 compares. The two
groups behave very differently:
-
Kinds 0 / 1 / 2 (
SendRecv,SendRecvCCE,AllReduce) — the on-chip / local path. Each is decomposed here into SBUF↔SBUF data exchange plusCoreBarriersynchronization. These are the collectives whose traffic stays inside the LNC group of on-die cores. -
Kind ≥ 3 (
ReduceScatter/AllGather/AllToAll/AllToAllV/Permutefamily) — the remote / ICI path. The collective op is not decomposed here. The run loop still callsinsertCoreBarrierper remote core (0x16299f7) to fence it, but the IT-48 op is left intact for a different cross-node lowering downstream. In the call graph, kinds 0/1/2 are the only ones that reach alowerXXXhelper; kind ≥ 3 reaches nothing butinsertCoreBarrier. [INFERRED] That the downstream consumer is the runtime collective library over the ICI fabric is not pinned anywhere in this binary.
The decomposition target is on-die SBUF↔SBUF transport: either the GPSIMD cross-core SB2SB copy
engine (InstGPSIMDSB2SB, IT 33) or core-to-core DMA (InstDMACopy), synchronized by
InstCoreBarrier (IT 87) and finalized by the PTCOM op InstTensorCompletion (IT 94).
The worker loop
// LowerLocalCollectives::run(vector<Module>&) @0x1624850
// reconstructed from disasm 0x1624850..0x162a6a0
void run(vector<unique_ptr<Module>> &modules) {
int lnc = modules[0]->passOptions[+0x1A4]; // 0x1624876
if (lnc == 1) return; // nothing to lower
auto precMap = getPrecedingCoreBarrierOrLocalCCMap(modules); // 0x16248ab
auto groups = extractAndGroupCCOps(modules); // 0x16248c4 -> vector<vector<CC*>>
vector<uint> allCores = {0 .. lnc-1}; // built via _M_realloc_insert loop
for (auto &group : groups) { // one group = the SPMD-symmetric copies
CC *cc0 = group[0]; // of ONE collective across all cores
uint core = cc0->getModule()->getNeuronCoreId(); // 0x1624a28
vector<DenseMap<MemoryLocation*,MemoryLocation*>> denseMaps[lnc]; // per-core remote map
vector<vector<Instruction*>> newInsts[lnc]; // ops per core
unsigned long semCounter = 0;
unsigned kind = *(uint32_t *)(cc0 + 0xf8);
switch (kind) { // dispatch @0x1629337
case 0: lowerSendRecv (group, newInsts, allCores, denseMaps, &semCounter); break; // 0x16296c2
case 1: lowerSendRecvCCE(group, newInsts, allCores, denseMaps, &semCounter); break; // 0x1629995
case 2: lowerAllReduce (group, newInsts, denseMaps, &semCounter); break; // 0x16298a2
default: /* kind >= 3, remote/ICI: fence only, op left intact */ // 0x162999f
for (size_t i = 0; i < group.size(); ++i)
insertCoreBarrier(group[i], allCores, newInsts, i, name, &semCounter); // 0x16299f7
}
// kind==0 special case (0x16296e2): unless enableRemoteSemaphoreDMA (.bss 0x3e05160)
// is set, add a CoreBarrier per remote core for the SendRecv swap.
for (size_t i = 0; i < group.size(); ++i)
updateDeps(group[i], newInsts[i]); // 0x162971d (dependency fix-up)
for (CC *cc : group)
cc->getBasicBlock()->removeInstruction(cc);// 0x1629742 <- original IT48 deleted
}
renumberCoreBarriers(modules, getSubgraphId(modules[0])); // 0x162979f — global CB ids
}
Two helpers do the work that makes the lowering joint across cores rather than independent per-core:
-
extractAndGroupCCOps(0x1618440): walks every basic block of every core module (getFirstWithin/getNextInstFlat), selects insts withInstructionType == 48(thecmp [I+0x58],0x30test), and clusters them into avector<vector<InstCollectiveCompute*>>— one inner vector per logical collective (the matching SPMD copies thatlnc_splittercloned across cores). The grouping is exactly what turns a lowering into a joint cross-core operation: the swap on core A and the swap on core B are members of one group. -
getRemoteCores(self)(0x1617940): loopscore = 0..lnc-1, pushescoreinto the result iffcore != self(cmp eax,ebx; je skip); re-readslnceach iteration fromPassOptions+0x1A4. Returns{0..lnc-1}\{self}— the set of peers each core must exchange with. It is also what feeds thevnc_remote_addr_mapRemoteLocalTargetminting.
getPrecedingCoreBarrierOrLocalCCMap (0x1615e30) and precededByLocalCC (0x1615d50) record,
per core, the most recent preceding InstCoreBarrier or already-lowered local CC (a local
flag at Inst+0x1b0), so a redundant barrier is not inserted when a region is already
fenced. (The MapVector<CC*,Inst*> is destroyed at 0x162984f.)
Kind 0 — lowerSendRecv: the GPSIMD-or-DMA SBUF swap
A SendRecv is a paired send+recv "swap" between two cores' SBUF tiles — the TP/FSDP token
swap. This is the only kind where the data-mover engine is chosen per instruction; the other
two are fixed.
The engine-selection gate (0x161dfc3..0x161e8c5):
// lowerSendRecv @0x161cc70 — engine selection
bool compatible =
InstGPSIMDSB2SB::isCompatible(lnc, module, srcAP, dstAP, inst); // 0x161dfc3
// libBIR-native SB2SB shape/context predicate (libBIR @0x2931d0).
// FIRST check inside isCompatible is `cmp rdi,0x2` (ctx==2): the
// cross-core context must be exactly 2 cores-per-LNC.
if (!compatible) {
// -> DMA fallback
insertElement<InstDMACopy>(bb, ip, name); // 0x161eed6 / 0x161e3f5
} else {
int bpp = AccessPattern::getNumElementsPerPartition()
* DtypeSize[dtype]; // 0x161e892
// DtypeSize LUT @.rodata 0x1e01ce0, indexed by AP+0x30, bound <= 0x13 = 19.
int ebpp = InstGPSIMDSB2SB::getEffectiveBytesPerPartition(bpp, dtype); // 0x161e8b3
// dtype byte-pad (x4 / x2 / x1) applied via `shl eax,2` etc. (libBIR @0x291a50)
if (ebpp > sendRecvToGPSIMDMaxBytesPerPartition) { // 0x161e8c5
insertElement<InstDMACopy>(bb, ip, name); // -> DMA fallback
} else {
// -> GPSIMD path
BasicBlock::checkInsertionPointValid(...); // 0x161d721
Inst *g = insertElement<InstGPSIMDSB2SB>(bb, ip, name); // 0x161d72f
*(uint8_t *)(g + 0x90) = 1; // 0x161d772 (lowered/local marker)
}
}
// Both paths also emit an InstTensorCopy when an AP repack is needed
// (insertElement<InstTensorCopy>).
So the GPSIMD-vs-DMA choice is a two-stage filter:
-
Shape/context predicate.
InstGPSIMDSB2SB::isCompatible(libBIR@0x2931d0) must hold — the source/dest access patterns must form a legal SB2SB shape and the cross-core context must be 2 cores (cmp rdi,0x2is the first predicate). If not compatible at all → DMA. -
Byte budget. Even when compatible, the effective bytes-per-partition (
getEffectiveBytesPerPartition, libBIR@0x291a50, applying the dtype byte-pad factor) must not exceed the pass-sidecl::optvaluesendRecvToGPSIMDMaxBytesPerPartition. Above it → DMA; at or below it → GPSIMD. Thecl::optis namedsendrecv-to-gpsimd-max-bpp(.rodata 0x1c88377) with help text "Bytes/partition under which local swapping SendRecvs will be mapped to GPSIMDSB2SB" (.rodata 0x1d7ff48).
This pass-side knob sits under the hard libBIR ceiling
InstGPSIMDSB2SB::MaxBytesPerPartition = 0x400 = 1024 (libBIR .rodata @0x784a00, byte-read
00 04 00 00; assert GPSIMDSB2SB.cpp:52). They are two distinct gates: libBIR's isCompatible
enforces the 1024 ceiling unconditionally, while the pass knob is a softer threshold a
deployment can set lower to bias more swaps toward DMA. The numeric default of
sendrecv-to-gpsimd-max-bpp lives in the cl::opt initializer and is not read out here — only
the knob's existence and semantics are established. [SPECULATIVE on the default value.]
Addressing & sync. createRemoteAP(srcAP, allAPs, coreIdx, …, isRemote, denseMaps)
(0x161cad0) clones srcAP and calls getMemoryLocation to bind it to the peer core's
MemoryLocation; setLocation rebinds it. insertCoreBarrier(group, allCores, …) fences the
swap. Unless enableRemoteSemaphoreDMA (.bss 0x3e05160) is set, the run loop adds an extra
per-remote-core CoreBarrier for the SendRecv swap (the kind-0-only branch at 0x16296f6).
Kind 1 — lowerSendRecvCCE: the DMA-via-CCE swap
SendRecvCCE is a SendRecv routed through the Collective-Communication-Engine DMA path with
no GPSIMD option. The call set is the DMA-engine sibling of lowerSendRecv: createRemoteAP,
insertCoreBarrier, cloneToOutput, a PhysicalAccessPattern ctor, and
insertElement<InstDMACopy> only — there is no isCompatible probe and no GPSIMDSB2SB
symbol anywhere in the body. The same remote-AP minting and CoreBarrier sync as kind 0, but
the data mover is unconditionally InstDMACopy (the CCE descriptor flavour is selected later by
lower_dma / LegalizeCCEDMA). The whole difference from lowerSendRecv is the absence of
every GPSIMD symbol.
Kind 2 — lowerAllReduce: reduce-scatter + all-gather
The on-chip AllReduce is decomposed into a DMA-based reduce-scatter followed by an
all-gather across the LNC cores' SBUF, sharded along a free axis so each core reduces
1/lnc of the tensor.
// lowerAllReduce @0x1621a80 — reconstructed; call order from 0x1621a80..0x1624850
void lowerAllReduce(group, newInsts, denseMaps, unsigned long *semCounter) {
// Phase A — seed / initial copy
insertElement<InstDMACopy>(...); // 0x16221ab
auto ap = srcAP.clone();
auto ml = getMemoryLocation(...); ap.setLocation(ml); // 0x16222ab / 0x1622355 / 0x162236b
auto arg = ap.cloneToArgument(inst);
auto out = ap.cloneToOutput(inst); // 0x16223c6 / 0x16223fa
legalizeAPs(lnc, srcInst, newInst); // 0x1622417
// Phase B — REDUCE (scatter): each core owns one slice
auto remotes = getRemoteCores(self); // 0x1622dd8
auto scatterDim = getScatterDimension(ap); // x4 over candidate APs // 0x1622f8c..0x16230aa
insertElement<InstDMACopy>(/* reduce copy */); // 0x16234b6
// cloneToArgument / clone / getMemoryLocation / setLocation / cloneToOutput / legalizeAPs
// Phase C — BARRIER between reduce and gather
remotes = getRemoteCores(self); // 0x16236a8
insertElement<InstCoreBarrier>(...); // 0x1623899
// Phase D — ALL-GATHER: broadcast each core's reduced slice to all peers
remotes = getRemoteCores(self); // 0x162392a
insertElement<InstDMACopy>(/* gather copy */); // 0x1623fe2
createPartialAP(ap, allAPs, totalSize, partSize, coreIdx, scatterDim, ..., denseMaps); // x2
// // 0x1624134 / 0x1624209
legalizeAPs(...); // 0x1624244
remotes = getRemoteCores(self); // x2 // 0x162426e / 0x1624443
}
-
getScatterDimension(0x1615d40, tail-jmp togetFirstSplittableNonPartitionDimension0x1615c70): skips the partition dim (MemoryLocation::getPartitionDim @0x5f2c00) and walks the free dims (cmp [ML+0x110],0x4numDims check), returning the first dim with extent ≥ 4 that divides evenly acrosslnccores — the axis the reduction is sharded on. It has exactly one caller,lowerAllReduce, which calls it four times across two AP arms: one arm yields the scatterDim, the other the gatherDim. -
createPartialAP(0x1620fb0): builds a per-core sub-tensor AP —setOffset(coreIdx * partStride)/setPattern/setType/setLocation— so each core's DMA touches only its1/lncslice;getMemoryLocationthen binds it to the owning core's memloc. This is what makes the reduce a scatter (each core owns one slice) rather than an all-to-one. Its splittable-axis argumenta5is exactly thegetScatterDimensionresult, which downstream becomes the 2-bitgather_dim/scatter_dimDGE descriptor field.
The AllReduce therefore never uses GPSIMD — it is always one InstDMACopy per replica-core
plus InstCoreBarrier (_cb_end) fences. The body contains no isCompatible probe and no
GPSIMDSB2SB symbol at all. [INFERRED] The likely reason is that the reduced payload generally
exceeds the GPSIMD per-partition budget and the reduce-scatter shape is not an SB2SB swap, but
the binary only shows the absence, not the rationale.
Remote addressing — the vnc_remote_addr_map
getMemoryLocation(allAPs, idx, core, denseMaps&) (0x161ab30) is the cross-core address
minter that produces the vnc_remote_addr_map RemoteLocalTargets:
denseMapsis avector<DenseMap<MemoryLocation*,MemoryLocation*>>— one remote map per core;DenseMapBase::operator[]caches src-memloc → remote-memloc so the same remote target is reused across the group.- It resolves the target by tensor id via
MemoryLocationSet::getMemlocByTensorIdwithin the remote core'sMemoryLocationSet. - It mints/labels remote targets with names carrying
_remote_<core>and_setsuffixes (.rodata "remote" @0x1c86837,"_remote_" @0x1c88352,"_set" @0x1c867ac). - It marks the target with
MemoryLocation::setIsWrittenByAnotherCore(true)so downstream allocators (coloring_allocator_dram_shared,sync_shared_allocations) andvnc_linkkeep the cross-core writer alive and resolve the physical VNC address. The diagnostic strings " has a tensor id of " (0x1c88328) and " and isAllocated is " (0x1c8833d) spell out the tensor-id and allocation reasoning.
legalizeAPs (0x161b9a0) repacks the AP pattern
(PhysicalAccessPattern::setPattern) against the partition dim after retargeting so the cloned
remote/partial APs are HW-legal SBUF/DRAM patterns. insertCoreBarrier (0x161a150)
inserts InstCoreBarrier (IT 87) over the given core-id vector — the
semaphore/barrier between exchange phases — and renumberCoreBarriers (0x5e9080) re-assigns
globally-unique barrier ids across all core modules after lowering.
insert_ptcom_flat — the point-to-point completion tokens
PTCOM = point-to-point COMpletion. The pass inserts an InstTensorCompletion (IT 94)
per shared MemoryLocation at a convergence block, named <…>-PTCOM (.rodata "-PTCOM" @0x1c87e87). It is the wire-level p2p sync token — a completion / wait marker, not a data
mover: a consumer core blocks on it until a producer core has finished writing the shared tensor.
It complements the CoreBarrier (the collective fence) with finer point-to-point
producer→consumer completion, so a reader need not wait on a full barrier. insertPTCOMs emits
InstTensorCompletion plus attachWholeTensorInputAP and no copy at all.
// InsertPTCOMFlat::run(vector<Module>&) @0x16ad1c0 (linear)
void run(vector<unique_ptr<Module>> &modules) {
int subgraphs = modules.size() / lnc; // div @0x16ad208
// if (subgraphs == 1 && ...) { log; proceed; } // 0x16ad210
auto blocks = identifyPTCOMBlocks(modules); // 0x16ad66b -> vector<pair<ML*,BB*>>
auto lastW = identifyLastWrites(blocks); // 0x16ad679 -> per-(ML,BB) last writer
auto extended = extendLastWrites(lastW, modules); // 0x16ad69c -> propagate across cores
insertPTCOMs(extended); // 0x16ad6a7 -> emit InstTensorCompletion
}
// Pass-enable gate: cl::opt enablePTCOM (.bss 0x3df6a40, string @0x1600b1) checked by the
// pipeline/PassAdder (not the lambda); off => pass omitted.
The token-placement mechanism, in order:
-
identifyPTCOMBlocks(0x16aa040): per sharedMemoryLocation, collects the set of BBs that write it (DenseMap<ML*, DenseSet<BB*>>), then uses CFG dominance to pick the insertion block = the first common postdominator of all writer BBs (CFG::getImmediatePostdominator,CFG::getFirstCommonPostdominator). It filters viameetsPTCOMRequirementsandCFG::getSCCForBB, and asserts"loc2index.count(loc) > 0"(.rodata 0x1c8931e) and"blocks.size() == 1"(0x1c89337) — exactly one convergence block per memloc. -
meetsPTCOMRequirements(free fn0x16a5840): a candidate BB qualifies only if it is not inside a non-trivial loop SCC (CFG::getSCCForBB). Completions are hoisted out of loops so the p2p token fires once after the producing region, not every iteration. -
findMemLocOnCore0(0x16a7540) /findCBOnCore0(0x16a8010): the "FLAT" addressing scheme. All per-core modules reference the canonical core-0 names:findMemLocOnCore0looks a shared memloc up by string name in core-0's module;findCBOnCore0matches aCoreBarrieron core 0 byModule::getName+ name-hash (_Hash_bytes,memcmp). "Flat" means a single flat namespace keyed on core 0, so cross-core producer/consumer pairs are correlated by shared name rather than by per-core pointer identity — which differs after thelnc_splitterJSON clone, asaveJson→loadJsonround-trip that preserves names but not pointers. -
identifyLastWrites(0x16a9280) /extendLastWrites(0x16abd80):identifyLastWritesfinds, per (ML, BB), the last writer instruction;extendLastWritespropagates thatLastWriteInfoacross the per-core modules (so a completion on core B references the producing write on core A), asserting"extendedWrite != nullptr"(.rodata 0x1c89378).findNewLastWrite(0x16a8df0) re-locates the writer after re-lowering, gated bymeetsPTCOMRequirements. The exactLastWriteInfostruct layout is not pinned. [SPECULATIVE on the field offsets.] -
insertPTCOMs(0x16acc50): for each (ML, insertInst) pair, insertInstTensorCompletionbeforeBasicBlock::getFirstTerminator(0x16acd1f/0x16acd3c), thenattachWholeTensorInputAP(newInst, ML)(0x16acd4b) so the completion waits on the whole shared tensor;bir::isDataPathEngineroutes it to the proper sync engine. The"_sb2sb"string (0x1c88425) ties the completions to the on-chip SB2SB transport produced by the kind-0 lowering.
Why each pass runs twice — dual placement
Both passes are registered at two static pipeline slots:
- Default (post-schedule) phase, contiguous:
80 lower_local_collectives → 81 extend_shared_lifetimes → 83 sync_shared_allocations → … → 95 insert_ptcom_flat. - Alternate (pre-
post_sched) phase, contiguous:111 lower_local_collectives → 112 insert_ptcom_flat.
The selector is the cl::opt --run-shared-allocation-before-post-sched, help text verbatim:
"Run lower_local_collectives and coloring_allocator_dram_shared before post_sched. Only
relevant for trn2." The same pass body is registered at two static slots; the
pipeline builder positions the collective lowering + its PTCOM completions either after
post_sched (default) or before post_sched (when the flag is set — Trn2 / CoreV3 only).
The two are mutually exclusive at runtime, not both-run:
lower_local_collectivesis idempotent across placements because itsrun()deletes the original IT-48 CC op (removeInstruction @0x1629742) once lowered. Whichever slot fires first consumes the work; the other findsextractAndGroupCCOpsreturns no IT-48 ops → no-op. Thelnc == 1early-return likewise makes single-core a no-op at either slot.insert_ptcom_flatis paired with whicheverlower_local_collectivesslot is active (80→95 vs 111→112) so the PTCOM completions are inserted in the same scheduling regime as the collective DMAs they fence (pre- vs post-RA / post-sched addresses).
On Trn2 the shared-DRAM allocation (coloring_allocator_dram_shared) and the collective
lowering must run before post_sched so the post-RA scheduler sees the real cross-core
DMA / GPSIMD traffic and the shared allocations are colored coherently; on other arches the
default keeps them after post_sched. The dual registration is the mechanism that lets one
static pipeline serve both orderings.
Limits of this reading
The pass structure on this page — the kind dispatch and its CC+0xf8 field, the two-stage
GPSIMD-vs-DMA gate with its libBIR predicates and the 1024-byte ceiling, the DMA-only
reduce-scatter/all-gather shape of lowerAllReduce, the vnc_remote_addr_map minting with
setIsWrittenByAnotherCore, and the PTCOM insertion at the writers' first common postdominator —
is read off libwalrus.so and libBIR.so directly.
Three things are not:
- The
cl::optnumeric defaults.sendrecv-to-gpsimd-max-bppand its siblings are named and their semantics are clear, but the initializer values were not extracted. [SPECULATIVE] - The
LastWriteInfostruct layout. The helper call chain is clear; the field offsets are not. [SPECULATIVE] - The downstream lowering of the kind-≥3 remote collectives. They are fenced here and left intact; the cross-node decomposition happens in a different translation unit that is not part of this binary. [INFERRED] that it is the runtime collective library over the ICI fabric.
One softer structural claim: that PTCOM placement is hoisted out of loop SCCs and that
cross-core producer/consumer pairs are correlated by flat core-0 names rather than pointers
follows from the call set and the lnc_splitter JSON round-trip, not from a single decisive
instruction.