Penguin Tensor / Buffer Node & Memory-Space Placement

All symbols and offsets on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310; cp311/cp312 are byte-identical Cython rebuilds). The abstract tensor lives in neuronxcc/starfish/penguin/ir/Tensor.cpython-310…so; the physical placement lives in neuronxcc/starfish/penguin/targets/tonga/TongaTensor.cpython-310…so. Both ship with debug info and are not stripped, so every method name below is a verbatim __pyx_pw_…/__pyx_k_… symbol. Treat every offset as version-pinned.

Abstract

A Penguin tensor is an SSA value with a shape, a dtype, and — eventually — a place to live. The middle-end splits that into two strata that a reimplementer must keep distinct. The abstract tensor (ir/Tensor.py, a ComputeValue subclass with 119 methods) carries shape, dtype, name, source-vs-device layout, and an IO role — but no physical address. The target-bound tensor (targets/tonga/TongaTensor.py) is the same SSA value re-classed into one of a small set of physical subtypes, each of which names a hardware memory and carries a BufferAllocation describing where in that memory the data sits. Placement in Penguin is therefore not a flag on a tensor — it is a change of class: the tiler (InferTongaTensor, [Part 5 layout passes]) lowers an ir/Tensor into a NeuronSBTensor, NeuronPSUMTensor, or a DRAM{2D,3D}BlockTensor, and that choice of subclass is the memory-space declaration.

This mirrors LLVM only loosely. In LLVM a pointer carries an addrspace(n) integer and the value type is orthogonal; here the address space is reified as the Python type of the node, and the byte address is a quasi-affine expression over the loop block axes rather than a constant. The two on-chip memories the subclasses target — SBUF (the 128-partition state buffer) and PSUM (the small array of fixed banks the PE array drains into) — are physically different address spaces with different allocation tuples, and the BufferAllocation serialize format [{bnk_idx}, {block_addr}]… exposes that bank/partition structure directly. The abstract tensor's only nod to placement is the NonLocal bit in its TensorType kind-flags, which marks a value the allocator is allowed to spill off-chip to DRAM.

This page recovers: the abstract Tensor serialization and field model; the TensorType kind bit-flags and the composite predicate masks built over them; the TongaTensor subclasses (DRAM{2D,3D}BlockTensor, NeuronLocalTensor, NeuronSBTensor, NeuronPSUMTensor, NeuronBlockTensor) and which memory each names; and the BufferAllocation family that holds the physical bank/partition address. It is a sibling of 5.1 — the Penguin node model, which covers the SSA Value/Instruction/Operator base that Tensor descends from.

For reimplementation, the contract is:

• The abstract Tensor serialize format and its five core fields (attr, dtype, shape, name, init), plus the TensorType kind enum and the seven composite flag-masks that gate every IO query.
• The two-strata model: abstract ir/Tensor has no address; physical placement is realized by re-classing into a TongaTensor subclass — NeuronSBTensor (SBUF), NeuronPSUMTensor (PSUM), DRAM{2D,3D}BlockTensor (off-chip), with NeuronLocalTensor the common on-chip base and NeuronBlockTensor the partition/block-tiled form.
• The BufferAllocation physical-address struct: its (bank_id, block_addr) fields, its [{bnk_idx}, {block_addr}]{alloc_bnk_shape}{alloc_blk_shape} serialization, and the PSUM 3-tuple vs SBUF 2-tuple distinction.
• The address-as-quasi-affine-expression model: the bank id and byte offset are functions of the block index, not constants.

The abstract tensor — ir/Tensor

Purpose

Tensor is the tensor-valued SSA def: a ComputeValue (and through it a Value/User, 5.1) that names a multi-dimensional buffer the program reads and writes. It is the operand and result type of every Operator and ISA Instruction. At this level the tensor is abstract: it knows its shape, dtype, name, and the layout it wants, but it has no bank, no partition, and no byte address. Everything physical is added later by re-classing into a TongaTensor subtype (§ The target-bound tensors).

Serialization and core fields

The textual form is the cleanest window onto the node's identity. The serialize template is interned verbatim:

"{attr}{dtype} {shape} {name}{init}"

So an abstract tensor is exactly five textual fields: an attribute prefix (attr — the IO/placement annotations, e.g. an Input/Const marker), a dtype, a shape, a name (the SSA name string, which is also the dep-edge hash key on the BIR side), and an optional initializer (constant data). The {name} and {init} template slots are interned only inside that format string, not as standalone field strings; the backing storage fields that are independently interned are: shape (@ 0x9cfc8, with orig_shape / new_shape / source_shape / access_shape variants), dtype (@ 0x9cff2, + dtype_size_in_bytes), layout (@ 0x9cf33, + src_layout / dst_layout / deviceLayout / deviceShape), and addrs (@ 0x9d00b).

NOTE — addrs (@ 0x9d00b) exists on the abstract tensor, but it holds address expressions (quasi-affine AffineExpr over loop axes — [5.x affine algebra]), not a resolved bank/byte. The resolved physical address is only present after a BufferAllocation is attached to a TongaTensor subclass. Do not read addrs as "this tensor is placed."

The layout model is two-faced: a src_layout/source_shape (the layout the HLO producer emitted) and a deviceLayout/deviceShape (the on-device layout the backend wants), bridged by a transpose view (transposeView / sourceTranspose / applyStaticTranspose). This is why a single tensor can be "logically NCHW, physically partition-major" without a copy: the view carries the permutation. (Full layout-solver treatment is out of scope here; see [Part 5 layout].)

The TensorType kind bit-flags

A tensor's IO role is a TensorType value — a bit-flag set, not a scalar enum. The five primitive members are interned verbatim in the ir/Tensor pool:

QUIRK — NonLocal is a placement permission, not a placement. It does not say "this tensor is in DRAM"; it says "the allocator is allowed to put it there." A tensor can be Output | NonLocal (a graph output the allocator spilled) or a purely on-chip Output. The actual memory is decided by which TongaTensor subclass the tiler picks (§ next), not by this bit. set_optionally_non_local / unset_optionally_non_local are the tiler's hooks to toggle the option without committing.

Each member occupies one bit position; membership is tested by TensorType.contains (interned verbatim as TensorType.contains @ 0x9b8f0, and fully-qualified as neuronxcc.starfish.penguin.ir.Tensor.TensorType.contains @ 0x944b0) — i.e. a bit-mask AND. To avoid recomputing common masks, the pool also interns seven composite flag-sets, all confirmed verbatim:

InputOrConst                       = Input | Const
InputOrNonLocal                    = Input | NonLocal
InputOrOutput                      = Input | Output
InputOrOutputOrConst               = Input | Output | Const
InputOrConstOrNonlocal             = Input | Const | NonLocal
InputOrOutputOrNonLocal            = Input | Output | NonLocal
InputOrOutputOrConstOrNonLocal     = Input | Output | Const | NonLocal

Each composite is the precomputed mask behind a predicate method. The query API on Tensor is interned verbatim:

GOTCHA — isInputAndOutput is conjunction; the rest are disjunction. Every …Or… method is a "does the flag set intersect this mask" test. isInputAndOutput is the only one that asks "are both bits set" — an aliased in-place IO tensor (donated parameter). A reimplementer who codes them all as (kind & mask) != 0 will mis-classify donated buffers.

The address-space tag — VNCAddrSpace

Alongside the kind-flags the tensor carries a VNCAddrSpace value (interned in ir/Tensor) — the address-space tag (the "VNC" = virtual Neuron core address space). This is the coarse space label (on-chip vs DRAM vs the collective/replica address window) that travels with the tensor independent of the fine-grained BufferAllocation. It is the nearest analog to an LLVM addrspace(n): a small enum the lowering threads down to BIR, distinct from the per-subclass physical class.

Splat / padded specializations

Two abstract specializations matter for placement:

• SingleValueTensor (7 methods: value, splat_value, is_const, convertToTensor, _reshape_value, serialize, verify_value) — a broadcast-scalar / splat: one value over a shape. It needs no per-element storage; the lowering can materialize it as a constant-splat rather than a real buffer. The related BroadcastScalar (ir umbrella) is the broadcast-of-scalar value.
• PaddedTensor / ZeroPaddedTensor (ir/PaddedTensor.py) — a tensor with implicit edge padding (down_paddings / up paddings), the pad-op result view. The padding is logical; the underlying buffer is still placed by the normal mechanism.

The target-bound tensors — tonga/TongaTensor

Purpose

This module realizes the physical memory model. After tiling, each abstract ir/Tensor is lowered to one TongaTensor subclass, and that choice declares the memory space. There is no memory_space= field; the type is the field. The module string table (@ 0x9193) and module-exec body (@ 0x168d2) intern all five subclass names verbatim, with their docstrings and method rosters.

The placement taxonomy

ComputeValue (ir/Tensor)                 — abstract, no address
   └─ NeuronTensor                       — "Tonga specific tensor type"   (docstring confirmed)
        ├─ DRAM2DBlockTensor             — off-chip (HBM/DRAM) buffer, 2-D block
        ├─ DRAM3DBlockTensor             — off-chip (HBM/DRAM) buffer, 3-D block
        └─ NeuronLocalTensor             — "Tonga specific tensor type on
             │                              statebuffer or psumbuffer"     (on-chip base)
             ├─ NeuronSBTensor           — SBUF  (state buffer)   engine tag "TongaSB"
             ├─ NeuronPSUMTensor         — PSUM  (matmul accumulator) tag "TongaPSum"
             └─ NeuronBlockTensor        — partition/block-tiled on-chip tensor
        └─ NeuronWeightTensor            — PE-array weight stream (set_split/split_axis)
             └─ IdentityWeightTensor     — the 128×128 identity (matmul-transpose)

The off-chip space is realized by two block-shaped DRAM tensor classes: the binary interns DRAM2DBlockTensor (@ 0xc7550) and DRAM3DBlockTensor (@ 0xc7530), with partition-diag strings "DRAM2DBlk partitions[%s]" / "DRAM3DBlk partitions[%s]".

GOTCHA — there is no standalone DRAM class to look for; every off-chip tensor is one of the two block-shaped types above.

The memory-space mapping is the central fact a reimplementer needs:

QUIRK — SBUF and PSUM are different address spaces, hence different classes. A reimplementer coming from a flat-memory IR will want one OnChipTensor with a space field. The binary refuses: NeuronSBTensor and NeuronPSUMTensor are siblings, because their allocation tuples differ in arity (SBUF is (start_partition, base_addr); PSUM is (bank_id, start_partition, base_addr) — the PSUM bank is a third coordinate SBUF does not have). The common behavior they do share (allocation lifecycle, has_physical_address, address linking) lives one level up in NeuronLocalTensor.

NOTE — the partition dim splits a NeuronLocalTensor's dims into time and space. The NeuronLocalTensor docstring (@ 0xc4ca0) states it verbatim: dims higher than partition_dim map to time (the schedule/loop order), and dims lower than or equal to partition_dim map to space (the physical 128-partition × byte layout). So the partition axis is the boundary between "iterated over in the loop nest" and "laid out across the hardware" — a reimplementer must place exactly the dims at/below partition_dim into the on-chip address, and leave the higher dims as loop iterations.

NeuronLocalTensor — the on-chip allocation API

NeuronLocalTensor (24 methods) is where on-chip placement actually happens. Its allocation lifecycle (all interned verbatim):

// NeuronLocalTensor — on-chip (SBUF|PSUM) allocation lifecycle
createAllocation(...)           // bind a concrete BufferAllocation
createSymbolicAllocation(...)   // bind a SymbolicBufferAllocation (address still a free var)
resetAllocation() / cleanAllocation() / cleanAllocationModule()
has_physical_address()          // true once a non-symbolic allocation is bound
_link_address() / _unlink_address() / cleanAddress()   // wire addr into the def-use graph
addr_str()                      // textual address
linearizePartitionAddrs(...)    // flatten the 2-D (partition × free) addr to 1-D
// size queries answered from the bound allocation:
alloc_size_in_bytes / alloc_bank_size_in_bytes / alloc_block_size_in_bytes
num_alloc_banks / num_alloc_blocks / num_alloc_element_per_bank

GOTCHA — a tensor can be placed in space but not yet addressed. createSymbolicAllocation binds a SymbolicBufferAllocation whose bank/byte are free variables; has_physical_address returns false until a concrete BufferAllocation replaces it. The two-phase split (choose the space by re-classing, then resolve the address later) is deliberate: the layout solver fixes the class early and the bank/byte allocator (Part 8) resolves the symbol late. Treat "has a NeuronSBTensor" and "has an SBUF address" as independent facts.

NeuronBlockTensor — the partition/block-tiled form

After tiling, an on-chip tensor is usually a NeuronBlockTensor: the 128-partition axis and the block-tiling outer axis are made explicit. Confirmed methods (33 total) split into three groups:

// geometry — the explicit partition × block shape
npartitions / nblocks
partition_size / partition_size_in_bytes / total_partition_size_in_bytes
partition_start / partition_ap / partitionDimIndices
block_shape / tile_size
tonga_rank / tonga_shape                 // the device-side (partition,block,free) rank/shape
calculate_partition_size(...)

// dim motion — move a logical dim onto the partition / block / free axis
moveToBlkDim / moveToFreeDim / moveToOuterFreeDim
expandPartitionDim / canDelinearizeParitionDim   // [sic — "Parition" is the real symbol]
linearizePartitionAddrs

// sharding (LNC)
access_by_lnc_comm / used_by_shard_reduce

The tonga_shape is the device view: (partition, [block…], [free…]), where the partition dim is the 128-wide SBUF/PE axis. moveToBlkDim / moveToFreeDim are how the layout solver re-assigns which logical dimension is the partition, the block, or a free dim — the P/F/B classification consumed downstream.

NOTE — the misspelling canDelinearizeParitionDim is the real interned symbol (Parition, missing a t). Cross-check against the binary, not against a corrected spelling — a grep for Partition will miss it.

Layout-requirement hints

NeuronLayoutDimHint (35 methods) / NeuronShardDimHint are per-dimension layout requirements the layout solver reads to decide P/F/B. Each dim can demand must_be_partition, must_be_free, or must_be_block (all interned verbatim), plus per-op hints (matmulLoadHint, transposeLoadHint/StoreHint, genericLoadHint/StoreHint, …). These are inputs to placement, not placement itself.

The physical address — BufferAllocation

Purpose

A BufferAllocation is the struct a NeuronLocalTensor holds once placed: it answers which bank, which partition, which byte, how big. Three variants are interned verbatim:

Fields and serialization

The base BufferAllocation interns these field/accessor names verbatim:

struct BufferAllocation {            // fields are Python __dict__ entries; names interned
    bank_id;                         // which PSUM bank (PSUM only)
    block_addr;                      // byte address within the bank/partition
    // size accessors (computed):
    alloc_size_in_bytes;
    alloc_bank_size_in_bytes;
    num_alloc_banks;
    num_alloc_blocks;
    allocated_bank_upper_bound;      // highest bank index this alloc touches
    allocated_byte_addr_upper_bound; // highest byte address this alloc touches
};
// methods: __init__, __str__/serialize, expandPartition,
//          enumerate_indices, is_static_allocation

The serialize template is interned verbatim (@ 0xfd32):

"[{bnk_idx}, {block_addr}]{alloc_bnk_shape}{alloc_blk_shape}"

So a serialized allocation reads as [bank, byte_addr] followed by the bank shape and block shape — i.e. the coordinate is a (bank, byte) pair and the two trailing shapes give how many banks and blocks the allocation spans. The diagnostic string "bank_id on " and "NeuronPSUMTensor must be initialized to…" confirm that PSUM placement validates a bound bank.

PSUM vs SBUF allocation tuples

The arity of the allocation differs by target memory. The address-map docstring (@ 0xc3700) states both forms verbatim: "For PSUM, it's a three-element tuple (bank_id, start_partition, base_addr)" and "For SBUF, it's a two-element tuple (start_partition, base_addr)".

QUIRK — the bank id and byte offset can be quasi-affine functions of the block index, not constants. The allocation is not necessarily a fixed (bank, byte); for a block-tiled tensor each block can map to a different bank/offset via a quasi-affine expression over the block axes (the same affine algebra the access patterns use). enumerate_indices walks those per-block addresses. A reimplementer who models a BufferAllocation as a constant base+size will be correct for the simple case and wrong for any tensor whose tiles round-robin across banks.

Where placement is decided vs resolved

ir/Tensor (abstract, NonLocal-permission only)
   │
   │  InferTongaTensor pass (tonga/passes)   ── chooses the SUBCLASS = the SPACE
   ▼
NeuronSBTensor / NeuronPSUMTensor / NeuronBlockTensor / DRAM
   │
   │  createSymbolicAllocation                ── space fixed, address still a free var
   ▼
SymbolicBufferAllocation   ── has_physical_address() == False
   │
   │  bank/byte allocators (Part 8)           ── resolve the symbol to a concrete bank+byte
   ▼
BufferAllocation           ── has_physical_address() == True

The space is declared by the layout/InferTongaTensor stage (the subclass choice); the address is resolved later by the backend allocators. The Penguin tensor node is the carrier of both decisions but commits them at different times.

Evidence summary

The five structural claims on this page and what pins each:

1. Abstract serialize "{attr}{dtype} {shape} {name}{init}" — the literal is present verbatim in the ir/Tensor string table (0x8e1a).
2. TensorType = {Input,Output,Const,InternalIO,NonLocal} + 7 composite masks + membership predicate — all five primitives, all seven …Or… composites, and TensorType.contains (@ 0x9b8f0) are interned verbatim in ir/Tensor. The composite masks imply bit-packing, but no kind_shift symbol exists in either module, so this page names no packer.
3. BufferAllocation serialize [{bnk_idx}, {block_addr}]{alloc_bnk_shape}{alloc_blk_shape} — the format string is at 0xc5c40 (ida) / 0xfd32 (disasm table) in TongaTensor.so; both alloc_bnk_shape and alloc_blk_shape tokens are independently interned.
4. The subclasses and their docstrings — NeuronLocalTensor, NeuronSBTensor, NeuronPSUMTensor, NeuronBlockTensor, NeuronWeightTensor, IdentityWeightTensor are all interned, and the full NeuronLocalTensor docstring (@ 0xc4ca0), including the time/space dim-mapping, is binary-direct.
5. PSUM 3-tuple vs SBUF 2-tuple — the address-map docstring (@ 0xc3700) states both tuples verbatim ("…three-element tuple (bank_id, start_partition, base_addr)" / "…two-element tuple (start_partition, base_addr)"). The bank_id/block_addr/start_partition fields and the "NeuronPSUMTensor must be initialized to…" diagnostic are independently interned.

The __dict__-resident field names are interned; their owning class is [INFERRED] only where a method roster does not pin it. Three tokens a reader might expect are simply not present in these modules and are deliberately not claimed here: a standalone DRAM class, standalone name/addr/init field strings (they appear only inside the serialize template), and kind_shift.

Related Components

Cross-References

• Penguin Node Model — 5.1; the Value/ComputeValue/Instruction/Operator SSA base that Tensor descends from, and the def-use graph the loads/stores ride on.
• SBUF / PSUM Bank Geometry — 1.05; the physical partition count, per-partition byte budget, and 2 KiB PSUM bank size the BufferAllocation addresses index into.
• DRAM / HBM Geometry & the DRAM Split — 1.06; the off-chip window a DRAM tensor / NonLocal-spilled value lands in.
• Part 8 — The libwalrus Backend — the bank/byte allocators that resolve a SymbolicBufferAllocation into a concrete BufferAllocation.
• Part 7 — BIR & the Simulator — the BIR MemoryLocation the lowered tensor address becomes once Penguin descends to BIR.