Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

NeuronCodegen Tensor-Op Forward Builders

All symbols and addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22, module neuronxcc/nki/compiler/backends/neuron/KernelBuilder.cpython-310-x86_64-linux-gnu.so (the cp311/cp312 twins are source-identical; their Op set matches). The .so is an UNSTRIPPED 14.5 MB Cython extension with DWARF debug_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 (selectTensorSelect or AffSelTensorScalarOp). The MX-quantize builder quantize_mxQuantizeMXOp lives here, correcting an earlier strand that placed it only in BirCodeGenLoop.

CORRECTION — the Cython class is GeneratedNeuronCodegen, not bare NeuronCodegen. Sibling strand P01 (and the S2-09 source survey) name the owning class NeuronCodegen. The symbol mangling in this binary is unambiguous: every method wrapper is __pyx_pw_9neuronxcc_9generated_3nki_8compiler_8backends_6neuron_13KernelBuilder_22GeneratedNeuronCodegen_<idx><name>. The 22 length-prefix decodes to the 22-character identifier GeneratedNeuronCodegen — the Cython-emitted concrete class that the methods actually compile into (a generated subclass; NeuronCodegen is the hand-written base). This page uses GeneratedNeuronCodegen when citing symbols and NeuronCodegen when speaking of the logical surface; they are the same builder. [CONFIRMED — symbol mangling, names.json]

Module / classKernelBuilder.cpython-310…so · class GeneratedNeuronCodegen (logical NeuronCodegen)
Directionnl/nisa-Python args → construct penguin.ir.<Op> node → self.insert
Insert entryself.insert (mdef idx 297, @0x165f00) → IRBuilder.add_named_instruction
Tile normalizecombine_tiles (@0x16ad90) + tile.canonical_par_indices / canonical_free_indices(...)
Per-op modulesneuronxcc.starfish.penguin.targets.generated.<Op> · ISA base …targets.tonga.TongaISAInst
Enum policypass-through; numeric ISA renumber deferred to BirCodeGenLoop codegen
Builder count~30 tensor-op methods (of 194 total __pyx_mdef_…GeneratedNeuronCodegen_<N><name>)
MX builderquantize_mx (idx 161, @0x1f9390) → QuantizeMXOp — lives here, not only in BirCodeGenLoop

NOTE — evidence base and its limits. Method addresses, mdef indices, and KernelBuilder.py source lines below are read from the binary's nm symbol table (__pyx_pw/__pyx_mdef per method) and DWARF decodedline, the Op-class names from the constructed pyx_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 methods rand2/rng/exponential/activate2/nonzero_with_count/*_read_accumulator, idx 1–17 — confirming the class name, the import modules, and EngineAccumulationType/combine_tiles/insert strings), 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_instruction is 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.insert reaches add_named_instruction for it. A reimplementer who hunts for a direct add_named_instruction(...) at each builder will not find one — it is always behind self.insert. [CONFIRMED — add_named_instruction symbol present; insert string 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:

Selectorkwarg(s)Type hint (.rodata)Carried onConfidence
ALU opopUnion[np.ufunc, ALUOpcode]TensorTensorOp / TensorReduceOp / ActivationOpSTRONG (type-hint)
TensorScalar opop0 / op1TSOpcodeTensorScalarPtrOpSTRONG (type-hint)
Accumulate cmdreduce_cmdOptional[EngineAccumulationType] (default Idle; reset member ResetReduce)activation / select-reduce / tensor-scalar-cacheCONFIRMEDEngineAccumulationType string present @0x24870
Activation funcop (of activation)np.ufuncActivationOpSTRONG (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 the np.ufunc/enum objects this builder stores verbatim. Builder = store; printer = name-map; BirCodeGenLoop = numeric ISA renumber (codegenAluOp 1→29; codegenAccumCmd Idle/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.

Familyidxbuilderpw addrpypenguin.ir.<Op>
Activation211activation0x1968c02882ActivationOp
213activation_accu0x9f8102960ActivationAccumulationOp
Tensor-Tensor181tensortensor0x15b8502262TensorTensorOp (cls=, default)
179binop0x177c802221delegates → tensortensor
215tensortensorscan0x262c603010TensorTensorScanOp (2-op scan)
Tensor-Scalar217tensorscalar0x1405303066TensorScalarPtrOp (op0/scalar0/reverse0)
203tensorvectorscalar0x293fd02733TensorScalarPtrOp (op0/op1 + operand0/1 + reverse0/1)
219tensorscalarcumulative0x1f59b03131TensorScalarCacheCumulative
221tensorscalarreduce0x1ca7f03184TensorScalarCacheReduce
Reduce173tensorreduce0x239de02098TensorReduceOp (op/reduce_dims/negate/keepdims)
175tensor_partition_reduce0x24b4202162PartitionReduceOp (reduce_all_axis)
Unary177unary_op0x1e4dd02204get_unary_func → activation; ActivationOp
223simple_unary0x260c503241cls-parameterized; raw InstTile + insert
225reciprocal0xe92703269ReciprocalOp (via simple_unary cls=)
Select / pred195select0x181b602507TensorSelect or AffSelTensorScalarOp
183affine_select0x18cf702312AffSelTensorScalarOp
199select_reduce0x2295402619SelectReduce (ex-CopyPredicatedReduce)
193range_select0x1544c02414RangeSelect / RangeSelectReduce
197tensor_copy_predicated0x193c502574TensorCopyPredicated
DVE top-8 / stats259max80x15e9303799SundaMax8
261nc_find_index80x122ce03838SundaMaxIndex8
263nc_match_replace80x1045403899SundaMatchReplace8
265nc_match_replace_indices80xffe903973MaxIndexAndMatchReplace
245bn_stats_inst0xae4303558SundaBNStats (Welford partial)
247bn_aggr_inst0xd5aa03586SundaBNAggr (Welford aggregate)
233index_value_inst0x13bfc03352IndexValueInst (+ IndexValueTile)
Misc / movement235dropout0x1c10003375DropoutMaskInst
161quantize_mx0x1f93901837QuantizeMXOp ⟵ correction (§7)
243memset0x19aac03514MemsetOp
227copy0x2268603296TensorCopyOp
229tensor_copy0x227b103311TensorCopyOp (explicit engine)
267tensor_copy_dynamic_src0x1187d0TensorCopyDynamicSrc
269tensor_copy_dynamic_dst0xef890TensorCopyDynamicDst

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.TongaISAInst and …tonga.TongaEnums, and per-op codegen classes under neuronxcc.starfish.penguin.targets.generated.<Op> (e.g. …generated.Exponential, …generated.Activate2). The Sunda… 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 strings tonga.TongaISAInst, targets.generated.*]

QUIRK — three builders own no Op of their own (pure delegates). binop routes a binary op through tensortensor; unary_op resolves a unary func and calls activation; reciprocal calls simple_unary with cls=ReciprocalOp. They appear in the roster but never self.insert directly — 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.power and mod are decomposed, not emitted as one op. tensorscalar(op0=np.power, ...) routes to power_tensorscalar and carries the guard "Schedule is not supported for np.power, use nisa.tensortensor or nisa.activation explicitly with schedule tuple!"; mod/fmod/remainder route to mod_from_remainder (a − floor(a/b)·b). Both lower to one-or-more TensorScalarPtrOp/TensorTensorOp nodes — 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 under EngineAccumulationType reduce_cmd; inner helper process_operand (mdef idx 22). Keeps a running PSUM accumulate.
  • tensorscalarreduce (@0x1ca7f0, py3184) → TensorScalarCacheReduce — tensor-scalar fused with a reduce into reduce_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_mxQuantizeMXOp — the MX correction

CORRECTION — the forward-build QuantizeMXOp lives here, in GeneratedNeuronCodegen. An earlier strand (P02 §10) found quantize_mx only as a BirCodeGenLoop macro re-trace and noted the printer (NkiCodegen) has no quantize_mx. Both observations are correct for those surfaces — but the genuine forward builder quantize_mx (mdef idx 161, @0x1f9390, py1837) is in this module and constructs a first-class QuantizeMXOp Penguin node. QuantizeMXOp is a first-class inline op, not a kernel-registry macro (consistent with P06). [CONFIRMED — ctor + self.insert in 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 — select is a hybrid that builds two different Op classes. A data predicate (a runtime tensor mask) builds TensorSelect directly. An affine iteration-space mask delegates to affine_select, which builds AffSelTensorScalarOp — a TensorScalar-class node, not a select node. Both ctor sites and the delegate call are present in the select body. A reimplementer who assumes select always 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 the predicates list. [CONFIRMED]
  • select_reduce (@0x229540, py2619) → SelectReduce — predicated select fused with a reduce (reduce_op ALU + reduce_cmd EngineAccumulationType + reduce_res; reverse_pred inverts 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 of select_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.

builderOpsemantics
max8 (@0x15e930)SundaMax8DVE top-8 (8 largest) per partition; locals max_topk_elements/min_topk_elements (top-8 vs bottom-8) + filter_indices
nc_find_index8 (@0x122ce0)SundaMaxIndex8index (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)SundaMatchReplace8find-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)MaxIndexAndMatchReplacecombined max-index + match-replace, fused
bn_stats_inst (@0xae430)SundaBNStatsBatchNorm/Welford partial stats (count/mean/M2 per partition tile)
bn_aggr_inst (@0xd5aa0)SundaBNAggrWelford 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 — memset is a compute-engine fill, not a DMA. The guard restricts the engine to Vector/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 of copy adding explicit NeuronEngine selection.
  • 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:

  1. Class is GeneratedNeuronCodegen, not bare NeuronCodegen. ✅ The wrapper symbols in names.json are __pyx_pw_…KernelBuilder_22GeneratedNeuronCodegen_<idx><name>; the 22 length prefix = the 22-char GeneratedNeuronCodegen. P01/S2-09's "NeuronCodegen" is the logical base; the compiled subclass is GeneratedNeuronCodegen. Issued as an in-place CORRECTION. [CONFIRMED]
  2. quantize_mxQuantizeMXOp is a forward builder in this module. ✅ idx 161 @0x1f9390 with a QuantizeMXOp ctor + self.insert (report body read); corrects P02 §10's "only in BirCodeGenLoop". The MatMulMX consumer and check_mx_scale validator corroborate it. [CONFIRMED — body; the partial IDA export does not re-read this body, so the address is STRONG]
  3. tensorscalar and tensorvectorscalar both build the single TensorScalarPtrOp. ✅ Both bodies construct TensorScalarPtrOp; the difference is one vs two (op, operand) pairs (update_op_and_ptr called once vs twice). [CONFIRMED both bodies]
  4. Enum marshalling is pass-through; renumbering is downstream. ✅ Type-hint strings Union[np.ufunc, ALUOpcode]/TSOpcode/Optional[EngineAccumulationType] pin the kwarg types; EngineAccumulationType is a CONFIRMED string in this binary (@0x24870). Numeric ISA mapping is BirCodeGenLoop.codegenAluOp/codegenAccumCmd (I04/I05). [STRONG]
  5. self.insert (@0x165f00) reaches IRBuilder.add_named_instruction.add_named_instruction symbol + insert string 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