KLR→BIR Codegen of the Dynamic / Runtime Access Pattern

Every symbol, address, struct offset, and vtable VA on this page applies to neuronx_cc 2.24.5133.0+58f8de22 (cp310). The two functions reconstructed here — KlirToBirCodegen::assembleDynamicInfo and KlirToBirCodegen::extractDynamicApToStandAlone — live in neuronxcc/starfish/lib/libwalrus.so (md5 1d93972b81e619ce6d178a0e4b9003b3, 64,973,024 bytes); the back-end codegen bodies are real here even where the nki_klr_sim driver carries only PLT stubs. VA == file offset for .text (verified at 0x62d660) and .rodata (0x1c72000); the bir::DynamicAPINFO setters and bir::Function::getRegisterByName are imported U symbols satisfied by libBIR.so. cp311/cp312 share the ABI but drift the VAs.

Abstract

A KLR access is dynamic when its address is not a compile-time constant — when the offset into SBUF/PSUM is a value the device only knows at execution time. There are exactly two shapes of runtime offset, and they are mutually exclusive: a scalar register-backed offset (a runtime integer read from a device register — a dynamic loop base, a data-dependent slice start) and a vector gather index (a tensor of indices that selects rows — the classic gather/scatter, and the MoE expert-routing path). This page reconstructs the KLR→BIR codegen slice that turns each shape into the symbolic address algebra the rest of the back-end materializes.

The center of the page is assembleDynamicInfo @0xf20230. Given a KLR operand whose static geometry has already been encoded into an InstArg (the kind-1 PhysicalAccessPattern that codegenAccess mints), it builds the bir::DynamicAPINFO companion that carries the runtime term. The decisive branch is the scalar-kind-7 vs vector-kind-3 split: the scalar offset becomes a bir::BirIntRuntimeValue (pelican Expr kind 7, the IntRuntimeValueBase with its regref at +0x40 documented in Pelican Index/Runtime); the vector offset becomes a pelican::IndirectArgExpr (kind 3, an arg_id handle). Both are wrapped in a QuasiAffineExpr and paired with an accumulated-shape coefficient, so the net address the device evaluates is addr = c + Σ_axis coef_axis · expr_axis.

The sibling function extractDynamicApToStandAlone @0xf1e2c0 is the index pre-materialization: for the gather case it hoists the index tensor out of the parent AP into its own standalone InstArg and strips the vector-offset flag, so the index is encoded exactly once and the parent AP is no longer dynamic at that slot. It is a codegen-time rewrite — it does not build a DynamicAPINFO and does not emit register-ALU ops. The two functions are siblings on the dynamic-DMA path, not nested: assemble… packages the address algebra onto the data AP; extract… hoists the index operand. Both bottom out in codegenAccess as the leaf encoder.

This is the KLR-codegen slice of the dynamic-shape thread. It mints the symbolic (kind-2) form; the Symbolic-AP → Register-ALU lowering (LowerAP::convertSymAP) later turns it into a kind-3 RegisterAccessPattern with real InstRegisterAlu micro-ops; the Dynamic-Shape Synthesis page is the end-to-end view. Front-end DMA bucketing is in front/pipeline.

1. What "dynamic" means at the KLR level

The predicate and the BAP payload

A KLR Access of kind==4 carries a klr::BirAccessPattern (BAP). A BAP is dynamic iff it holds a runtime offset, and the predicate is a one-line byte OR (isAccessDynamic(BAP) @0xf14030, CONFIRMED — the decompiled body reads result = *(u8*)(*a1 + 64); LOBYTE(result) = *(u8*)(*a1 + 88) | result):

bool isAccessDynamic(klr::BirAccessPattern *bap) {  // 0xf14030
    return bap[+0x40]      // SCALAR  runtime offset present
         | bap[+0x58];     // VECTOR  runtime offset present
}

The two wrappers that reach it gate on the KLR object kind first (CONFIRMED from the symbol-table delegation chain):

• isAccessDynamic(klr::Access) @0xf16540 — true only if Access.kind == 4 (it is a BAP).
• isAccessDynamic(klr::TensorRef) @0xf16760 — true only if TensorRef.kind == 1 (an SBUF/PSUM ref); a DRAM kind==4 ref can never be dynamic, then it delegates to the inner Access.

The two presence bits are mutually exclusive. A dynamic BAP carries either a scalar runtime offset or a vector gather index, never both, and the invariant is enforced before any work begins (_validateOnlyOneOfScalarAndVectorDynamicOffsetIsProvided @0xf140a0, CONFIRMED disasm movzx ecx,[rax+0x40]; cmp [rax+0x58],cl; jz fault):

void _validateOnlyOneOfScalarAndVectorDynamicOffsetIsProvided(BAP *bap) {  // 0xf140a0
    assert(bap[+0x40] != bap[+0x58]);   // EXACTLY ONE of {scalar, vector}
}

QUIRK — the XOR is on raw presence bytes. The check is a byte inequality, not a logical XOR over normalized booleans: bap[+0x40] != bap[+0x58]. If both were 1 (or both 0) it fires. Because each byte is written as a 0/1 flag this behaves as XOR, but a reimplementer must store the flags as canonical 0/1 — any other truthy value would make the inequality pass spuriously when both are "present."

The klr::BirAccessPattern field map

Reconstructed from the [rsi+NN]/[rax+NN] loads in both function bodies (STRONG — no dedicated KLR-struct strand maps the full BAP header, so unread slots are not pinned; the read slots below are CONFIRMED against the disasm of assembleDynamicInfo and extractDynamicApToStandAlone):

The castToBirAp(TensorRef) @0xf168d0 → castToBirAp(Access) @0xf166b0 chain is the typed accessor that recovers a BAP from a ref: it asserts ref.kind==1 ("Logic fault: Attempt to cast a TensorRef that is not an BirAP"), takes the inner Access at ref+0x08, asserts Access.kind==4, and returns the BAP at Access+0x08.

2. assembleDynamicInfo — the dynamic-info packager

Signature and role

// _ZN9neuronxcc7backend16KlirToBirCodegen19assembleDynamicInfoE…  @0xf20230  [CONFIRMED]
bir::InstArg KlirToBirCodegen::assembleDynamicInfo(
        std::unique_ptr<bir::DynamicAPINFO> &dyn,   // OUT: the companion to fill
        std::shared_ptr<klr::TensorRef>      ref,   // the dynamic operand
        bir::InstArg                        &arg);  // the ALREADY-built STATIC InstArg

The signature is byte-confirmed against nm -DC. The caller has already encoded the operand's static geometry into arg — an InstArg is {MemLoc*, SmallVector<APPair,4> pattern, partition, dtype, valid} (the kind-1 form; see codegen-accesspattern). assembleDynamicInfo builds the bir::DynamicAPINFO companion that carries the runtime address algebra, then returns a fresh InstArg by value — the operand the caller binds via addArgument. The companion hangs off PhysicalAccessPattern+0x1D8 and serializes to the dynamic_ap_info key.

Prologue — common to every branch

// 0xf20230..0xf20338  [CONFIRMED]
BAP *bap = castToBirAp(ref);                                   // 0xf20280
_validateOnlyOneOfScalarAndVectorDynamicOffsetIsProvided(bap);// 0xf202b4

dyn->setActualPattern(arg.pattern @ arg+8);                   // 0xf202d7  ⭐
    // actual_ap = the STATIC-resolved SmallVector<APPair> already in `arg`:
    // the concrete (step,num) geometry codegenAccess produced. DynamicAPINFO+0x00.

PelicanContext *ctx = *(*(this+0x38) + 0xA0);                 // 0xf202e6
    // this+0x38 = Module/Function;  +0xA0 = its PelicanContext.

// Build ONE pelican IndirectArgExpr (kind 3) up front:
auto *idx = (Expr*) tc_new(0x28);                             // 0xf202f9
pelican::Expr::Expr(idx, /*kind=*/3, ctx);                    //   esi=3  (CONFIRMED imm)
idx->vptr = &_ZTVN7pelican15IndirectArgExprE + 0x10;          // 0xf20309 (off_3DAA110)
idx[+0x20] = 1;                                               // 0xf20310  arg_id = 1
inc_ref(idx);

switch (presence) {                                           // 0xf20334 / 0xf2033e
    case bap[+0x58]:  goto VECTOR;    // §2.2 gather index tensor
    case bap[+0x40]:  goto SCALAR;    // §2.3 register-backed offset
    default: throw "Logic fault, this branch should never be taken"; // 0x1d1d900
}

The arg_id = 1 store is CONFIRMED (mov qword [rbp+0x20], 1 @0xf20310). The IndirectArgExpr is built unconditionally in the prologue and reused as the per-axis term in both the vector path and the scalar codegenAccess sub-path; only the scalar register-name path mints a different expr (kind 7). The default arm is a redundant fail-safe — the XOR validator already excludes it.

2.2 Vector branch — a runtime gather index (kind 3)

// loc_F20860..  [CONFIRMED]
VECTOR:
  Access *offAcc = bap[+0x48];                  // 0xf20860 the gather-index inner Access
  codegenAccess(offAcc);                        // 0xf20887 ENCODE the index AP (result
                                                //   dropped; resolves the index MemLoc)
  assert(bap[+0x60] /*indirect_dim*/ == 0);     // 0xf208a0
      // else throw "When using vector offset, the only supported indirectDim is 0"
      //   (0x1d1d8c0). A gather index addresses the W/partition axis only.
  dyn[+0x74] = 0;                               // 0xf208b2  indirect_dim = 0

  // indirect_dim_max_index = number of gatherable rows (the index range):
  if (MemLoc_kind(arg.MemLoc) == 1)             // 0xf208ce scalar-name MemLoc
      dyn[+0x70] = (*(ML+0x40))[0];             //   loc_F20B18
  else {
      dyn[+0x70]  = **(ML+0x100);               // 0xf208ec  ML+0x100 "dims"[0] = elem count
      bytes       = **(ML+0x108);               // 0xf208f9  ML byte size
      dt          = MemoryLocation::getDtype(ML);// 0xf20906 assert dt<=0x13 else "Unknown dtype"
      dyn[+0x70]  = bytes / dtype_bytes[dt];    // 0xf20922  ÷ per-dtype size
  }

  // per-axis term {IndirectArgExpr(idx), accumulated-shape coef}:
  coef = dyn[+0x70];
  dyn->aff_expr @+0x50 = { { QuasiAffineExpr(RefPtr(idx)), coef } };  // 0xf20940 / 0xf209a9
  dyn->setStaticOffsetPortion(bap[+0x10]);      // 0xf209b9  c = static residue
  populate_return_InstArg();

The per-dtype size table is dtype_bytes[] = {1,1,2,1,1,1,1,1,4,4,2,2,2,2,4,4,4,4,8,8} (20 entries, indices 0..0x13), read at qword_1DE98C0 and byte-verified directly from .rodata. indirect_dim_max_index is recovered from the index tensor's own extent — either its leading "dims" entry, or its byte size divided by the element size — and bounds the gatherable rows.

QUIRK — the gather index can only address the partition axis. indirect_dim is asserted == 0 for the vector path and the DynamicAPINFO field is hard-written 0. The error string "the only supported indirectDim is 0" is the guard. A reimplementer who wants to gather along an inner free axis cannot express it through this path — the vector offset is structurally a W/partition-axis selector.

2.3 Scalar branch — a register-backed runtime offset (kind 7)

// 0xf20348..  [CONFIRMED]
SCALAR:
  dyn[+0x74] = 0;                               // 0xf20365  indirect_dim = 0
  Access *offAcc = bap[+0x38];                  // the scalar-offset inner Access
  if (offAcc->kind == 1) goto REGISTER_NAME;    // 0xf20371  a bare register name
  else                   goto CODEGEN_ACCESS;   // a full Access

§2.3a Register-name path (loc_F20580..0xf2075b, CONFIRMED) — the offset is a bare register name:

REGISTER_NAME:
  std::string name = std::string(offAcc.bytes @ +0x8 .. +0x10);   // 0xf205bf
  Register *r = Function->getRegisterByName(name);                // 0xf205e4
  if (!r) throw "Cannot find register: " + name + " for scalar offset";  // 0x1c7e0a0/0x1c7e0b7

  auto *rv = (Expr*) tc_new(0x48);                                // 0xf20619
  pelican::Expr::Expr(rv, /*kind=*/7, ctx);                       //   esi=7  (CONFIRMED imm)
  rv->vptr = &_ZTVN3bir18BirIntRuntimeValueE + 0x10;             // 0xf2062d
  rv[+0x40] = r;                                                  // 0xf20634  regref = Register*  ⭐
  rv[+0x20] = (i64)-1;                                            // 0xf20638  no static index

  coef = _computeAccumulatedShape(*bap+0x28 /*Shape*/, bap[+0x60] /*dim*/);  // 0xf2068f
  dyn->aff_expr @+0x50 = { { QuasiAffineExpr(RefPtr(rv)), coef } };          // 0xf206b3
  dyn->setStaticOffsetPortion(bap[+0x10]);                        // 0xf20753  c
  populate_return_InstArg();

rv is a bir::BirIntRuntimeValue — the bir leaf grafted onto pelican kind 7 (IntRuntimeValueBase). Its regref at +0x40 is the same field the Pelican Index/Runtime page documents as the dynamic-shape spine: a Register* parked inside an Expr node, treated downstream as an opaque loop-invariant scalar. The mov [r15+0x40], r13 store @0xf20634 is CONFIRMED, and the vtable VA 0x3d8de40 matches nm -DC (vtable for bir::BirIntRuntimeValue).

§2.3b Codegen-Access path (0xf2037a..0xf204d6, CONFIRMED) — the scalar offset is a full Access rather than a register name:

CODEGEN_ACCESS:
  codegenAccess(offAcc);                                         // 0xf203b8 resolve its MemLoc
  coef = _computeAccumulatedShape(*bap+0x28, bap[+0x60]);        // 0xf2040f
  dyn->aff_expr @+0x50 = { { QuasiAffineExpr(RefPtr(idx)), coef } };  // 0xf2043b/0xf204c1
      // reuses the prologue IndirectArgExpr (kind 3) — NOT a fresh kind-7 node.
  dyn->setStaticOffsetPortion(bap[+0x10]);                       // 0xf204d1  c
  populate_return_InstArg();

QUIRK — "scalar" does not always mean kind 7. The scalar branch splits on offAcc->kind == 1. Only the register-name sub-path mints the kind-7 BirIntRuntimeValue; the full-Access sub-path reuses the prologue's kind-3 IndirectArgExpr. The kind-7-vs-kind-3 distinction is therefore not "scalar vs vector" at the top level — it is "register-name scalar" (kind 7) vs "everything else" (kind 3, minted once and shared by the vector path and the scalar full-Access path).

2.4 The DynamicAPINFO it builds

Field-by-field, every offset is written by assembleDynamicInfo and matches the independently-mapped DynamicAPINFO layout byte-for-byte:

tensor_indirect_arg_id (+0xF0) and num_tensor_indirect_indices (+0xF4) are not written here — they belong to the SymbolicAccessPattern companion produced downstream.

The net runtime address algebra (the pelican semantics) is:

addr = c  +  Σ_axis  coef_axis · expr_axis

where expr_axis is a runtime value (IndirectArgExpr index or BirIntRuntimeValue register), coef_axis = _computeAccumulatedShape is the per-axis element stride, and c is the static residue. This is the offset-expression the device evaluates per access.

3. _computeAccumulatedShape — the per-axis coefficient

// 0xf152e0  [CONFIRMED]
long _computeAccumulatedShape(std::shared_ptr<klr::Shape> shape, int dim) {
    assert(0 <= dim && dim <= shape->numDims + 1);   // else range fault sub_6D6FDE
    node = shape.dims_head;                          // shape+0x08 intrusive dim list
    for (k = 0; k < dim; k++) node = node->next;     // walk to `dim`
    long coef = 1;
    for (; node != list_sentinel; node = node->next)
        coef *= node->extent;                        // node+0x10 (u32) dim extent
    return coef;
}

The accumulated stride is the product of all dim extents from dim inward — the row-major stride of axis dim in element units. It is paired as the long half of each aff_expr entry, so addr += coef · expr. The dim argument is the BAP's indirect_dim (bap+0x60). (CONFIRMED — the walk-then-multiply loop is in the decompiled body.)

The offset-expr vector assigners

sub_F13410 @0xf13410 and sub_F1FB60 @0xf1fb60 are both std::vector<std::pair<bir::QuasiAffineExpr, long>>::operator= — a deep copy-assignment with a 40-byte element stride (a QuasiAffineExpr is 32 B = {RefPtr<Expr>@+0, SmallVector<LoopAxis*>@+0x8} plus the long coef@+0x20, padded to 5 qwords). They refcount-bump each RefPtr<Expr> via the nanobind intrusive counter. sub_F1FB60 builds the stack working vector; sub_F13410 installs it into the companion's persistent aff_expr field at DynamicAPINFO+0x50. (CONFIRMED — the += 40 loop stride and the QuasiAffineExpr destructor on old elements identify the type.)

4. extractDynamicApToStandAlone — the index pre-materializer

Signature and role

// _ZN9neuronxcc7backend16KlirToBirCodegen28extractDynamicApToStandAloneE…  @0xf1e2c0  [CONFIRMED]
bir::InstArg KlirToBirCodegen::extractDynamicApToStandAlone(
        std::shared_ptr<klr::TensorRef> ref);

This is the gather-index pre-materialization. For a dynamic vector-offset access it pulls the index tensor out of the parent AP and encodes it as a self-contained Access whose address is the runtime offset, returning the InstArg the consuming op binds. It then strips the vector-offset flag so the parent AP is no longer dynamic at that slot and the index is encoded exactly once.

// full 104-ins body  [CONFIRMED]
bir::InstArg extractDynamicApToStandAlone(std::shared_ptr<klr::TensorRef> ref) {
  assert(isAccessDynamic(ref));                       // 0xf1e303
      // else throw "Compiler Logic Fault: cannot extract dynamic ap from a static access!"
      //   (0x1d1d768)
  BAP *bap = castToBirAp(ref);                         // 0xf1e352

  if (bap[+0x40] /*scalar flag*/)                      // 0xf1e36b — must be VECTOR, not scalar
      throw "When unique_indices=True, it's expected that the vector_offset is set. "
            "Check the nisa.dma_copy call";            // 0x1d1d7b0
  if (!bap[+0x58] /*vector flag*/)                     // 0xf1e3c8
      throw "Logic Fault: dynamic access should have vector offset.";  // 0x1d1d818

  Access *offAcc = bap[+0x48];                         // the gather-index inner Access (= bap+72)
  InstArg result = codegenAccess(offAcc);              // 0xf1e3ad ⭐ ENCODE as standalone InstArg

  if (bap[+0x58]) {                                    // 0xf1e3c8 re-test vector flag
      ctrl = bap[+0x80];                               // old gather-index shared_ptr ctrl
      bap[+0x58] = 0;                                  // STRIP the vector-offset flag
      if (ctrl) release(ctrl);                         // drop its old reference
  }
  return result;
}

Every error string above is verified at its cited .rodata address, and the decompiled C shows the exact branch structure: the scalar guard fires the unique_indices=True string, the missing-vector guard fires "should have vector offset," codegenAccess reads the vector-offset Access at bap+72 (0x48), and the tail clears bap+0x58 after releasing bap+0x80.

GOTCHA — the two error strings are inverted guards on the same gather contract. The scalar flag present here is an error ("unique_indices=True … vector_offset is set"); the vector flag absent is also an error ("should have vector offset"). Together they assert: a unique-indices nisa.dma_copy gather must arrive with the vector offset set and the scalar offset clear. A reimplementer must not treat extractDynamicApToStandAlone as a general dynamic-AP extractor — it is the gather-index hoist specifically, and it rejects the scalar/register path.

Relation to assembleDynamicInfo

The two are siblings on the dynamic-DMA path, not nested (CONFIRMED by their disjoint bodies and shared codegenAccess leaf):

• assembleDynamicInfo packages the runtime offset into a symbolic companion (DynamicAPINFO, attached to the data AP).
• extractDynamicApToStandAlone hoists the index tensor out as a separate operand and strips the vector flag from the parent.

A unique-indices gather DMA uses extract… to materialize the index operand, and the data side uses assemble… (or the symbolic AP) for its address algebra. extract… does not build a DynamicAPINFO and does not emit register-ALU ops.

5. The kind-1 / kind-2 / kind-3 lifecycle

assembleDynamicInfo mints the symbolic (kind-2) form. The register-computed address is produced downstream. The full lifecycle, KLR codegen → lowering → encoder:

So the division of labor is sharp: assembleDynamicInfo (codegen) mints the symbolic offset_expr; LowerAP::convertSymAP (L32) materializes it into register-ALU ops; the encoder stamps the kind-3 register-mode descriptor. extractDynamicApToStandAlone runs alongside assembleDynamicInfo at codegen, hoisting the index so the symbolic AP it feeds references a clean standalone operand. extract… is therefore a codegen-time pre-materialization of the operand, not of the address compute.

6. The BIR-JSON wire form

The DynamicAPINFO is not a per-op top-level JSON key. It is nested in the access-pattern object under the common-header ins / indirection_ins arrays.

CORRECTION — an earlier draft attributed the serialization to bir::DynamicAPINFO::toJson @0x268c00. There is no such symbol in libwalrus.so: the DynamicAPINFO class exports only getCanonicalPattern, setActualPattern, isPartitionContiguous, setStaticOffsetPortion, and getNumElementsPerPartition (the full method set in _names.json) — no member named toJson. The address 0x268c00 (body frame 0x1268c00) is in fact neuronxcc::backend::CoreV2GenImpl::generateIndirectLoadSave(bir::InstDMA&, bool), an unrelated CoreV2 encoder, not a serializer. The field→key mapping below is recovered from the DynamicAPINFO member layout (the offsets each setter writes), not from a named toJson body, and the wrapping serializer was not pinned to a specific symbol.

The field→key mapping (from the member layout):

"actual_ap"                  ← DynamicAPINFO+0x00   (setActualPattern; the static AP)
"aff_expr" + "coef"          ← DynamicAPINFO+0x50   (the pair<QuasiAffineExpr,long> vector)
"offset_expr"  (QuasiAffineExpr) ← the per-axis runtime expr (kind 3 / kind 7)
"c"                          ← DynamicAPINFO+0x68   (setStaticOffsetPortion)
"indirect_dim"               ← DynamicAPINFO+0x74   (= 0 here; the W axis)
"indirect_dim_max_index"     ← DynamicAPINFO+0x70   (gatherable-row count)
"num_tensor_indirect_indices"← +0xF4 ; "tensor_indirect_arg_id" ← +0xF0  (set downstream)

The register half (reg_ap_offset / const_ap_offset / is_regloc_offset) appears later, on the kind-3 RegisterAccessPattern that lower_ap mints — not in assembleDynamicInfo. The regref key is the BirIntRuntimeValue's backing register.

7. Where the two flavours fire

The scalar/vector split maps onto two distinct front-end constructs:

• Scalar (kind 7, BirIntRuntimeValue) — a runtime scalar offset: a dynamic loop bound or a per-iteration base index read from a device register (the nl.* dynamic-slice / dynamic-shape access). The address is c + coef·reg; the register is resolved by name (getRegisterByName) and lowered to register-ALU ops by lower_ap. This is the dynamic-shape / data-dependent addressing path.
• Vector (kind 3, IndirectArgExpr) — a runtime index tensor: the classic gather/scatter. The index addresses the W/partition axis (indirect_dim forced 0); indirect_dim_max_index bounds the gatherable rows. This is the MoE expert-routing path — the per-token expert id (a vector of indices) selects which expert's weight rows to gather, exactly the nisa.dma_copy "unique_indices" gather that extractDynamicApToStandAlone's "Check the nisa.dma_copy call" guard names.

The codegen leaves that produce these are codegenNcDmaCopy (dynamic case calls assembleDynamicInfo), codegenNcLocalGather / codegenNcNGather (InstIndirectCopy / InstGather), and the GenericIndirectLoad/Save family. See front/pipeline for the DMA-bucketing front-end and Dynamic-Shape Synthesis for the whole thread.

8. Confidence and what is not pinned

The hard verification ceiling: every offset, immediate, vtable VA, and error string on this page is CONFIRMED against the cp310 libwalrus.so binary. The two STRONG items (the unread BAP slots, the precise shared_ptr-ctrl field cleared on hoist) and the single INFERRED item (the front-end use-site attribution) are the only places a reader should keep a reservation.