DMA Legalization — Strided / CCE / Generic-Indirect

All symbols and addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310 wheel; cp310/11/12 are byte-identical for the C++ core logic). The three passes live in neuronxcc/starfish/lib/libwalrus.so (.text/.rodata base 0x62d660/0x1c72000, VA == file offset; .data adds a +0x400000 delta). The BIR instruction enums (bir::InstructionType) and AccessPattern layout live in libBIR.so. Treat every address as version-pinned. See Build & Version Provenance.

Abstract

Three backend passes stand between the high-level BIR a kernel arrives in and the descriptor-legal IR the DMA encoders can serialize. Each takes one over-general DMA opcode and rewrites it into a sequence the hardware DGE can actually issue:

• legalize_strided_dma (order 18, gated on the gen-stride-dma CLI flag) rewrites a partition-strided InstLoad (IT 19) into a contiguous DMA into a fresh SBUF tile plus an InstTensorCopy that re-applies the stride out of SBUF — but only when the materialized tile is /32-partition-aligned and fits the per-arch SBUF byte budget.
• legalize_cce_dma (order 19, unconditional) splits a collective-compute-engine InstDMACopy (IT 32) whose source count exceeds board.device.core.MaxCceDmaSource into a chain of InstDMACopys, each consuming at most that many sources, splitting the per-source scale-constant vector alongside. This is the pass that makes every CCE DMA satisfy the verifier's checkCCEDMALegality gate.
• lower_generic_indirect (order 10, unconditional) expands a "generic" multi-dimensional indirect access (GenericIndirectLoad IT 42 / GenericIndirectSave IT 44) into the concrete ReadVarAddr → TensorScalarPtr → Memset → TensorScalarPtr → IndirectLoad(43)/IndirectSave(45) IR sequence the indirect-gather encoders consume.

The bar for this page is that a reimplementer can reconstruct exactly when each pass fires and exactly what it emits — the split thresholds (/32 alignment, SBUF byte budget, MaxCceDmaSource) are the decisive mechanism, and they are the headline content. What each pass does not do is equally important: none of them assigns a DGE queue or picks a wire descriptor. Queue/engine assignment is the later assign_hwdge_engine / alloc_queues passes; the 16B-vs-20B-vs-MX descriptor choice is made downstream at codegen from the lowered op's dimensionality and dtype. These three passes only make the IR DMA concrete and hardware-shaped.

This page is the earlier legalization stage. The later materialization (lower_dma, the DGE descriptor build) is a separate, planned page (8.34).

All seven body-frame addresses, the three factory invokes, and the dtype tables are read directly off the IDA disassembly and Hex-Rays decompilation of libwalrus.so. CONFIRMED.

The shared pass skeleton

All three passes are the same shape, so it is worth pinning once. Each is a neuronxcc::backend::BackendPass subclass registered through a register_generator_<name>__ factory lambda. The factory __cxa_demangles the pass's own class name (e.g. "N9neuronxcc7backend18LegalizeStridedDMAE"), allocates the 96-byte BackendPass object, and seeds its logging::Logger with the registry string ("legalize_strided_dma" / "legalize_cce_dma" / "lower_generic_indirect"). CONFIRMED — factory bodies read end-to-end at the three invoke addresses above.

BackendPass::run(bir::Module&):

// run() — common to all three passes [CONFIRMED, decompiled all three]
void run(bir::Module &M) {
  // push a logging attribute scope { "module_name" = M.getName() }
  Function *main = M.getFunction("main");        // "main" visited first
  visitFunction(main);
  for (Function *F : M.functions())
      if (F != main) visitFunction(F);
  // (legalize_cce_dma only) post-pass marker — see below
}

The CRTP IRVisitor::visit walks each BasicBlock, recurses into the structured-control bodies of IT 105 Loop / IT 106 DynamicForLoop / IT 108 DoWhile, and switches on the instruction node's IT field (node +80) to dispatch the one DMA opcode that pass owns. The three opcodes are mutually disjoint:

IT 43 (IndirectLoad) is intentionally not dispatched by lower_generic_indirect — it is already the concrete op the pass produces. CONFIRMED (run() + IRVisitor::visit decompiled for all three; opcode ownership matches the bir::InstructionType ordinals).

legalize_strided_dma — strided load → SBUF tile + TensorCopy

What it owns

A partition-strided DMA load is an InstLoad (IT 19) whose partition dimension is not contiguous — i.e. the access has an inter-partition gap. The hardware DGE cannot issue such a load directly; the descriptor wants a partition-contiguous source. The pass's job is to manufacture one: stage the data into a freshly-allocated SBUF tile via a contiguous DMA, then re-apply the original strided pattern with an on-chip InstTensorCopy.

This pass runs only under the gen-stride-dma CLI gate (order 18); the other two are unconditional. CONFIRMED via the P-3-04 pass-order walk.

The arch budget read

run() reads a per-arch SBUF-partition byte capacity from the hardware model and stows it on the visitor as the legalization budget. The dereference chain is getArchModel(Module::getArch()) → (+8) → (+16) → (+40) — the same Board → Device → Core singleton tree the SBUF/PSUM geometry page documents. The impl object the visitor carries is:

// LegalizeStridedDMAImpl object layout [STRONG — wiring reconstructed from run()]
struct Impl {
  void   *loggerScope;   // +0
  uint64  copyCounter;   // +8   init 0; ++'d per emitted copy for a unique decimal suffix
  uint64  sbufByteBudget;// +12  per-arch SBUF partition byte capacity
  Function *curFn;        // +16
};

The rewrite

The AccessPattern field offsets below are the decompiler byte-offsets and match the bir::AccessPattern layout: dtype at +0x30, MemoryLocation* at +0x38, the APPair[] stride/num pairs at +0x50 (each pair 16 B: step@+0, num@+8), the pattern dimension count at +0x58 (1..4), the byte offset at +0xD0.

// LegalizeStridedDMAImpl::visitInstLoad(bir::InstLoad&)  @0xb5ddc0  [CONFIRMED]
void visitInstLoad(InstLoad &I) {
  AccessPattern &inAP  = I.getArgument<PhysicalAccessPattern>(0);
  AccessPattern &outAP = I.getOutput<PhysicalAccessPattern>(0);

  // R0 — already-legal short-circuit: partition dim has no inter-partition gap
  if (inAP.isPartitionContiguous()) return;          // PLT off_3DCF5E0 → libBIR.so

  // R1 — output must be a placeable MemoryLocation
  MemoryLocation *outLoc = *(MemoryLocation**)((char*)&outAP + 56);
  assert(outLoc != nullptr);
  assert(*(int*)((char*)outLoc + 272) == 4);         // LLVM isa<> cast tag, not MemoryType

  // R2 — destination base partition must be on a /32 (DMA-quadrant) boundary
  int basePart = (outAP.offset / outLoc->elemStride) + outLoc->getBasePartition();
  if ((basePart & 0x1F) != 0) return;                // not /32-aligned → leave as-is

  // R3 — linearized element span of the input access pattern
  uint64 span = 1;
  for (int i = 1; i < inAP.num_dims; ++i)            // num_dims at inAP+0x58
      span += inAP.APPair[i].step * (inAP.APPair[i].num - 1);

  // R4 — SBUF byte-budget fit
  int dt = inAP.dtype;                               // inAP+0x30
  assert(dt <= 0x13);                                // else "Unknown dtype"
  uint64 bytes = DtypeSize[dt] * (inAP.offset_elems + span);   // DtypeSize = qword_1DE10A0
  if (bytes > this->sbufByteBudget) return;          // does not fit → leave as-is (R4 fail)

  // R5 — DECOMPOSE: contiguous DMA into a fresh SBUF tile + InstTensorCopy back out
  unsigned n = ++this->copyCounter;                  // unique decimal suffix
  MemoryLocation scratch = curFn->addMemoryLocation("<dtype>-tile");  // 1-D, span elems
  *(int*)((char*)&scratch + 216) = 16;               // MemoryType = SB (16)
  InstTensorCopy &tc = BB.insertElement<InstTensorCopy>("<orig>-copy");
  // tc input  ← copy of inAP's pattern (offset = inAP.offset)
  // tc output ← scratch SBUF location, contiguous
  // rewrite ORIGINAL load: source AP now reads the contiguous scratch tile; dst kept
  PhysicalAccessPattern::resetEvaluatedAps(inAP);
  PhysicalAccessPattern::resetEvaluatedAps(outAP);
}

The three thresholds, in plain terms

If R2 or R4 fails, the load is left unmodified — deferred to the verifier / later legalization rather than rewritten. The /32 figure is the DMA-quadrant constant (128 partitions ÷ 4); the SBUF byte budget is a per-arch immediate stamped into the Statebuf record (see SBUF/PSUM geometry). CONFIRMED for the rule structure and the 0x1F mask; the runtime budget value lives in the arch-model ctor, not in a shipped JSON, so its numeric value is INFERRED per arch.

The DtypeSize table qword_1DE10A0 (.rodata 0x1de10a0, 20 qwords) is {1,1,2,1,1,1,1,1,4,4,2,2,2,2,4,4,4,4,8,8} (byte size per Dtype 0..19). CONFIRMED.

Cross-lib note. AccessPattern::isPartitionContiguous resolves through PLT slot off_3DCF5E0 into libBIR.so (its body is not in libwalrus.so .text); its truth-semantics (true = partition dim has no inter-partition gap) are inferred from the symbol name and call site. STRONG.

legalize_cce_dma — split a CCE DMACopy over MaxCceDmaSource

What it owns

A collective-compute-engine DMA is an InstDMACopy (IT 32) for which bir::isCCEDMA is true. The hardware CCE engine can fan in only a bounded number of sources per copy. That bound is board.device.core.MaxCceDmaSource — and it is the same field the verifier's checkCCEDMALegality reads, asserting I.num_arguments() <= MaxCceDmaSource (NeuronAssertion code 329 — that is the assertion code, not the cap value). This pass exists precisely so that every CCE DMA satisfies that gate. CONFIRMED — identical ArchModel.device.core+0x50 dereference in both the pass and the verifier.

The post-pass marker

After visiting all functions, run() iterates every Function and sets FunctionAttribute #21 = "stage_legalize_cce_dma_completed" (written through the boost::variant attribute map at Function+240). This is the pipeline-bookkeeping flag the later stages key off — part of the stage_*_completed attribute cluster. CONFIRMED (attribute enum value 21; string present in .rodata).

The split

// LegalizeCCEDMAImpl::visitInstDMACopy(bir::InstDMACopy&)  @0xb61280  [CONFIRMED, full body read]
void visitInstDMACopy(InstDMACopy &I) {
  if (!bir::isCCEDMA(I)) return;                      // gate @0x607110
  Module &M  = *I.getFunction()->getModule();
  int archMax = ArchModel.device.core.MaxCceDmaSource;          // *(...core + 0x50)
  int nSrc    = count(I.args) + count(I.constants);
  if (archMax >= nSrc) return;                        // already legal

  // shape/dtype legality diagnostics over the source APs
  int   outDt    = outAP.dtype;
  Shape outShape = outAP.getShape();
  for (AccessPattern &inAP : I.sourceAPs()) {
      if (inAP.dtype != outDt)    reportError("Input AP dtype does not match output AP dtype");
      if (inAP.shape != outShape) reportError("Input AP shape does not match output AP shape");
  }
  bool hasConst = (outAP.argKind == 4) && (I.constants.size() != 0);  // CCE w/ immediate scalars

  // SPLIT LOOP — chunk sources into groups of <= MaxCceDmaSource
  int chunk = 0;
  while (count(I.sources) > archMax) {
      InstDMACopy *clone = I.clone();                 // vtable slot 18 (vt+0x90) clone
      clone->rename("<name>_imm_output_" + str(chunk));
      // move the first archMax source ARGUMENTS into the clone
      Instruction::spliceArguments(/*to*/clone, /*from*/&I, /*count*/archMax);
      // split the per-source constant (scale) vector alongside, when present:
      //   copy the first archMax floats to clone; erase them from I (SmallVector<float>::erase);
      //   the partial 1.0f scale (0x3F800000) is appended per chunk on the trailing remainder.
      MemoryLocation partLoc = M.getFunction(...)->addMemoryLocation("<name>_imm_output_" + str(chunk));
      *(int*)((char*)&partLoc + 224) = 7;             // partial-output kind marker = 7
      PhysicalAccessPattern partAP(clone);
      partAP.clone(outAP); partAP.offset = 0;
      partAP.cloneToArgument(&I);
      partAP.cloneToOutput(clone);
      // re-link clone into the BB list + symbol table (hash-table remove/insert)
      ++chunk;
  }
  // remainder (<= archMax sources) stays as the original InstDMACopy.
}

The threshold, in plain terms

The single split condition is nSrc > MaxCceDmaSource, where nSrc = #arguments + #constants. Each emitted chunk consumes at most MaxCceDmaSource sources; the per-source scale-constant SmallVector<float> is partitioned in lockstep, and each chunk gets a freshly-minted "<name>_imm_output_<N>" partial-output MemoryLocation tagged kind 7 at field +224. No queue or DGE-engine assignment happens here. CONFIRMED — full 920-line body read; helpers (clone via vtable slot 18 / vt+0x90, spliceArguments, addMemoryLocation, symbol-table relink) identified by call site.

Tag caveat. Field +224 = 7 (partial-output kind) is a distinct tag from the +216 MemoryType field that the strided pass sets to 16 (SB). The strided pass's MemoryType=16 is a named enum (SB); the cce pass's +224=7 is observed-not-enum-named. STRONG / +224=7 value INFERRED as a kind tag.

The exact numeric value of MaxCceDmaSource lives in the per-arch EngineInfo/Target ctor in .text (not in any shipped JSON); only the dereference path is pinned. GAP.

lower_generic_indirect — generic indirect → concrete IR sequence

What it owns

The front end (penguin/HLO) emits a "generic" indirect access — GenericIndirectLoad (IT 42) or GenericIndirectSave (IT 44) — carrying a multi-dimensional index expression and unflattened address tensors. The hardware indirect DMA wants a single 1-D index, a 2-D address tensor, and one element/one address per partition. This pass bridges that gap, expanding the generic op into a concrete IR chain and forcing every index/offset/address AP to the hardware shape. CONFIRMED — IRVisitor::visit dispatches IT 42 → codegenIndirectLoadSave, IT 44 → visitInstGenericIndirectSave (a trivial forwarder, thunk 0x5ecf20 → codegenIndirectLoadSave).

Input invariants asserted

Before lowering, the pass asserts (these strings are present verbatim in .rodata):

• the instruction is GenericIndirectLoad (42) or GenericIndirectSave (44);
• "Indirect dimensions must be one less than instruction argument number";
• "Compiler middle end must have already flattened IndexAP to be only one AP" — the index AP arrives already 1-D;
• "GenericIndirectLoad/Save must have 1 element per partition for IndexAP" (getNumElementsPerPartition(indexAP) == 1);
• "indirect DMA's AddrAP should be 2D";
• "Addr Tensor partition dim must match its tensor shape".

These exactly mirror the verifier's checkIndirectShape / checkTensorIndirection ruleset — the pass produces precisely the shape the verifier then re-checks. CONFIRMED (all assertion strings located in libwalrus.so).

The lowering chain

// LowerGenericIndirectImpl::codegenIndirectLoadSave(Instruction*, BasicBlock*)  @0xb580e0  [CONFIRMED]
AccessPattern codegenIndirectLoadSave(Instruction *inst, BasicBlock *bb) {
  bool load = (inst->IT == 42);
  // (a) per-indirect-dim flat byte address: for each indirect dim (outer→inner) emit an
  //     InstTensorScalarPtr "<base>_TSPAddGenericAddr" accumulating step*idx; the flat
  //     element multiplier = product of trailing dim sizes.
  // (b) InstReadVarAddr (IT 41) "<base>_ReadVarAddr": read the target DRAM tensor base
  //     var address into the 2-D AddrAP; scratch/dst memloc tagged "<base>_DstTensor".
  // (c) element size: dt = DstTensorMemLocSet.getDtype(); assert dt <= 0x13;
  //     elemBytes = DtypeSize[dt]  (table qword_1DE0E20, identical mirror).
  //     If the target DRAM tensor is flattened to 1-D, IndirectDims must also be 1-D
  //     ("target DRAM tensor is flattened, IndirectDims must be expanded...").
  // (d) InstTensorScalarPtr "<base>_TSPAddAddr": add the byte address (immediate ptr-type
  //     Argument, setType 15) to the AddrAP. op-mode fields 240=6 / 288=4.
  // (e) InstMemset "<base>_MemsetTileOffset": zero a per-tile offset tensor (immediate
  //     setType 14); "_DstTensor" memloc, dims updated to 4-wide.
  // (f) InstTensorScalarPtr "<base>_TSPAddOffset": add the tile offset (op-mode 240=4).
  // (g) FINAL concrete indirect op:
  if (load) {
    InstIndirectLoad &out = emit<InstIndirectLoad>("<base>_IndirectLoad");   // IT 43
    addMemoryLocation("<base>_ScratchSpaceTensor");                          // 1-wide dim
    // args = { addr-result AP, original index AP, offset AP }; output cloned from inst output
  } else {
    InstIndirectSave &out = emit<InstIndirectSave>("<base>_IndirectSave");   // IT 45
    // args = { offset-result AP, addr AP, output AP }; output ← inst arg
  }
  return sub_B54210(finalAP);   // finalize the surviving AP for load vs save
}

So the generic op (IT 42/44) becomes, in order: ReadVarAddr(41) → per-dim TensorScalarPtr address accumulation → Memset tile-offset → TensorScalarPtr offset add → IndirectLoad(43) / IndirectSave(45), with the index, offset, and address APs all forced to "1 element / 1 address per partition, contiguous data". CONFIRMED. The impl carries a worklist SmallVector (this+1 begin / this+4 size / this+5 cap) to which codegenIndirectLoadSave appends inst before lowering. STRONG.

Descriptor selection is downstream — not here

This pass emits the IR-level concrete ops IndirectLoad (43) and IndirectSave (45). It does not choose the wire descriptor. The encoder picks among INDIRECT16B (3-D AP, assignAccess3D), INDIRECT20B (4-D AP, assignAccess4D), and MXINDIRECT16B (MX scale present, assignIndirectPatternForMX) downstream at codegen, from the lowered op's AP dimensionality, dtype, and DGE type — see Indirect-Gather Descriptors and the IndirectLoad/Save codegen op 0xC4/0xD6. lower_generic_indirect's entire job is to make the IR indirect DMA concrete and hardware-shaped; the 16B-vs-20B (dim count) and MX (scale) selection happens after. STRONG — pass body produces 43/45; the descriptor encoders are codegen-resident.

The DtypeSize table this pass reads, qword_1DE0E20 (.rodata, 20 qwords), is the byte-identical mirror {1,1,2,1,1,1,1,1,4,4,2,2,2,2,4,4,4,4,8,8}. CONFIRMED.

How the three relate

Three observations tie the set together. Opcode ownership is disjoint — strided owns Load (19), cce owns DMACopy (32), generic-indirect owns GenericIndirect (42/44) — so no DMA op is touched by two of them. All three honor the same arch-model root getArchModel() → Board+0x8 → Device+0x10 → Core for their HW limits (strided reads the SBUF byte budget; cce reads MaxCceDmaSource; generic-indirect reads only the dtype tables). And none assigns a queue or DGE engine — bir::DGEType stays Unassigned(3) from the InstDMA ctor; HWDGE-engine and queue assignment are the later assign_hwdge_engine / alloc_queues passes. STRONG.

Confidence & gaps

• CONFIRMED: the three-pass set, their order (10/18/19), the gating of order 18 on gen-stride-dma, every body-frame and thunk address, the three factory invokes, opcode ownership (IT 19/32/42-44), the /32 (0x1F) strided alignment mask, the nSrc > MaxCceDmaSource cce split condition, the full cce split body, the generic-indirect lowering chain and all its named IR ops and string literals, FunctionAttribute #21, and the three byte-identical DtypeSize mirrors.
• STRONG: the strided Impl object layout (offset wiring reconstructed from run()), isPartitionContiguous semantics (cross-lib PLT, name + call site), the generic-indirect worklist, the downstream descriptor-selection claim (pass produces 43/45; encoders are codegen-resident).
• INFERRED / GAP: the numeric values of the per-arch SBUF byte budget and board.device.core.MaxCceDmaSource — both live in EngineInfo/Target ctors in .text, not in any shipped JSON; only the dereference paths are pinned. The cce +224 = 7 partial-output value is an observed tag, not an enum-named constant.

Cross-References

• 2.21 DMA-Family Encoding & Descriptors — the descriptor wire format these passes legalize toward; the IndirectLoad/Save codegen ops (0xC4/0xD6) that consume the IT 43/45 this page produces.
• 2.07 Indirect-Gather Descriptors — INDIRECT16B / INDIRECT20B / MXINDIRECT16B, the downstream descriptor the generic-indirect lowering feeds; the 16B-vs-20B-vs-MX selection this page defers.
• SBUF / PSUM Bank Geometry — the Board → Device → Core arch-model tree and the per-arch SBUF byte budget the strided pass tests against.
• The Compile Pipeline at a Glance — where these three passes sit in the overall backend order.
• 8.34 lower_dma (DGE materialization) — planned — the later pass that materializes the legalized DMA into concrete DGE descriptors; this page is the earlier legalization stage.