TensorCopyDynamic Generators

All addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22, cp310 wheel (the canonical "split_full" IDA target). cp311/cp312 are python_abi_equivalent duplicates covered by cp310; their offsets differ.

Abstract

A dynamic tensor copy is the penguin-IR node that moves data through an address that is only known at runtime — the building block behind XLA DynamicSlice (indexed read) and DynamicUpdateSlice (indexed write). Where an ordinary TensorCopy carries a fully compile-time access pattern (AP), a dynamic copy carries a runtime-address triple — a second tensor of index values (generic_addrs), the dimension those indices address (generic_dim), and the index→offset scale (offset_scale) — plus the AP free-indices that materialize the per-descriptor address register (addr_free_indices). At lowering this becomes an indirect DMA descriptor: a gather when the runtime address is on the read side, a scatter when it is on the write side.

Three Cython extension modules under neuronxcc/starfish/penguin/targets/generated/ implement this. TensorCopyDynamicBaseGen is the abstract parent emitted by brewer.py; it holds all the machinery — the constructor that pins both operands to AccessMode.LOAD, the operand tuple, the AP-index enumeration, and the BIR serializer. TensorCopyDynamicSrcGen (gather / read-dynamic) and TensorCopyDynamicDstGen (scatter / write-dynamic) are thin instabrew/main.py-emitted subclasses whose only behavioral divergence is one method: rhs_str, the IR-dump label. Everything else delegates upward. These are Gen classes in the brewer sense — generator classes wrapping an instruction definition, i.e. penguin-IR node producers, not codegen passes.

The page documents the base machinery first (the runtime-address payload, the LOAD-LOAD operand model, AP-index threading, serialization), then the two thin subclasses and the gather/scatter polarity they encode, then the lowering chain into the dynamic-DMA backend.

For reimplementation, the contract is:

• The runtime-address triple (generic_addrs, generic_dim, offset_scale) plus addr_free_indices, and which of those become operands vs. AP indices vs. serialized scalars.
• The LOAD/LOAD operand model: both the data tensor (src) and the index tensor (generic_addrs) are read inputs; the dynamic destination of a scatter is the instruction's result, never an extra operand.
• The src⇄gather / dst⇄scatter distinction and how it is encoded (class identity + rhs_str label + the brewer instruction definition), and how it maps to indirect-read vs. indirect-write DMA descriptors.

TensorCopyDynamicBaseGen — the base machinery

Purpose

TensorCopyDynamicBaseGen is the abstract parent that carries the entire dynamic-copy payload and behavior. It is the only one of the three modules with the full attribute set (src, generic_addrs, generic_dim, offset_scale, addr_free_indices), the AP-index machinery, and the serializer. The base-class string pool confirms its lineage and imports: it embeds the marker Generated by brewer from the definition in neuronxcc/instabrew/brewer.py at line 3384, the source path …/generated/TensorCopyDynamicBaseGen.py, the base instruction label NeuronUnaryInst, and the penguin-IR names …ir.Tensor, …ir.Access, AccessMode, set_access_mode, …targets.tonga.APIndex/TongaInst.

NOTE — "module name == class name" is the brewer convention. The qualified class is neuronxcc.starfish.penguin.targets.generated.TensorCopyDynamicBaseGen.TensorCopyDynamicBaseGen; likewise for *SrcGen/*DstGen.

The runtime-address payload

The five constructor parameters partition into three roles. Knowing the partition is the whole point of the data model:

All five names are confirmed against the binary: the string pool carries generic_addrs/_generic_addrs, generic_dim/_generic_dim, offset_scale/_offset_scale, addr_free_indices/_addr_free_indices, and the __init__ keyword-argument table (__pyx_pyargnames) lays them out in exactly this order after self/src.

QUIRK — the runtime-ness lives entirely in (generic_addrs, generic_dim, offset_scale, addr_free_indices). Drop them and you have a static TensorCopy. generic_addrs == None is a legal state (the serializer emits a none placeholder, §serialize) — i.e. the node can degrade to a static-address fallback with no second operand.

Algorithm — __init__

// __pyx_pw_…BaseGen_1__init__  @ 0x15930
// signature: (self, src, generic_dim, offset_scale,
//             generic_addrs=None, addr_free_indices=None)
function BaseGen_init(self, src, generic_dim, offset_scale,
                      generic_addrs=None, addr_free_indices=None):
    self.src = src
    self.src.set_access_mode(AccessMode.LOAD)            // 0x15930+~445/489: src is READ
    self.generic_addrs = generic_addrs
    self.generic_addrs.set_access_mode(AccessMode.LOAD)  // 0x15930+~567/596: index tensor is READ
    super().__init__(self.src, ...)                      // NeuronUnaryInst ctor — links src as the data input
    self.generic_dim   = generic_dim
    self.offset_scale  = offset_scale
    self.addr_free_indices = list(addr_free_indices)     // PySequence_List(...) if truthy
                            if addr_free_indices else []  // else PyList_New(0)
    self._link_operands()                                // wire operand→use edges into the def-use graph

Two distinct set_access_mode(AccessMode.LOAD) call sites appear in the decompiled body (one per operand), reached through the set_access_mode global and the AccessMode module global resolving .LOAD. The list(addr_free_indices) if … else [] idiom is visible as the PySequence_List branch vs. an empty-list branch. _link_operands is the zero-arg FastCall that registers the operand→use edges.

GOTCHA — for the base class neither operand is a STORE. The dynamic destination of a scatter (DstGen) is not modeled as a STORE operand here — it is the instruction's result tensor, materialized by the runtime address. A reimplementation that adds a third "dst" operand to operands() is wrong (see Dst operands below: it appends nothing).

Algorithm — operand model & queries

// operands  @ 0xff80
function operands(self):
    ops = [self.src]                 // operand 0 — always
    if self.generic_addrs:           // operand 1 — only when an index tensor is present
        ops.append(self.generic_addrs)
    return tuple(ops)                // (src) or (src, generic_addrs)

// loadTensor  @ 0x14c10  — "does this instruction READ tensor t?"
function loadTensor(self, t):
    if isinstance(self.src, Access) and self.src.tensor == t:                   return True
    if isinstance(self.generic_addrs, Access) and self.generic_addrs.tensor == t: return True
    return super().loadTensor(t)     // delegate to NeuronUnaryInst

loadTensor confirms the dual-LOAD model from the consumer side: the instruction reads both the data tensor and the index tensor. Access is the wrapper around ir.Tensor; .tensor unwraps it.

Algorithm — AP-index enumeration

enumerate_ap_indices is a Python generator (Cython splits it into a public wrapper @ 0xf9c0 and an inner __pyx_gb_…_12generator body @ 0x10e30; the inner body references the globals has_ap_indices, src_par_indices, addr_free_indices, free_indices, fixing the yield order):

// generator body  @ 0x10e30
function enumerate_ap_indices(self):
    if self.has_ap_indices:
        yield from self.src_par_indices    // partition indices of the src AP (only when present)
    yield from self.addr_free_indices      // the DYNAMIC-ADDRESS free indices
    yield from self.free_indices           // the remaining (base) free indices

QUIRK — addr_free_indices is yielded in line with the ordinary partition/free indices — the runtime-address indices are first-class AP indices. Whatever pass iterates enumerate_ap_indices (index renumbering, remove_ap_index, and the symbolic-AP register materialization) walks the dynamic-address registers through the same path as static ones. There is no separate "dynamic index" iterator; uniformity is the design.

The standard brewer instr-API mutators all touch this unified index list: remove_ap_index @ 0x12660 (raises "<idx> not found for remove_ap_index." — the error string is in the pool), plus the brewer-standard updateAPIndicies / _updateAllIndicesList / replaceUseOfWith operand-rewrite helpers. The addr_free_indices are spliced into the same renumber/rewire path, so a downstream register-materialization pass renumbers them uniformly.

Algorithm — serialize (BIR emission)

// serialize  @ 0x1a0c0
function serialize(self, ctx):
    ctx.var(self.src)                                 // emit the data operand var
    if self.generic_addrs:
        ctx.var(self.generic_addrs)                   // emit the index-tensor var
    else:
        ctx.none()                                    // else a 'none' placeholder
    serialize_ap_indices(ctx, self.addr_free_indices) // emit the dynamic-address AP indices
    kwargs = { generic_dim:  self.generic_dim,        // via ctx.number(...)
               offset_scale: self.offset_scale,
               shape/dtype:  ctx... + self.src.dtype } // element geometry carried through
    return super().serialize(ctx, **kwargs)

The decompiled body resolves ctx.{var,none,number}, self.{src,generic_addrs,generic_dim,offset_scale,shape}, the module global serialize_ap_indices, and src.dtype, then assembles a PyDict of kwargs and tail-calls super().serialize(ctx, **kwargs). The serialized BIR record is exactly the dynamic-copy descriptor handed to DMA lowering: data var, index var (or none), the dynamic-address AP indices, generic_dim, offset_scale, and shape/dtype.

Function Map

TensorCopyDynamicSrcGen — gather (read-dynamic)

Purpose

TensorCopyDynamicSrcGen is the gather form: it READS from a source whose address is a runtime value. generic_addrs + generic_dim + offset_scale supply the runtime read offset — for each runtime index i, src[generic_addrs[i] * offset_scale] along generic_dim is read and copied to a static destination. This is the IR for DynamicSlice / indexed read; it lowers to an indirect-read DMA descriptor.

It is a thin subclass. The string pool gives its brewer origin (…instabrew/main.py at line 760), its parent import TensorCopyDynamicBase (the hand-written wrapper, §hierarchy), and the same TongaISAInst / support.LogContext / NeuronEngine imports as Dst. It declares only six user methods, five of which delegate straight up.

Algorithm — the methods

// SrcGen — all five delegate; only rhs_str diverges
function __init__(self, *args, **kwargs):   return super().__init__(*args, **kwargs)   // @0xbca0
function operands(self):                    return tuple(list(super().operands()))     // @0xb4b0
function serialize(self, ctx):              return super().serialize(ctx)              // @0xce80 (empty kwargs)
function verify(self):                      return super().verify()                    // @0xc590
function replaceUseOfWith(self, src, dst):  return super().replaceUseOfWith(src, dst)  // @0xa5d0

// rhs_str  @ 0xda90  — the SRC-specific override (richer than Dst)
function rhs_str(self):
    return ("dynamic_copy_src %s generic_dim:[%s] generic_addrs: %s"
            % (self.src.use_str(), self.generic_dim, self.generic_addrs.use_str()))

The rhs_str reconstruction is anchored by three confirmed string fragments in the Src pool — "dynamic_copy_src ", " generic_dim:[", "] generic_addrs: " — joined as a 6-segment unicode format. The gather rhs explicitly prints the runtime read addressing (generic_dim + generic_addrs.use_str()), because for a gather the runtime address lives on the READ side and so belongs in the right-hand side of the IR dump.

TensorCopyDynamicDstGen — scatter (write-dynamic)

Purpose

TensorCopyDynamicDstGen is the scatter form: it WRITES into a destination whose address is a runtime value. The data is src (a normal LOAD); the runtime destination address comes from generic_addrs + generic_dim + offset_scale — for each runtime index i, the element of src is stored to dst[generic_addrs[i] * offset_scale] along generic_dim. This is the IR for DynamicUpdateSlice / indexed write / scatter; it lowers to an indirect-write DMA descriptor.

Brewer origin: …instabrew/main.py at line 787. Same imports and same thin six-method shape as Src.

Algorithm — the methods

// DstGen — five delegate; rhs_str is the only divergence
function __init__(self, *args, **kwargs):   return super().__init__(*args, **kwargs)   // @0xbb40
function operands(self):                    return tuple(list(super().operands()))     // @0xb350 — NO extra dst operand
function serialize(self, ctx):              return super().serialize(ctx)              // @0xd5b0
function verify(self):                      return super().verify()                    // @0xccc0
function replaceUseOfWith(self, src, dst):  return super().replaceUseOfWith(src, dst)  // @0xa470

// rhs_str  @ 0xc430  — the DST-specific override (terse)
function rhs_str(self):
    return "dynamic_copy_dst %s" % self.src.use_str()

QUIRK — the scatter rhs prints only the data source ("dynamic_copy_dst %s", the single confirmed format string). It does not re-print generic_dim / generic_addrs, in deliberate contrast to the gather rhs above. The runtime address of a scatter is the instruction's output / dest side, so it is implied by the result and omitted from the right-hand side. This asymmetry is the cleanest binary-visible signal of the src⇄gather / dst⇄scatter polarity.

NOTE — in replaceUseOfWith(self, src, dst) the parameter names src/dst are brewer's old-use / new-value operand-rewrite names — not the copy's data-src / addr-dst. The rewrite delegates to the base def-use machinery.

Class hierarchy and the src/dst encoding

TongaInst / TongaISAInst / NeuronUnaryInst         (penguin/tonga base instruction classes)
        ▲
        │  brewer.py L3384 — the abstract instr definition; ALL the machinery
 TensorCopyDynamicBaseGen   [targets/generated/ — this task]
        ▲
        │  hand-written wrapper (device/lowering hooks); imported by Src/Dst
        │  as `TensorCopyDynamicBase` — NOT a standalone .so in targets/generated
 TensorCopyDynamicBase
        ▲                          ▲
        │                          │
 TensorCopyDynamicSrcGen      TensorCopyDynamicDstGen
  (gather / read-dynamic)      (scatter / write-dynamic)
  main.py L760                 main.py L787
  rhs_str "dynamic_copy_src…"  rhs_str "dynamic_copy_dst %s"

The src-vs-dst distinction is therefore not a runtime flag — it is encoded by (a) class identity, (b) the rhs_str label and what it prints, and (c) the brewer instruction definition the subclass re-declares (engine hint, instance class, the access-mode polarity of the copy target). TensorCopyDynamicBase — the non-Gen wrapper Src/Dst actually inherit from — is hand-written and lives outside these three .so; its exact body is INFERRED here.

NOTE — NeuronEngine (STRONG). Both Src and Dst import NeuronEngine at module level, but it is referenced in no method body — it is consumed in the brewer-emitted class body as the per-instruction engine hint (e.g. a default-engine or engine-eligibility attribute). This is the attribute the DGE-level / engine-selection logic reads when assigning the dynamic copy to a DMA engine. Both variants carry it; the gather/scatter polarity plus this engine attribute drive engine assignment. The exact member value is a class-body constant, not in any traced method (INFERRED).

Lowering — dynamic IR node → indirect DMA

The producer chain (STRONG, cross-referenced)

front-end dynamic-shape op (DynamicSlice / DynamicUpdateSlice)
  → TensorCopyDynamicSrcGen / TensorCopyDynamicDstGen   (THIS task: the IR-node producer)
  → penguin IR NeuronUnaryInst carrying the dynamic-AP payload
  → symbolic-AP register materialization (operates on addr_free_indices + generic_addrs)
  → dynamic-AP lowering selects the indirect AP kind
  → dynamic-DMA / DGE path materializes the indirect descriptor
  → engine assignment per the NeuronEngine class attribute

The on-wire descriptor from BaseGen.serialize carries exactly the fields the DMA lowering needs:

src ⇄ gather, dst ⇄ scatter (STRONG)

• TensorCopyDynamicSrcGen → the read address is runtime → indirect / gather DMA. The descriptor's source address is generic_addrs[i] * offset_scale along generic_dim. (This is why Src's rhs_str prints generic_addrs — the runtime addr is on the read side.)
• TensorCopyDynamicDstGen → the write address is runtime → indirect / scatter DMA. The descriptor's dest address is generic_addrs[i] * offset_scale along generic_dim. (This is why Dst's rhs_str does not print it — the runtime addr is on the output side.)

The two map to the gather vs. scatter indirect-DMA descriptor variants the backend materializes. The rhs_str asymmetry is directly binary-confirmed; the gather/scatter ⇔ src/dst mapping itself is STRONG (an inference from which side carries the runtime address, fully consistent with every observed string and operand fact).

Confidence ladder

• CONFIRMED — full Cython bodies for BaseGen __init__/operands/loadTensor/verify/serialize + the enumerate_ap_indices generator; both subclasses' __init__/operands/serialize/verify/rhs_str/replaceUseOfWith; all field names (generic_addrs/generic_dim/offset_scale/addr_free_indices); the dual AccessMode.LOAD on src + generic_addrs; the brewer source markers; the rhs_str label strings.
• STRONG — gather/scatter ⇔ src/dst-dynamic mapping; NeuronEngine as the engine-selection hint; generic_addrs/generic_dim/offset_scale as the runtime-address triple; addr_free_indices as the symbolic-AP register set; the producer→DMA-lowering chain.
• INFERRED — the exact NeuronEngine member value (class-body constant, in no method body); the hand-written body of the TensorCopyDynamicBase wrapper (outside these three .so); the precise super().__init__ argument splat beyond self.src (Cython obscures the *args).

Cross-References

• DGE-Level Dynamic DMA — the DGE levels and engine selection that consume the NeuronEngine hint on these nodes
• Symbolic-AP Register Materialization — lowers addr_free_indices + generic_addrs into address registers
• Indirect Descriptors — the gather/scatter indirect-DMA descriptor format these nodes lower to
• Penguin IR Node Model — the def-use graph _link_operands wires these operands into
• Penguin Target Abstraction — the Tonga / NeuronEngine target layer the brewer instrs bind to