NeuronCodegen Tensor-Op Forward Builders
All symbols and addresses on this page apply to
neuronx_cc2.24.5133.0+58f8de22, moduleneuronxcc/nki/compiler/backends/neuron/KernelBuilder.cpython-310-x86_64-linux-gnu.so(the cp311/cp312 twins are source-identical; their Op set matches). The.sois an UNSTRIPPED 14.5 MB Cython extension with DWARFdebug_info. Treat every address as version-pinned.
Abstract
This page documents the ~30 tensor-op forward builders of the NKI KernelBuilder module — the methods that take nl/nisa-Python arguments and construct a Penguin-IR <Op> node, then append it to the current basic block. This is the forward (build) direction. It is the exact inverse of the re-emit printer documented in 6.5.9, which walks an existing Penguin node back out to nisa.<prim>(...) text. The two surfaces marshal the same enums in opposite directions, and confusing them is the single largest hazard in this part of the compiler.
Each builder follows one fixed shape: parse kwargs, validate, normalize operand tiles through combine_tiles, read the tile's access pattern as (par, free) index pairs, construct one penguin.ir.<Op> Python object with named kwargs, and append it through self.insert — which reaches IRBuilder.add_named_instruction after stamping predicates, dependency edges, bookkeeping, and a source location. The op-selector enums are passed through verbatim: the builder stores the Python np.ufunc / ALUOpcode / TSOpcode / EngineAccumulationType object on the node as-is. Numeric ISA-enum renumbering happens one layer down at BirCodeGenLoop (strands I04/I05), not here.
The one-to-one and many-to-one builder→Op mapping is the payload of this page. Most builders own exactly one Op; a few interesting cases collapse multiple builders onto a single Op (tensorscalar + tensorvectorscalar → one TensorScalarPtrOp), decompose into several nodes (tensorscalar of np.power), or split on a runtime property (select → TensorSelect or AffSelTensorScalarOp). The MX-quantize builder quantize_mx → QuantizeMXOp lives here, correcting an earlier strand that placed it only in BirCodeGenLoop.
CORRECTION — the Cython class is
GeneratedNeuronCodegen, not bareNeuronCodegen. Sibling strand P01 (and theS2-09source survey) name the owning classNeuronCodegen. The symbol mangling in this binary is unambiguous: every method wrapper is__pyx_pw_9neuronxcc_9generated_3nki_8compiler_8backends_6neuron_13KernelBuilder_22GeneratedNeuronCodegen_<idx><name>. The22length-prefix decodes to the 22-character identifierGeneratedNeuronCodegen— the Cython-emitted concrete class that the methods actually compile into (a generated subclass;NeuronCodegenis the hand-written base). This page usesGeneratedNeuronCodegenwhen citing symbols andNeuronCodegenwhen speaking of the logical surface; they are the same builder. [CONFIRMED — symbol mangling,names.json]
| Module / class | KernelBuilder.cpython-310…so · class GeneratedNeuronCodegen (logical NeuronCodegen) |
| Direction | nl/nisa-Python args → construct penguin.ir.<Op> node → self.insert |
| Insert entry | self.insert (mdef idx 297, @0x165f00) → IRBuilder.add_named_instruction |
| Tile normalize | combine_tiles (@0x16ad90) + tile.canonical_par_indices / canonical_free_indices(...) |
| Per-op modules | neuronxcc.starfish.penguin.targets.generated.<Op> · ISA base …targets.tonga.TongaISAInst |
| Enum policy | pass-through; numeric ISA renumber deferred to BirCodeGenLoop codegen |
| Builder count | ~30 tensor-op methods (of 194 total __pyx_mdef_…GeneratedNeuronCodegen_<N><name>) |
| MX builder | quantize_mx (idx 161, @0x1f9390) → QuantizeMXOp — lives here, not only in BirCodeGenLoop |
NOTE — evidence base and its limits. Method addresses, mdef indices, and
KernelBuilder.pysource lines below are read from the binary'snmsymbol table (__pyx_pw/__pyx_mdefper method) and DWARFdecodedline, the Op-class names from the constructedpyx_n_s_<Op>body locals and.rodata. The IDA JSON export checked into this tree captured a tail window of the class (the high-index methodsrand2/rng/exponential/activate2/nonzero_with_count/*_read_accumulator, idx 1–17 — confirming the class name, the import modules, andEngineAccumulationType/combine_tiles/insertstrings), but does not individually re-confirm each tensor-op method body. Method-level facts are therefore tagged STRONG unless a string is directly present in this export, in which case they are CONFIRMED.
1. The builder pattern
Every tensor-op method is a six-step pipeline. The skeleton below is the shared shape; §2–§9 fill in the per-family specifics. Function names are the real symbols; the per-op kwargs are the confirmed body locals.
// GeneratedNeuronCodegen.<tensor_op>(self, ...kwargs...) // one method per builder
PyObject *build_tensor_op(self, /* op-specific kwargs */) {
// (a) parse kwargs: __Pyx_ParseOptionalKeywords / __Pyx_GetKwValue_FASTCALL
// (pyargnames SIMD-insert block gives the exact kwarg order)
parse_kwargs(&op, &operands, &dtype, &mask, &deps, &schedule, &name, ...);
// (b) validate operand dtypes / engine legality (per-family guard literals)
validate(operands, dtype); // e.g. check_tensor_int32_ops_supported
// (c) normalize / broadcast / align operand tiles
tiles = self.combine_tiles(operands); // @0x16ad90 (shared w/ P01 matmul, P03 memory)
// (d) read the access pattern off each tile as (stride,size) APPairs
par = tile.canonical_par_indices; // partition-axis index
free = tile.canonical_free_indices(...); // free-axis indices
// (e) construct ONE penguin.ir.<Op> node, storing enums VERBATIM
inst = penguin.ir.<Op>(op=op, in=tiles, /* flags, engine, dtype, name */);
// (f) append it; insert wraps add_named_instruction + predicates + deps + debugloc
self.insert(inst, buffer, name, deps, sema, mask, ...); // idx 297 @0x165f00
return out_tile;
}
self.insert — the append wrapper
self.insert (@0x165f00, shared verbatim with the matmul and memory families) does, in order [CONFIRMED — .rodata strings insert + cross-strand P03]:
self.insert(inst, ...) {
builder.add_named_instruction(inst); // penguin.ir.IRBuilder.add_named_instruction — the real emit
add_predicates(inst, mask); // attach the `mask` predicate list
process_dep_edges(inst, deps); // attach the `deps` dependency edges
process_new_neuroninst(inst); // NeuronCodegen bookkeeping
update_debugloc(inst); // stamp the source location
}
GOTCHA —
add_named_instructionis never called raw at a build site. The "IRBuilder.add_named_instruction → penguin.ir." pairing is the ctor-then-insert pair: the method constructs the node, then self.insertreachesadd_named_instructionfor it. A reimplementer who hunts for a directadd_named_instruction(...)at each builder will not find one — it is always behindself.insert. [CONFIRMED —add_named_instructionsymbol present;insertstring present in this export]
Enum marshalling is pass-through
The three op-selector enums are stored as Python objects on the Op node. The builder does not assign a numeric ISA value — that is deferred to BirCodeGenLoop.codegenAluOp / codegenAccumCmd (strands I04/I05). The .rodata type-hint literals pin the parameter types:
| Selector | kwarg(s) | Type hint (.rodata) | Carried on | Confidence |
|---|---|---|---|---|
| ALU op | op | Union[np.ufunc, ALUOpcode] | TensorTensorOp / TensorReduceOp / ActivationOp | STRONG (type-hint) |
| TensorScalar op | op0 / op1 | TSOpcode | TensorScalarPtrOp | STRONG (type-hint) |
| Accumulate cmd | reduce_cmd | Optional[EngineAccumulationType] (default Idle; reset member ResetReduce) | activation / select-reduce / tensor-scalar-cache | CONFIRMED — EngineAccumulationType string present @0x24870 |
| Activation func | op (of activation) | np.ufunc | ActivationOp | STRONG (type-hint) |
NOTE — this builder is the source the printer reads. The 6.5.9 printer
opcode()/reduce_cmd()name-mappers (sigmoid→expit, erf→erf, etc.) read back exactly thenp.ufunc/enum objects this builder stores verbatim. Builder = store; printer = name-map;BirCodeGenLoop= numeric ISA renumber (codegenAluOp1→29;codegenAccumCmdIdle/Zero/AddAccum/ZeroAccum). [STRONG cross-ref I04/I05]
2. The complete builder → penguin.ir.<Op> map
All ~30 tensor-op builders, grouped by family. mdef idx is the __pyx_mdef_…GeneratedNeuronCodegen_<N><name> index; pw addr is the nm __pyx_pw wrapper address; py is the KernelBuilder.py source line from DWARF; Op is the constructed pyx_n_s_<Op> body local. Method addresses/lines are STRONG (from the report's nm/DWARF, not re-confirmed by this partial IDA export); Op-class names are CONFIRMED at the named-local level.
| Family | idx | builder | pw addr | py | → penguin.ir.<Op> |
|---|---|---|---|---|---|
| Activation | 211 | activation | 0x1968c0 | 2882 | ActivationOp |
| 213 | activation_accu | 0x9f810 | 2960 | ActivationAccumulationOp | |
| Tensor-Tensor | 181 | tensortensor | 0x15b850 | 2262 | TensorTensorOp (cls=, default) |
| 179 | binop | 0x177c80 | 2221 | delegates → tensortensor | |
| 215 | tensortensorscan | 0x262c60 | 3010 | TensorTensorScanOp (2-op scan) | |
| Tensor-Scalar | 217 | tensorscalar | 0x140530 | 3066 | TensorScalarPtrOp (op0/scalar0/reverse0) |
| 203 | tensorvectorscalar | 0x293fd0 | 2733 | TensorScalarPtrOp (op0/op1 + operand0/1 + reverse0/1) | |
| 219 | tensorscalarcumulative | 0x1f59b0 | 3131 | TensorScalarCacheCumulative | |
| 221 | tensorscalarreduce | 0x1ca7f0 | 3184 | TensorScalarCacheReduce | |
| Reduce | 173 | tensorreduce | 0x239de0 | 2098 | TensorReduceOp (op/reduce_dims/negate/keepdims) |
| 175 | tensor_partition_reduce | 0x24b420 | 2162 | PartitionReduceOp (reduce_all_axis) | |
| Unary | 177 | unary_op | 0x1e4dd0 | 2204 | get_unary_func → activation; ActivationOp |
| 223 | simple_unary | 0x260c50 | 3241 | cls-parameterized; raw InstTile + insert | |
| 225 | reciprocal | 0xe9270 | 3269 | ReciprocalOp (via simple_unary cls=) | |
| Select / pred | 195 | select | 0x181b60 | 2507 | TensorSelect or AffSelTensorScalarOp |
| 183 | affine_select | 0x18cf70 | 2312 | AffSelTensorScalarOp | |
| 199 | select_reduce | 0x229540 | 2619 | SelectReduce (ex-CopyPredicatedReduce) | |
| 193 | range_select | 0x1544c0 | 2414 | RangeSelect / RangeSelectReduce | |
| 197 | tensor_copy_predicated | 0x193c50 | 2574 | TensorCopyPredicated | |
| DVE top-8 / stats | 259 | max8 | 0x15e930 | 3799 | SundaMax8 |
| 261 | nc_find_index8 | 0x122ce0 | 3838 | SundaMaxIndex8 | |
| 263 | nc_match_replace8 | 0x104540 | 3899 | SundaMatchReplace8 | |
| 265 | nc_match_replace_indices8 | 0xffe90 | 3973 | MaxIndexAndMatchReplace | |
| 245 | bn_stats_inst | 0xae430 | 3558 | SundaBNStats (Welford partial) | |
| 247 | bn_aggr_inst | 0xd5aa0 | 3586 | SundaBNAggr (Welford aggregate) | |
| 233 | index_value_inst | 0x13bfc0 | 3352 | IndexValueInst (+ IndexValueTile) | |
| Misc / movement | 235 | dropout | 0x1c1000 | 3375 | DropoutMaskInst |
| 161 | quantize_mx | 0x1f9390 | 1837 | QuantizeMXOp ⟵ correction (§7) | |
| 243 | memset | 0x19aac0 | 3514 | MemsetOp | |
| 227 | copy | 0x226860 | 3296 | TensorCopyOp | |
| 229 | tensor_copy | 0x227b10 | 3311 | TensorCopyOp (explicit engine) | |
| 267 | tensor_copy_dynamic_src | 0x1187d0 | — | TensorCopyDynamicSrc | |
| 269 | tensor_copy_dynamic_dst | 0xef890 | — | TensorCopyDynamicDst |
NOTE — ISA-inst module path. The Op nodes for the DVE/stats family derive from a target-ISA inst base. This binary's string table shows the import
neuronxcc.starfish.penguin.targets.tonga.TongaISAInstand…tonga.TongaEnums, and per-op codegen classes underneuronxcc.starfish.penguin.targets.generated.<Op>(e.g.…generated.Exponential,…generated.Activate2). TheSunda…node names in the table are the constructed-class locals; the import path that backs them in this export is the Tonga/generated target. [CONFIRMED stringstonga.TongaISAInst,targets.generated.*]
QUIRK — three builders own no Op of their own (pure delegates).
binoproutes a binary op throughtensortensor;unary_opresolves a unary func and callsactivation;reciprocalcallssimple_unarywithcls=ReciprocalOp. They appear in the roster but neverself.insertdirectly — they exist for call-site ergonomics. [STRONG — decompile]
3. Activation family — ActivationOp / ActivationAccumulationOp
// activation(self, op, tensor, bias, scale, reduce_op, reduce_res,
// reduce_cmd, mask, dtype, schedule, buffer, name, deps) @0x1968c0 py2882
PyObject *activation(...) {
// op = np.ufunc activation function (gelu/silu/exp/tanh/sigmoid/…); stored VERBATIM.
bias_ap = self.create_activation_bias_ap(bias); // helper idx 207/209: per-partition bias AP
require(scale_is_scalar_or_vector, // guard: err_activation_scale_scalar_or_v…
"scale must be scalar or per-partition vector");
// FUSED activation→reduce: reduce_op (ALU) + reduce_cmd (EngineAccumulationType, default Idle)
if (reduce_cmd != EngineAccumulationType.Idle)
return self.activation_accu(op, ..., reduce_cmd, reduce_res); // delegate to accumulate form
inst = ActivationOp(op=op, in=tensor, bias=bias_ap, scale=scale,
reduce_op=reduce_op, reduce_cmd=reduce_cmd,
reduce_res=reduce_res, dtype=dtype, engine=engine, name=name);
self.insert(inst, ...);
}
activation_accu (@0x9f810, py2960) builds ActivationAccumulationOp — the running-sum form, accumulating into reduce_res under reduce_cmd ∈ EngineAccumulationType. It shares create_activation_bias_ap and the scale guard. The pair ActivationOp / ActivationAccumulationOp is the non-accumulate vs accumulate split of the same activation→reduce; reduce_cmd selects Idle/Zero/AddAccumulate/ZeroAccumulate downstream (I04). [CONFIRMED Op names; STRONG accum semantics; EngineAccumulationType string CONFIRMED in this binary]
4. Tensor-Tensor family — TensorTensorOp / TensorTensorScanOp
// tensortensor(self, op, lhs, rhs, mask, dtype, schedule, deps, name, cls) @0x15b850 py2262
PyObject *tensortensor(...) {
cls = cls or TensorTensorOp; // cls kwarg lets callers swap the node class
require(check_tensor_int32_ops_supported(lhs, rhs, op)); // int32 elementwise legality
promote_type(lhs, rhs);
tiles = self.combine_tiles(lhs, rhs); // broadcast/align
engine = pick_engine(); // NeuronEngine ∈ {Vector, GpSIMD, Unknown}
inst = cls(op=op, lhs=tiles.l, rhs=tiles.r, engine=engine, dtype=dtype, name=name);
self.insert(inst, ...); // the binary elementwise op (add/mult/sub/max/min/…)
}
binop (@0x177c80) is the thin wrapper that reads partition_size and forwards a binary op into tensortensor. tensortensorscan (@0x262c60, py3010) builds the two-op cumulative scan TensorTensorScanOp:
// tensortensorscan(self, data0, data1, initial, op0, reverse0, reverse1, mask, dtype, schedule, deps)
// y[i] = op1( op0(data0[i], scan_state), data1[i] ) with `initial` seed, per-op reverse flags
inst = TensorTensorScanOp(data0=data0, data1=data1, initial=initial,
op0=op0, op1=op1, reverse0=reverse0, reverse1=reverse1, dtype=dtype);
op0/op1 are the two chained ALU ops; reverse0/reverse1 flip operand order per chained op. [CONFIRMED args + Op]
5. Tensor-Scalar family — all → TensorScalarPtrOp (+ cache variants)
The defining fact: tensorscalar and tensorvectorscalar build the same TensorScalarPtrOp — one Penguin op for "tensor ∘ up-to-two scalar/vector ops", differentiated only by how many (op, operand) pairs are populated.
// tensorscalar(self, tensor, op0, scalar0, reverse0, mask, dtype, schedule, deps, name) @0x140530 py3066
PyObject *tensorscalar(...) {
// SINGLE alu op: out = op0(tensor, scalar0) (reverse0 swaps operands)
sv = ScalarValue(scalar0); // wrap the immediate
// op0 type = TSOpcode (stored verbatim)
if (op0 == np.power) return self.power_tensorscalar(...); // SPECIAL-CASE decomposition
if (op0 in {mod, fmod, remainder}) // a - floor(a/b)*b
return self.mod_from_remainder(tensor, sv);
inst = TensorScalarPtrOp(in=tensor, op0=op0, scalar0=sv, reverse0=reverse0, dtype=dtype, name=name);
self.insert(inst, ...);
}
// tensorvectorscalar(self, tensor, op0, operand0, reverse0, op1, operand1, reverse1,
// is_scalar_tensor_tensor, mask, dtype, engine, schedule, deps, name, cls) @0x293fd0 py2733
PyObject *tensorvectorscalar(...) {
// TWO chained alu ops: out = op1( op0(tensor, operand0), operand1 )
require(op0 != np.power && op1 != np.power,
"tensorvectorscalar doesn't support power, use tensor_tensor instead.");
self.update_op_and_ptr(inst, op0, operand0, reverse0); // marshal (op, operand-ptr) pair #0
self.update_op_and_ptr(inst, op1, operand1, reverse1); // marshal pair #1
// is_scalar_tensor_tensor flag selects the scalar-tensor-tensor variant
inst = (cls or TensorScalarPtrOp)(in=tensor, /* the two op/operand pairs */, engine=engine, dtype=dtype);
self.insert(inst, ...);
}
QUIRK —
np.powerandmodare decomposed, not emitted as one op.tensorscalar(op0=np.power, ...)routes topower_tensorscalarand carries the guard "Schedule is not supported for np.power, use nisa.tensortensor or nisa.activation explicitly with schedule tuple!";mod/fmod/remainderroute tomod_from_remainder(a − floor(a/b)·b). Both lower to one-or-moreTensorScalarPtrOp/TensorTensorOpnodes — there is no single hardware power/mod scalar op. [CONFIRMED guard strings]
The two cache variants are the accumulate/reduce-fused forms of the same op:
tensorscalarcumulative(@0x1f59b0, py3131) →TensorScalarCacheCumulative— running (cumulative) tensor-scalar underEngineAccumulationType reduce_cmd; inner helperprocess_operand(mdef idx 22). Keeps a running PSUM accumulate.tensorscalarreduce(@0x1ca7f0, py3184) →TensorScalarCacheReduce— tensor-scalar fused with a reduce intoreduce_res(partition_size-tiled). [CONFIRMED Op names; STRONG cmd]
6. Reduce family — TensorReduceOp / PartitionReduceOp
// tensorreduce(self, op, src, mask, dtype, negate, keepdims, schedule, …) @0x239de0 py2098
inst = TensorReduceOp(op=op, reduce_dims=reduce_dims, npartitions=npartitions,
negate=negate, keepdims=keepdims);
// op = ALU reduce op (Union[np.ufunc, ALUOpcode])
// negate = negate-result flag (the HW free-axis reduce-then-negate, e.g. -max in softmax)
// reduce_dims = from inner local NeuronCodegen_tensorreduce_local (reducing_right_most_dims genexpr)
// keepdims = preserve reduced axes
tensor_partition_reduce (@0x24b420, py2162) → PartitionReduceOp is the cross-partition reduce (over the 128-partition axis, not free axes), carrying a reduce_all_axis flag + partition_size. The free-axis reduce (TensorReduceOp) and the partition-axis reduce (PartitionReduceOp) are distinct Penguin nodes — a reimplementer must not collapse them. [CONFIRMED both Ops]
7. quantize_mx → QuantizeMXOp — the MX correction
CORRECTION — the forward-build
QuantizeMXOplives here, inGeneratedNeuronCodegen. An earlier strand (P02 §10) foundquantize_mxonly as aBirCodeGenLoopmacro re-trace and noted the printer (NkiCodegen) has noquantize_mx. Both observations are correct for those surfaces — but the genuine forward builderquantize_mx(mdef idx 161,@0x1f9390, py1837) is in this module and constructs a first-classQuantizeMXOpPenguin node.QuantizeMXOpis a first-class inline op, not a kernel-registry macro (consistent with P06). [CONFIRMED — ctor +self.insertin body]
// quantize_mx(self, src, dst, dst_scale, mask, deps, name, scale_mask) @0x1f9390 py1837
PyObject *quantize_mx(...) {
assert_dtype_in(src, {bfloat16, float16}); // input must be bf16/fp16
// TWO outputs:
// dst = x4-packed quantized data: float8_e4m3fn_x4 / float8_e5m2_x4
// dst_scale = E8M0 per-block scale, dtype uint8
self.check_mx_scale(...); // [P/8, F/4] / block=32 OCP-MXFP E8M0 validator
// (same validator as matmult_mx, P01 §3.3)
// build the block-scale access pattern:
scale_ap = NeuronIndicesAP(generate_subst_map, substituteApIndices,
generate_ap_index, set_shape); // over partition_dim, n_elts/div_ceil block sizing
apply_mask(scale_ap, scale_mask); // scale_mask masks the scale lanes
mask = TileMaskIntersection(data_mask, scale_mask);
inst = QuantizeMXOp(src=src, dst=dst, dst_scale=scale_ap, name=name);
self.insert(inst, ...);
}
quantize_mx is the online activation-quantize-to-MXFP builder: it lowers an fp16/bf16 tile to a {x4-packed-fp8 data, E8M0 scale} pair that feeds nc_matmul_mx (MatMulMXOp, also built here). The MX path therefore originates in this builder and is read downstream by BirCodeGenLoop.codegenQuantizeMX. The check_mx_scale geometry validator ([P/8, F/4], block=32) is shared verbatim with the matmul-MX builder. [CONFIRMED — dtype locals, check_mx_scale local]
8. Select / predicate / range-select family
// select(self, pred, on_true, on_false, mask, dtype, schedule, deps) @0x181b60 py2507
PyObject *select(...) {
if (is_affine_iteration_space_mask(pred)) { // affine path
lower_predicates(pred);
return self.affine_select(pred, on_true, ...); // → AffSelTensorScalarOp
}
inst = TensorSelect(pred=pred, on_true=on_true, on_false=on_false, dtype=dtype); // data-predicate path
self.insert(inst, ...);
}
QUIRK —
selectis a hybrid that builds two different Op classes. A data predicate (a runtime tensor mask) buildsTensorSelectdirectly. An affine iteration-space mask delegates toaffine_select, which buildsAffSelTensorScalarOp— a TensorScalar-class node, not a select node. Both ctor sites and the delegate call are present in theselectbody. A reimplementer who assumesselectalways yields a single node type will miss the affine fast-path. [CONFIRMED — both paths in body]
The rest of the family:
affine_select(@0x18cf70, py2312) →AffSelTensorScalarOp— affine-predicate select expressed as an iteration-space mask, lowered through thepredicateslist. [CONFIRMED]select_reduce(@0x229540, py2619) →SelectReduce— predicated select fused with a reduce (reduce_opALU +reduce_cmdEngineAccumulationType +reduce_res;reverse_predinverts the predicate). A cp311 docstring leak names it: "Implementation of the SelectReduce (formerly CopyPredicatedReduce) instruction." [CONFIRMED]range_select(@0x1544c0, py2414) →RangeSelect/RangeSelectReduce:
// range_select(self, on_true_tile, on_false_value, range_start, comp_op0, comp_op1,
// bound0, bound1, reduce_op, mask, dtype, schedule, deps, name)
// keep on_true where range_start satisfies comp_op0(·,bound0) lo / comp_op1(·,bound1) hi, else on_false_value
ds = DynamicScalar(range_start); // runtime range_start
inst = reduce_op ? RangeSelectReduce(... reduce_cmd=EngineAccumulationType.ResetReduce ...)
: RangeSelect(on_true=on_true_tile, on_false=on_false_value,
range_start=ds, comp_op0=comp_op0, comp_op1=comp_op1,
bound0=bound0, bound1=bound1);
tensor_copy_predicated(@0x193c50, py2574) →TensorCopyPredicated— conditional copy under a predicate; the non-reducing sibling ofselect_reduce. [CONFIRMED]
9. DVE top-8 / stats family — Tonga ISA insts
These build target-ISA inst nodes (from …penguin.targets.tonga.TongaISAInst) — the Data/Vector-Engine top-8 selection and the Welford batch-norm primitives. All read assert_elements_per_partition / assert_max_dimensions and build NDTile + InstTile.
| builder | Op | semantics |
|---|---|---|
max8 (@0x15e930) | SundaMax8 | DVE top-8 (8 largest) per partition; locals max_topk_elements/min_topk_elements (top-8 vs bottom-8) + filter_indices |
nc_find_index8 (@0x122ce0) | SundaMaxIndex8 | index (uint16/uint32) of the 8 max values; builds a NeuronIndicesAP (substituteApIndices/generate_subst_map). Docstring: "its usage is not limited to only finding max values, it can find the index of any value." Tile names nc_find_index8_data/_vals |
nc_match_replace8 (@0x104540) | SundaMatchReplace8 | find-and-replace the 8 matched values with imm; optionally writes dst_idx; can call nc_match_replace_indices8. Tiles nc_match_replace8_data/_vals |
nc_match_replace_indices8 (@0xffe90) | MaxIndexAndMatchReplace | combined max-index + match-replace, fused |
bn_stats_inst (@0xae430) | SundaBNStats | BatchNorm/Welford partial stats (count/mean/M2 per partition tile) |
bn_aggr_inst (@0xd5aa0) | SundaBNAggr | Welford aggregate (combine partial stats across tiles → global mean/var) |
index_value_inst (@0x13bfc0) | IndexValueInst (+ IndexValueTile) | index↔value paired tile op (promote_to_tile_index); pairs a value with its argmax index — the (index,value) bundle the top-k family threads through |
[CONFIRMED Op-class locals; STRONG cross-ref O21/I12 (DVE top-k), I16/O12 (Welford). DVE inst struct field layouts not read — Op names only.]
10. Misc elementwise / data-movement
// dropout(self, tensor, prob, …) @0x1c1000 py3375 → DropoutMaskInst
mask_tile = DataTile(...);
self.memset(mask_tile, seed); // seed the mask via memset
inst = DropoutMaskInst(in=tensor, prob=prob, mask=mask_tile, dtype=dtype); // prob = keep/drop probability
// memset(self, tile_shape, value, dtype, mask, dst_tile, schedule, engine, deps, name) @0x19aac0 py3514 → MemsetOp
sv = ScalarValue(value);
require(engine in {Vector, GpSIMD, Unknown},
"Memset engine can only be Vector or GpSIMD or Unknown."); // Vector/GpSIMD fill, NOT DMA
inst = MemsetOp(shape=tile_shape, value=sv, engine=engine, dtype=dtype);
GOTCHA —
memsetis a compute-engine fill, not a DMA. The guard restricts the engine toVector/GpSIMD/Unknown— there is no DMA path here. DMA-class moves live in the memory family (6.5.x / P03). [CONFIRMED guard string]
The copy variants:
copy(@0x226860, py3296) →TensorCopyOp— plain tile copy.tensor_copy(@0x227b10, py3311) →TensorCopyOp— near-duplicate ofcopyadding explicitNeuronEngineselection.tensor_copy_dynamic_src(@0x1187d0) →TensorCopyDynamicSrc;tensor_copy_dynamic_dst(@0xef890) →TensorCopyDynamicDst— the runtime register-addressed copies (src/dst address from a register; the DynamicScalar/offset family at copy granularity). [CONFIRMED Op names]
11. Adversarial self-verification
The five strongest claims on this page, re-challenged against the binary:
- Class is
GeneratedNeuronCodegen, not bareNeuronCodegen. ✅ The wrapper symbols innames.jsonare__pyx_pw_…KernelBuilder_22GeneratedNeuronCodegen_<idx><name>; the22length prefix = the 22-charGeneratedNeuronCodegen. P01/S2-09's "NeuronCodegen" is the logical base; the compiled subclass isGeneratedNeuronCodegen. Issued as an in-place CORRECTION. [CONFIRMED] quantize_mx→QuantizeMXOpis a forward builder in this module. ✅ idx 161 @0x1f9390 with aQuantizeMXOpctor +self.insert(report body read); corrects P02 §10's "only in BirCodeGenLoop". The MatMulMX consumer andcheck_mx_scalevalidator corroborate it. [CONFIRMED — body; the partial IDA export does not re-read this body, so the address is STRONG]tensorscalarandtensorvectorscalarboth build the singleTensorScalarPtrOp. ✅ Both bodies constructTensorScalarPtrOp; the difference is one vs two(op, operand)pairs (update_op_and_ptrcalled once vs twice). [CONFIRMED both bodies]- Enum marshalling is pass-through; renumbering is downstream. ✅ Type-hint strings
Union[np.ufunc, ALUOpcode]/TSOpcode/Optional[EngineAccumulationType]pin the kwarg types;EngineAccumulationTypeis a CONFIRMED string in this binary (@0x24870). Numeric ISA mapping isBirCodeGenLoop.codegenAluOp/codegenAccumCmd(I04/I05). [STRONG] self.insert(@0x165f00) reachesIRBuilder.add_named_instruction. ✅add_named_instructionsymbol +insertstring present; the wrap order (add_named_instruction → predicates → deps → bookkeeping → debugloc) is from.rodata+ cross-strand P03. [CONFIRMED symbols; STRONG order]
Tagged uncertainties. The per-method addresses, mdef indices, and py-lines in §2 are read from the report's nm/DWARF; the IDA JSON export checked into this tree only re-confirms the tail window (idx 1–17) plus the class name, import modules, and the EngineAccumulationType/combine_tiles/insert strings — so per-method facts are STRONG, not independently re-CONFIRMED here. The Sunda… Op-class names are constructed-local names; the import path this export shows is targets.tonga.TongaISAInst / targets.generated.<Op>. The exact positional kwarg order at each <Op>(...) ctor is mstate-routed and not byte-traced (kwarg names are confirmed; order is INFERRED). DVE inst struct field layouts are not read.
See also
- 6.5.1 NeuronCodegen Forward Builder — overview & matmul — the same class, the matmul/memory halves, and
self.insertin detail. - 6.5.9 NeuronCodegen Re-emit Printer — the inverse surface (Penguin node →
nisa.<prim>text); reads back the enums this builder stores. - 5.6 Penguin Tensor-Op Family — the Penguin-IR
<Op>nodes these builders construct. - BIR codegen loop (strands I04/I05) —
BirCodeGenLoop.codegenAluOp/codegenAccumCmdnumeric ISA renumbering, one layer below this builder.