VN-Splitter and Shrink-ML — Allocator-Prep Footprint Reshaping

All symbols and addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310; cp310/11/12 are byte-identical Cython rebuilds). Everything here lives in neuronxcc/starfish/lib/libwalrus.so. For .text (0x62d660+) and .rodata (0x1c72000+) the virtual address equals the file offset; 0x5e9020–0x62d650 is the .plt thunk band, so every body address cited here is the real high body (0xd4xxxx–0xd7xxxx for VN-Splitter, 0x16bxxxx–0x182xxxx for Shrink), never its thunk. .data carries a +0x400000 delta. Other wheels differ — treat every address as version-pinned.

Abstract

Two consecutive walrus backend passes reshape the live-range graph before the coloring allocator runs, so that coloring can succeed without spilling. They allocate nothing themselves; they edit the set of MemoryLocation nodes the allocator will later color. Pass 24 (vn_splitter) makes every SBUF/PSUM node fit one SB tile and breaks access-disjoint clusters into independent nodes; pass 25 (shrink_ml, the ShrinkDN class) trims each node's physical extent down to the bytes actually touched. Together they hand the allocator a disjoint, minimal, SB-fitting node set. This is the seam after the pipeline (pass-pipeline-optlevels) and before the colorers (allocator-drivers).

The single most important structural fact — and the one that corrects a naive reading — is that vn_splitter is not a splitter; it is a split-then-fold body. VNSplitterPass::run (@0xd73890) calls two transforms in sequence: VNSplitter::runTransform (the SPLIT phase, called first at d73d6a) and then VerticalFusion::runTransform (the FOLD phase, called second at d74162). Both call sites are directly visible in the dispatch, in straight-line order with no reorder branch between them. The splitter carves a tensor along a dimension into N disjoint _VN_ sub-tensors; vertical fusion then folds producer→consumer op chains vertically to collapse intermediate-tensor liveness over the result. The two are inverse footprint moves — split increases node count by carving, fold reduces it by merging — and the pass runs split-first, so the fold operates on the chains the splitter produced (and any residual chains). The foldIntoPredecessor PassOption gates the fold stage.

CORRECTION — order is SPLIT-then-FOLD, not fold-then-split. An earlier revision of this page asserted the pass folds first and splits second. That is backwards. In VNSplitterPass::run (@0xd73890) the call sequence is VNSplitter::runTransform (SPLIT) at d73d6a first, then VerticalFusion::runTransform (FOLD) at d74162 second — straight-line, no order-selecting branch. The page's own embedded disassembly (§1) already labelled these two sites correctly; only the surrounding prose had the order inverted. The PGA driver runVNSplitterOnce (@0xd6f440) runs the same split(d6fad2)-then-fold(d6febf) order, so there is no order difference between the two drivers. The why of the post-split fold is INFERRED (collapse the chains the split produced / residual chains); only the order is CONFIRMED.

The splitter's piece count N is not a constant. The Profile-Guided Auto-Tuning loop (ProfileGuidedAutoTuning::runVNSplitterOnce(int vn_limit, float ratio, int perSplitLimit)) drives the same fold+split body with a tuned (vn_limit, ratio, perSplitLimit) triple; the float ratio bounds tolerated cross-piece duplication and the two integer caps bound the virtual-node count and pieces-per-node. So the split factor is a knob the autotuner searches, not a fixed property of the tensor — documented further in the PGA autotuner page (planned). The shrink_ml half is unconditional: it minimizes the per-partition live byte-extent of every shrinkable node, leaving reinterpret-cast and shifted-partition nodes intact.

For reimplementation, the contract is:

• The split↔fold dispatch: VNSplitterPass::run → VNSplitter::runTransform then VerticalFusion::runTransform, the fold gated by foldIntoPredecessor.
• The split decision: how analyze clusters access-patterns into ApGroups by overlapping partition sets, and the isValidForSB partition + byte gate that decides "splittable".
• The split rewrite: how split emits <orig>_VN_<k> sub-tensors, re-points cloned access-patterns, and rebases offsets (normalizeOffsetToNewNode).
• The PGA knob: (vn_limit, ratio, perSplitLimit) as the tuned split-factor triple, and the three reject-strings that pin each role.
• The footprint shrink: ShrinkDN::analyze probes writers for the true live window, shrinkDN re-packs partition stride and trims free dims; the set-level shrinkTensor annotates (shrinkDim, liveN).
• The invariants handed to the colorer (SB-fittability, use-def locality, clean graph, minimal footprint).

1. Pass identity and the split-then-fold dispatch

Both passes are BackendPass subclasses registered into the walrus pass set by generator lambdas. The registrars and pass-order positions are CONFIRMED from the registry strings and the S2-06 §3 pass-order map (vn_splitter=24, shrink_ml=25), corroborated by the metric VNSplitterDeadNodesCount cited in S2-07 §3.3.

The shrink_ml ↔ ShrinkDN identity is CONFIRMED two ways: the class-name string "ShrinkDN" is co-resident (×36 occurrences) with the pass-name "shrink_ml" (×13), and shrink_dn.cpp is the only source TU under shrink_ml/. The run body's own symbol confirms the mangled name _ZN9neuronxcc7backend8ShrinkDN3runERN3bir6ModuleE.

The structural correction is in the vn_splitter body. VNSplitterPass::run @0xd73890 does not call one splitter — it calls two transforms, split first then fold, and the dispatch is directly readable in the disassembly: [CONFIRMED]

d73d6a:  call   5f4930 <…VNSplitter12runTransformEv@plt>        // SPLIT phase
d74162:  call   60ab60 <…VerticalFusion12runTransformEv@plt>    // FOLD phase
d74249:  call   5ee600 <…VerticalFusionD1Ev@plt>                // dtors
d74252:  call   5ed870 <…VNSplitterD1Ev@plt>

The VNSplitter::runTransform call at d73d6a runs first; the VerticalFusion vtable (_ZTVN9neuronxcc7backend14VerticalFusionE) and typeinfo (_ZTIN…14VerticalFusionE) are referenced at d73e93/d73ea3, after which VNSplitterPass::run constructs a real VerticalFusion object and runs its runTransform at d74162. The same two transforms in the same split→fold order are driven by ProfileGuidedAutoTuning::runVNSplitterOnce (@0xd6f440: split at d6fad2, fold at d6febf); the autotuner re-runs the whole split+fold body per trial.

QUIRK — vn_splitter is a two-stage body, not a monolithic splitter. A reimplementer who treats vn_splitter as "split a tensor" will miss half the pass. The first stage is the virtual-node split (VNSplitter::runTransform at d73d6a); the second stage is VerticalFusion (the fold) at d74162, gated by the foldIntoPredecessor option. The split carves the over-sized nodes first; the fold then collapses the producer→consumer chains over the resulting (and any residual) node set. VerticalFusion::runTransform's internals are a separate symbol and are not traced on this page [GAP] — only its position in the dispatch and its gating option are pinned.

1.1 PassOptions knobs

The following keys are parsed from .rodata (single occurrence each, so each is a distinct option): vn_limit, vnLimit, perSplitLimit, noSplitDRAM, foldIntoPredecessor, splitSB/SplitSB, VNSplitterOnce. The autotuner entry runVNSplitterOnce(int vn_limit, float ratio, int perSplitLimit) consumes the three tuning knobs as a positional triple; its mangled member-function signature …ProfileGuidedAutoTuning…(f,i)…i,f,i… confirms the (int, float, int) parameter shape. [STRONG — signature pinned by the mangled symbol and the dedicated reject-strings below; runVNSplitterOnce body not separately traced.]

2. VN-Splitter — the SPLIT phase

2.1 runTransform — the orchestrator (@0xd5a3d0) [CONFIRMED]

runTransform first caches the SB geometry from the arch model, then runs a two-pass analyze/split sweep over each function's MemoryLocation set.

// VNSplitter::runTransform()  @0xd5a3d0
void VNSplitter::runTransform() {
    Arch arch = Module::getArch();              // d5a532
    ArchModel *am = getArchModel(arch);         // d5a53a
    // SBModel cached from the ArchModel chain am[+8]->[+0x10]->[+0x28]:
    this->sbBytesPerPartition = am_sb[+0];      // this+0x30  (per-partition SBUF budget)
    this->sbPartitionCount    = am_sb[+8];      // this+0x34  (e.g. 128)
    // option bytes off the ArchModel (this+0xa8 = ArchModel*):
    bool noSplitDRAM        = am[+0xa8][+0x150]; // cmpb $0,0x150(%rax)  @d5a63f
    bool foldIntoPredecessor= am[+0xa8][+0x151]; // bit at 0x151        @d5a5ef

    // PASS A — decide which nodes split (per bir::Function via allocs()):
    for (MemoryLocation &ml : Function::allocs()) {     // Function::allocs @0x5f4a30, d5a66e
        if (excluded_by_user(ml)) continue;             // "VNSplitter: Skipping <name> by user request."
        if (skip_non_SB_or_DRAM(ml, noSplitDRAM)) continue;
        bool splittable = analyze(&ml, ratio, vn_limit, perSplitLimit);  // d5af9b
        log("INFO (VNSplitter)  Analyzing <ml>");
        if (splittable)
            log("<ml> is splittable, shape=(<P> partitions, <B> bytes)!");
    }

    // PASS B — split + dead-node sweep:
    collect_internal_vnodes();                          // "Collected all the internal vnodes: …"
    for (MemoryLocation &ml : splittable_set) {
        kind = ml.tensorKind;
        bool in  = isTensorKindInput(kind);             // d5bcc1
        bool out = isTensorKindOutput(kind);            // d5bcd7
        split(&ml, ratio, foldIntoPredecessor, deadSet, keepMemLocs, keepInsts);
    }
    // erase the originals the split made redundant:
    for (auto &d : deadSet)
        if (is_MemoryLocation(d) && !keepMemLocs.count(d)) {
            ml = Function::getMemoryLocationByName(name(d));
            Function::removeMemoryLocation(ml, /*recursive=*/true);
        }
    log("INFO (VNSplitter): Split <N> nodes … into <P> pieces.");
    log("(SplitSB) total_sb_nodes = …, can_sb_nodes = …, total dead nodes = …, "
        "vn_limit = …, foldIntoPredecessor = …");
    addMetric(BackendMetricType, "VNSplitterDeadNodesCount", <u32>, Aggregation);  // d5aa24/d5aa9c/d5ac3b
}

The two option-byte reads are CONFIRMED in the disassembly: cmpb $0x0,0x150(%rax) at d5a63f is the noSplitDRAM test, and the 0x151 byte at d5a5ef is foldIntoPredecessor. The VNSplitterDeadNodesCount metric (string at VA 0x1d95cd9, mangled backend::VNSplitterDeadNodesCount) is emitted from three addMetric sites; its value equals the count of MemoryLocations removed this run (= the "total dead nodes" figure in the SplitSB log). [CONFIRMED metric exists and is emitted; value semantics STRONG.]

2.2 analyze — the split-shape decision (@0xd53020) [CONFIRMED]

analyze(MemoryLocation* ml, float ratio, int vn_limit, int perSplitLimit) is a four-stage funnel: pre-filter unsplittable tensors, sweep access-patterns for legality, cluster legal APs into ApGroups, then test each group against the SB budget.

// VNSplitter::analyze(ml, ratio, vn_limit, perSplitLimit)  @0xd53020
bool VNSplitter::analyze(MemoryLocation *ml, float ratio, int vn_limit, int perSplitLimit) {
    // (a) PRE-FILTER:
    if (isDRAMOnly(ml)) return false;   // DRAM-resident tensors are not SB-split (unless noSplitDRAM allows)
    if (isMx(ml)) special_handling();   // MX (block-scaled) tensors take a separate path

    // (b) AP-LEGALITY SWEEP over writers<AccessPattern>() and readers<AccessPattern>():
    for (AccessPattern *ap : ml.writers_and_readers()) {
        if (ap.isShiftedPartitionAccess())                  // reject: "due to non-contiguous AP." /
            return false;                                   //         "due to transpose (differing partiton step …)"
        if (PhysicalAccessPattern::hasCrossChannelAccess(ap))
            return false;                                   // reject: "due to access across channels."
        // other reject reasons share this string cluster:
        //   "due to reintepret_cast.", "due to partial copy (node size = …)",
        //   "due to CollectiveCompute access or BIRKernel access or indirect access.",
        //   "due to ISA pattern vector assignment needs non-TensorIndirect AP",
        //   "due to flow dependencies: …"
        // instr-type gates: copy(0x2e)/memset(0x2d)/transpose(0x2b)-class opcodes recognised;
        //   AP MemoryType and AP-kind@+0x18==1 select which APs count.
    }

    // (c) GROUPING — use-def clustering into ApGroups:
    collectAccesses(apMap, apaList);            // d52910: one APAccesses per AP (partition set + element range)
    for (APAccesses &a : apaList)
        if (a.is_read())
            findMatchingWrites(a);              // d4bb90: writes whose partition sets overlap this read
    for (APAccesses &a : apaList) {
        ApGroup *g = findMatchingGroup(...);    // d4de60 / findLastGroupWithOverlap d4d380
        if (!g) g = open_new_group();           // APs that never co-access a partition -> SEPARATE group
        g.addReadAPs / g.addWriteAPs(a);
    }

    // (d) PER-GROUP SHAPE + SB-FIT:
    bool any_splittable = false;
    for (ApGroup &g : groups) {
        g.collectStats();                        // memoised: highestElementInPartition, baseOffset, extents
        u32 bytesPP = g.getBytesPerPartition();  // = highestElementInPartition * dtypeBytes[ml.Dtype]
        auto dims   = g.getDimensions();         // [freeExtentBytes, highestElemPerPartition]
        bool valid  = g.isValidForSB({this->sbBytesPerPartition, this->sbPartitionCount}, strict);
        // SPLIT TRIGGER (synthesised): split when APs decompose into >1 disjoint-partition group,
        //   OR a single group overflows the SB per-partition budget, such that each piece fits SB.
        // ratio bounds tolerated cross-piece DUPLICATION; vn_limit caps running virtual-node count;
        //   perSplitLimit caps pieces per node.
        any_splittable |= ...;
    }
    return any_splittable;
}

The three PGA knobs are each pinned by a dedicated reject-string [STRONG]:

• ratio — bounds duplication. Logged as "elts/partition duplicated (<ratio> x)". A split that would replicate more than ratio× elements across pieces is rejected.
• vn_limit — caps the running virtual-node count. Reject-string "; exceed vn limit = … v.s. limit=<vn_limit>".
• perSplitLimit — caps pieces per node. Reject-string "Do not split the vn if this would result in any of the splits having to be placed in DRAM (due to perSplitLimit exceeded)" (CONFIRMED present in the binary).

These are exactly the three positional arguments of runVNSplitterOnce(int vn_limit, float ratio, int perSplitLimit), which is why the autotuner can sweep the split factor: tightening ratio or perSplitLimit reduces N, loosening them lets the splitter carve more pieces. The split factor N is therefore PGA-tuned, not fixed.

2.3 The ApGroup accumulator and APAccesses

An ApGroup is the per-logical-sub-tensor accumulator — one group becomes one _VN_ sub-tensor. Fields reconstructed from offsets [STRONG]:

getDimensions() (@0xd4dde0) returns the 2-entry shape SmallVector<u32>{ group+0x4c, group+0x04 } = [freeExtentBytes, highestElementInPartition]. getBaseOffset() (@0xd4dcb0) is CONFIRMED to collectStats() then return this+0x08 (mov 0x8(%rbx),%eax; ret).

An APAccesses is the use-def grouping unit, one per PhysicalAP: it holds the PhysicalAP*, a set<int> partitionIndices, and an element range (the ctor @0xd4be50 calls getElementRange()). The grouping predicate is overlaps() (@0xd4a5a0), defined as !is_set_disjoint(a.partitions, b.partitions) (is_set_disjoint @0xd49120) — two APs overlap iff their accessed-partition sets intersect. APs that never share a partition land in separate groups, so they become separate sub-tensors with non-conflicting live ranges. This is INV-2 (use-def locality).

2.4 isValidForSB — the SB-fit gate (@0xd501e0) [CONFIRMED]

The gate packs the SBModel into a single 64-bit register: low 32 bits = sbBytesPerPartition budget, high 32 bits = sbPartitionCount. This is directly visible in the prologue — mov %rsi,%rax; sar $0x20,%rsi extracts the high half. [CONFIRMED]

// ApGroup::isValidForSB(SBModel sb, bool strict)  @0xd501e0
bool ApGroup::isValidForSB(SBModel sb /* esi: low32=byteBudget, high32=partCount */, bool strict) {
    u32 e = getHighestElementInPartition();              // this+0x04
    // walk write & read APs, count those with AP[+0x18]==1 (placed/real-AP kind),
    //   compute used bytesPerPartition.

    // PARTITION GATE  @d50532:
    //   if  partitionDim >= SBModel.partitionCount  ->  FAIL
    if (partitionDim >= sb.partitionCount /* 0x8(%rax) */) return false;   // cmp 0x8(%rax),%edx; jge -> fail
        // diagnostic: "ERR: channel count > … channels."

    // BYTE GATE:
    //   if  bytesPerPartition > SBModel.sbBytesPerPartition  ->  FAIL
    //   log "apGroup id <id> : isValidForSB = <bool>; exceed vn limit = <n> v.s. limit=<lim>"
    return (bytesPerPartition <= sb.sbBytesPerPartition);
}

The partition gate is CONFIRMED in the disassembly: at d50526 the group loads mov 0x50(%rbx),%rax (the SBModel record), and d50536: cmp 0x8(%rax),%edx; jge fails when the partition dim (%edx from 0x6c(%rsp)) reaches the partition count at 0x8 of the SBModel record. This is the precondition that lets the colorer assume every node fits one SB tile (INV-1). The running per-AP byte sum is SIMD-inlined, so the gate semantics (partition-count + byte-budget) are CONFIRMED but the exact per-AP sum form is STRONG.

2.5 split — sub-tensor creation (@0xd57c80) [CONFIRMED]

For each valid ApGroup, split creates a new physical MemoryLocation named <orig>_VN_<k> (the literal "_VN_" is in .rodata), re-points every AP of the group at it, rebases offsets, and records the original for removal.

// VNSplitter::split(ml, ratio, foldIntoPred, deadSet, keepMemLocs, keepInsts)  @0xd57c80
void VNSplitter::split(MemoryLocation *ml, float ratio, bool foldIntoPred,
                       set<variant<MemoryLocation*,Instruction*>> &deadSet,
                       DenseSet<MemoryLocation*> &keepMemLocs,
                       DenseSet<Instruction*> &keepInsts) {
    for (ApGroup &g : valid_groups) {     // or simple DRAM->SB case when groups.size()==1
        string newName = ml.name() + "_VN_" + to_string(k);
        MemoryLocation *vn = Function::addMemoryLocation(
            newName,
            partitionCount,            // from group / partition dim
            bytes,                     // group.bytesPerPartition * partitions
            ml.getDtype(),
            ml.getPartitionDim(),
            g.getDimensions(),         // tensorShape SmallVector<…,4>
            /*flags*/, /*optional*/, ml.getDebugInfo(),
            ml.getBankId(), ml.getAddress(), ml.getBasePartition(),
            ml.getMemoryType(),        // SB(16)/PSUM(32) per ml  (see D-D05 bit-flags)
            ml.getMemoryAddressSpace());

        for (PhysicalAP *ap : g.{writeAPs, readAPs}) {
            PhysicalAP *ap2 = _deepcopy(ap);                 // neuronxcc::_deepcopy clone
            // ap2 = cloneToOutput(inst) / cloneToArgument(inst) / PhysicalAP(inst) ctor
            ap2.setLocation(vn, b, b);                       // bind AP to the new sub-tensor
            int newOff = g.normalizeOffsetToNewNode(ap.offset);
            ap2.setOffset(newOff);                           // rebase into sub-tensor
            // patch Pattern dims via getPatternRef(); MemoryLocation::updateDimensions(jj,m)
        }
        deadSet.insert(variant(ml));                          // record original for removal
        // log "INFO (VNSPlitter): Moved <ml> from DRAM to SB…" /
        //     "WARNING: ML <name> has unused partitions! old: … ; new: …"
    }
}

The offset rebase is CONFIRMED exact at normalizeOffsetToNewNode (@0xd4dc80): sub 0x8(%rdi),%esi then idivl 0x50(%rdi) — i.e. (ap.offset − group.baseOffset[+0x08]) / group.partitionStride[+0x50]. The group stats accumulate via addAPtoStats (@0xd49200): partition index (ap.offset[+0xD0] − base) / stride, range stat group[+0x54] = max(·, getRange[0]) (via pmaxsd), free extent group[+0x4c] = max(·, ap.basePtr[+8] + adj), and group[+0x04] = max(·) + 1 (highest element index + 1).

getBytesPerPartition (@0xd4dc40) = getHighestElementInPartition() × dtypeBytes[ml.Dtype], indexed into the dtype byte-size table at 0x1de6c40 (20×int64: u8=1, u16=2, fp32/i32=4, i64/fp64=8, …, Dtype field ≤ 0x13, matching the BIR AccessPattern Dtype range in D-E12).

2.6 Dead-node elimination and VNSplitterDeadNodesCount [STRONG]

A node is dead after a split when its every reader/writer AP has been re-pointed to the new _VN_ sub-tensors, leaving it with no real use — that is the original over-sized MemoryLocation and any wholly-subsumed copy/memset instruction. split collects these into deadSet (std::set<variant<MemoryLocation*, Instruction*>>); the DenseSet<MemoryLocation*> / DenseSet<Instruction*> arguments are the keep-alive sets (referenced elsewhere, must survive). runTransform's tail resolves each dead entry by name (getMemoryLocationByName) and erases it via removeMemoryLocation(ml, /*recursive=*/true). The VNSplitterDeadNodesCount metric equals the count removed this run, which is the "total dead nodes" figure in the SplitSB log — confirming S2-07 §3.3. This is INV-3 (clean graph).

3. Shrink-ML / ShrinkDN — the footprint shrink

Where the splitter carves nodes to fit, shrink_ml trims each node to its true size. A MemoryLocation was sized for its declared shape; ShrinkDN discovers the live window — the largest byte-range any access actually touches in a partition — and physically rewrites the node down to that window. The pass runs at order 25, after the splitter, so each _VN_ piece is shrunk independently.

3.1 ShrinkDN::run — the driver (@0x16c4db0) [CONFIRMED]

// ShrinkDN::run(bir::Module&)  @0x16c4db0
void ShrinkDN::run(bir::Module &m) {
    // iterate every MemoryLocation: filter Storage -> MemoryLocationSet -> MemoryLocation
    for (MemoryLocation &ml : module_memory_locations(m)) {
        if (ml.hasReinterpretCast()) continue;   // reinterpret nodes left untouched
        analyze(ml);                             // compute live extents
        shrinkDN(ml);                            // rewrite extents
    }
}

3.2 ShrinkDN::analyze — the liveness probe (@0x16c07a0) [CONFIRMED]

// ShrinkDN::analyze(bir::MemoryLocation &ml)  @0x16c07a0
void ShrinkDN::analyze(bir::MemoryLocation &ml) {
    if (ml.isMx()) { /* special / skip */ }
    u32 maxUsed = 0;
    for (AccessPattern *ap : StorageBase::writers<AccessPattern>(ml)) {
        if (!AccessPattern::classof(ap)) continue;
        if (ap.isShiftedPartitionAccess()) bail();    // cannot shrink a shifted layout
        if (!isShrinkable(ap)) bail();
        u32 used = AccessPatternUtils::getElementRangeInPartition(ap);   // 108a130
                 // = getRange(ap.Pattern) * dtypeBytes[ml.Dtype]   (table @0x1df0e60, idx Dtype@+0x30 ≤0x13)
        maxUsed = max(maxUsed, used);
    }
    // record maxUsed as the shrink dimension/extent (consumed later via getShrinkDim)
    // assert shrink_dn.cpp:216
}

The metric ShrinkDN minimises is the per-partition live byte-extent: the largest range any access touches in a partition. isShrinkable(ap) (@0x16bef10) is AccessPattern::isPartitionContiguous() — only contiguous-partition APs may shrink; shifted/non-contiguous nodes bail and stay at their declared size. [CONFIRMED footprint metric; STRONG on exact result storage.]

3.3 ShrinkDN::shrinkDN — the extent rewrite (@0x16c38a0) [CONFIRMED]

// ShrinkDN::shrinkDN(bir::MemoryLocation &ml)  @0x16c38a0
void ShrinkDN::shrinkDN(bir::MemoryLocation &ml) {
    u32 newExtent = analyze_result;                     // live elements-per-partition
    for (AccessPattern *ap : writers_and_matching_readers(ml)) {
        updateChannelStep(ap, newStep, partitionIdx);   // re-pack partition stride
        shrinkAP(ap);                                    // trim free-dim extents to newExtent
    }
    MemoryLocation::updateDimensions(jj, m);             // ×2: write reduced stride + free size back
    // propagate ml.getDebugInfo(); per-set rewrite via indirect callback fn(MemoryLocation*,…)
    // assert shrink_dn.cpp:444
}

// updateChannelStep(ap, newStep, idx)  @0x16bdd20  (exact)
void updateChannelStep(PhysicalAP &ap, u32 newStep, u32 idx) {
    ap.getPatternRef().Pattern[0].step = newStep;       // channel (partition) step
    ap.setOffset(newStep * idx + oldBase);              // re-pack partitions packed-contiguous
    // assert shrink_dn.cpp:292
}

// shrinkPAP(ap, n)  @0x16c2f00  (exact) — invoked by shrinkAP @0x16c3020
void shrinkPAP(PhysicalAP *ap, u32 n) {
    Pattern v = copy(ap.Pattern[0]);                    // keep the partition dim
    for (dim = 1; dim < numDims; ++dim)
        v.Pattern[dim] = canonical{step, num} from const@0x1dd1fe0, free-count = n;
    ap.setPattern(v);                                   // PhysicalAccessPattern::setPattern
    // -> the AP now addresses only `n` live elements
}

updateChannelStep re-packs the partition stride so partitions stay contiguous after the free dim shrinks (setOffset = newStep * idx + oldBase); shrinkPAP keeps the partition dim and rewrites the free dims to the canonical {step, num} with the free count set to the shrunk extent n. shrinkAP (@0x16c3020) computes the live count via getElementRangeInPartition + getMemoryLocationSet, guarded by isPartitionContiguous (×5), then calls shrinkPAP(ap, liveCount). Together these write the reduced extents back through MemoryLocation::updateDimensions, so the tensor occupies fewer bytes per partition. This is INV-4 (minimal footprint).

3.4 shrinkTensor — the set-level analyzer [STRONG]

shrinkTensor(bir::Function&) (@0x1826e00) is the logical counterpart that operates at the MemoryLocationSet level rather than on physical extents. For each MemoryLocationSet, it inspects SymbolicAccessPattern::getBlockAddrs() and Instruction::getLoopnest() over the set's APs (a hashtable<u64, Instruction*> of block→writer), derives the shrink dimension and the count of live blocks, then records:

MemoryLocationSet::setShrinkDim(dim);   // producer
MemoryLocationSet::setLiveN(n);         // live block count

shrinkTensor(bir::Module&) (@0x1829e90) just loops functions. The cross-pass channel is getShrinkDim() — only a thunk lives in libwalrus (0x5fb650 for MemoryLocation, 0x5efe20 for MemoryLocationSet); the real body and the shrink-dim selection heuristic live in libBIR and are not traced here [GAP]. So shrinkTensor is the analysis that annotates each set with (shrinkDim, liveN), and ShrinkDN is the transform that reads getShrinkDim() and physically rewrites the per-MemoryLocation APs. The producer/consumer split (setShrinkDim vs getShrinkDim) proves the linkage. [STRONG]

QUIRK — VncLink is not part of this chain. VncLink (run(Module&) @0x16477c0, run(vector<unique_ptr<Module>>&) @0x1648de0) is the per-Neuron-Core link pass that joins per-LNC module shards. It shares no callgraph edge with VNSplitter/ShrinkDN; the only overlap is the "VN" prefix, which here means virtual-neuron-core, not virtual-node. A reimplementer must not conflate the two. [STRONG]

4. Allocator-prep role — invariants handed to the colorer

The coloring allocator (allocator-drivers) and SBSizeLegalization color MemoryLocations onto a fixed SB grid of sbPartitionCount × sbBytesPerPartition. The two prep passes establish the preconditions that make coloring tractable:

Ordering rationale. vn_splitter(24) runs before shrink_ml(25): the splitter carves each over-sized node so each piece is independently shrinkable; shrink_ml then trims each piece to its live extent; then the disjoint, minimal, SB-fitting node set is handed to the colorer. Within vn_splitter the order is split-then-fold: the virtual-node split runs first, then VerticalFusion folds the producer→consumer chains over the splitter's output (CONFIRMED by call order; the precise rationale for folding after splitting is INFERRED).

The PGA loop closes over this. Because runVNSplitterOnce re-runs the whole split+fold body with a tuned (vn_limit, ratio, perSplitLimit), the autotuner is effectively searching the split factor N against the downstream colorer's success/cost — tightening the knobs to fewer, larger pieces or loosening them to more, smaller pieces. The split↔fold duality plus this tuned knob is what makes vn_splitter an adaptive footprint reshaper rather than a fixed transform.

5. Confidence summary and gaps

• CONFIRMED by direct disassembly: the split↔fold dispatch (VNSplitterPass::run calling VNSplitter::runTransform at d73d6a first, then VerticalFusion::runTransform at d74162); the isValidForSB partition gate (cmp 0x8(%rax),%edx; jge); the packed-SBModel register split (sar $0x20,%rsi); normalizeOffsetToNewNode (off − base)/stride (sub/idivl); getBaseOffset returning this+0x08; the two option-byte reads at +0x150/+0x151; the perSplitLimit reject-string; the _VN_, SplitSB, is splittable, and VNSplitterDeadNodesCount strings; the ShrinkDN::run symbol and namespace.
• STRONG (pinned by strings/signature but with SIMD-inlined or partly-inferred arithmetic): the (int, float, int) PGA signature and the per-knob role mapping; the isValidForSB running byte-sum form; the ApGroup field map at +0x4c/+0x54; the shrinkTensor ↔ ShrinkDN producer/consumer linkage; VncLink being unrelated.
• GAP (not traced on this page): VerticalFusion::runTransform internals — the actual fold heuristic — are a separate symbol; the getShrinkDim selection heuristic lives in libBIR (only thunks here); the exact runVNSplitterOnce body and the autotuner's search policy (planned PGA page).
• Honest re-verify ceiling. The structural skeleton (which functions, what they gate, the invariants) is solid against the binary. The two heuristics not in libwalrus (vertical-fold condition, shrink-dim selection) are named and positioned but not reconstructed; a reimplementer would need the libBIR bodies and VerticalFusion::runTransform to fill them.