Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Lowering On-Chip Collectives — lower_local_collectives / insert_ptcom_flat

All symbols and addresses on this page apply to neuronx_cc 2.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 in neuronxcc/starfish/lib/libwalrus.so (build-id 92b4d331a42d7e80bb839e03218d2b9b0c23c346, stripped; all names resolve from .dynsym, there is no .symtab). For .text (0x62d660+) and .rodata (0x1c72000+) the virtual address equals the file offset; .data carries a +0x400000 delta. The IR-side handles (InstGPSIMDSB2SB::isCompatible, getEffectiveBytesPerPartition, the MaxBytesPerPartition ceiling) live in libBIR.so (build-id a9b1ea38…). The pre-dumped .c sidecars for these two passes are PLT re-export thunks (_0x5e9xxx); the real bodies disassemble from .text at 0x1615xxx/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):

SymbolAddressRole
run(bir::Module&)0x1615680per-module shim (stores 0x2, no-op)
run(vector<unique_ptr<Module>>&)0x1624850real worker (~0x5e50 bytes)
getRemoteCores(unsigned long self)0x1617940{0..lnc-1}\{self}
getScatterDimension(PhysicalAccessPattern const&)0x1615d40tail-jmp alias → getFirstSplittableNonPartitionDimension 0x1615c70
extractAndGroupCCOps(...)0x1618440clusters IT-48 ops into SPMD groups
lowerSendRecv(...)0x161cc70CollectiveKind 0
lowerSendRecvCCE(...)0x161f130CollectiveKind 1
lowerAllReduce(...)0x1621a80CollectiveKind 2
createRemoteAP(...)0x161cad0clone AP onto a peer core
createPartialAP(...)0x1620fb0per-core sub-tensor AP (the scatter slice)
getMemoryLocation(...)0x161ab30remote-memloc minter (vnc_remote_addr_map)
insertCoreBarrier(...)0x161a150emit InstCoreBarrier (IT 87)
updateDeps(...)0x1619620dependency fix-up before deletion
sendRecvToGPSIMDMaxBytesPerPartition.bss 0x3e050a0cl::opt byte-budget value

InsertPTCOMFlat (factory register_generator_insert_ptcom_flat__ @0x16ae030):

SymbolAddressRole
run(bir::Module&)0x16a5710per-module shim (stores 0x2, no-op)
run(vector<unique_ptr<Module>>&)0x16ad1c0real worker
identifyPTCOMBlocks(...)0x16aa040per-memloc convergence block (postdom)
identifyLastWrites(...)0x16a9280per-(ML,BB) last writer
extendLastWrites(...)0x16abd80propagate last-write across cores
insertPTCOMs(...)0x16acc50emit InstTensorCompletion (IT 94)
findMemLocOnCore0(...)0x16a7540flat core-0-name lookup
findCBOnCore0(...)0x16a8010match a CoreBarrier on core 0
meetsPTCOMRequirements(...)0x16a5840free fn — non-loop-SCC gate
enablePTCOM.bss 0x3df6a40pass-enable cl::opt
enableRemoteSemaphoreDMA.bss 0x3e05160extra-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 plus CoreBarrier synchronization. These are the collectives whose traffic stays inside the LNC group of on-die cores.

  • Kind ≥ 3 (ReduceScatter / AllGather / AllToAll / AllToAllV / Permute family) — the remote / ICI path. The collective op is not decomposed here. The run loop still calls insertCoreBarrier per 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 a lowerXXX helper; kind ≥ 3 reaches nothing but insertCoreBarrier. [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 with InstructionType == 48 (the cmp [I+0x58],0x30 test), and clusters them into a vector<vector<InstCollectiveCompute*>> — one inner vector per logical collective (the matching SPMD copies that lnc_splitter cloned 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): loops core = 0..lnc-1, pushes core into the result iff core != self (cmp eax,ebx; je skip); re-reads lnc each iteration from PassOptions+0x1A4. Returns {0..lnc-1}\{self} — the set of peers each core must exchange with. It is also what feeds the vnc_remote_addr_map RemoteLocalTarget minting.

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:

  1. 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,0x2 is the first predicate). If not compatible at all → DMA.

  2. Byte budget. Even when compatible, the effective bytes-per-partition (getEffectiveBytesPerPartition, libBIR @0x291a50, applying the dtype byte-pad factor) must not exceed the pass-side cl::opt value sendRecvToGPSIMDMaxBytesPerPartition. Above it → DMA; at or below it → GPSIMD. The cl::opt is named sendrecv-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 to getFirstSplittableNonPartitionDimension 0x1615c70): skips the partition dim (MemoryLocation::getPartitionDim @0x5f2c00) and walks the free dims (cmp [ML+0x110],0x4 numDims check), returning the first dim with extent ≥ 4 that divides evenly across lnc cores — 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 its 1/lnc slice; getMemoryLocation then 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 argument a5 is exactly the getScatterDimension result, which downstream becomes the 2-bit gather_dim/scatter_dim DGE 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:

  • denseMaps is a vector<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::getMemlocByTensorId within the remote core's MemoryLocationSet.
  • It mints/labels remote targets with names carrying _remote_<core> and _set suffixes (.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) and vnc_link keep 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 shared MemoryLocation, 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 via meetsPTCOMRequirements and CFG::getSCCForBB, and asserts "loc2index.count(loc) > 0" (.rodata 0x1c8931e) and "blocks.size() == 1" (0x1c89337) — exactly one convergence block per memloc.

  • meetsPTCOMRequirements (free fn 0x16a5840): 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: findMemLocOnCore0 looks a shared memloc up by string name in core-0's module; findCBOnCore0 matches a CoreBarrier on core 0 by Module::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 the lnc_splitter JSON clone, a saveJsonloadJson round-trip that preserves names but not pointers.

  • identifyLastWrites (0x16a9280) / extendLastWrites (0x16abd80): identifyLastWrites finds, per (ML, BB), the last writer instruction; extendLastWrites propagates that LastWriteInfo across 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 by meetsPTCOMRequirements. The exact LastWriteInfo struct layout is not pinned. [SPECULATIVE on the field offsets.]

  • insertPTCOMs (0x16acc50): for each (ML, insertInst) pair, insert InstTensorCompletion before BasicBlock::getFirstTerminator (0x16acd1f / 0x16acd3c), then attachWholeTensorInputAP(newInst, ML) (0x16acd4b) so the completion waits on the whole shared tensor; bir::isDataPathEngine routes 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_collectives is idempotent across placements because its run() deletes the original IT-48 CC op (removeInstruction @0x1629742) once lowered. Whichever slot fires first consumes the work; the other finds extractAndGroupCCOps returns no IT-48 ops → no-op. The lnc == 1 early-return likewise makes single-core a no-op at either slot.
  • insert_ptcom_flat is paired with whichever lower_local_collectives slot 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::opt numeric defaults. sendrecv-to-gpsimd-max-bpp and its siblings are named and their semantics are clear, but the initializer values were not extracted. [SPECULATIVE]
  • The LastWriteInfo struct 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.