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

BatchNorm-Family Encoding — Stats / Aggregate / Gradients / Backprop

All symbols and addresses on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310 wheel; cp310/11/12 share the C++ core logic — see the parity note in Build & Version Provenance). The five encoder bodies, the dtype LUT, the IMM_SRC scalar encoder and the pybind assert roster live in libwalrus.so (build-id 92b4d331…, VA == file offset); the bir::Inst* ordinals and the NEURON_ISA_TPB_*_BN_STRUCT shapes come from libBIR.so (a9b1ea38) and the neuron_isa_tpb_pybind.cpython-310-*.so assert table. Treat every address as version-pinned.

Abstract

This page is the byte-for-byte field map of the five batch-norm instructions the Pool engine emits onto the wire, produced by neuronxcc::backend::CoreV2GenImpl: visitInstBNStats, visitInstBNStatsAggregate, visitInstBNGradients, visitInstBNBackprop, and visitInstBNBackprop2. Between them they encode the full Welford / two-pass batch-norm statistics datapath — the forward partial-stats reduction, the cross-tile merge, the backward gradient reduction, and the dx apply — onto the 64-byte bundle the rest of the ISA shares (2.1 the bundle).

Three things make this family distinct from the elementwise ops. First, BNStats is mixed-width: the input access pattern is a TENSOR4D (the reduced ofmap) landing at +0x0C, but the output is a TENSOR2D — six fp32 per partition, two Welford triples — landing at +0x30 (Family-C anchors, 2.1 BNStats mixed-width note). Second, the two backprop ops emit two bundles each: a control / param-load bundle that carries the coefficients, N and Mean, followed by a data / back-prop bundle that carries the x / dy / dx tensors — and the param-load bundle is written first on the wire. Third, the BNGradients op forks by arch level: a pre-v5 AP/scalar-coefficient form (opcode 0x63) and a core_v5+ register-or-AP-coefficient form (opcode 0x94), with the older form FATAL-asserting on core_v5.

The bar for this page: a reader can byte-encode any BN-family instruction by hand, knowing for each control byte its offset, width, semantic, value source, the pybind assert string that names it, and the disassembly store-site that pins it. Every field row carries a confidence tag — CONFIRMED (exact store disassembled this pass, offset cited), STRONG (predicate / LUT / opcode cross-checked), INFERRED (zero-init implied by the bulk movups zero-fill, no direct store). The Welford math itself is the silicon's; the encoder only seeds the per-lane counts and chooses which operand lands in which descriptor slot. No field name is fabricated; where the binary exposes only a 16-char-truncated assert slot, the reconstructed tail is shown in brackets.

At a glance

Every BN bundle is a std::array<unsigned char, 64>: emplace_back into the engine's SmallVector, four movups xmm0 to zero all 64 bytes, header-stamped by setupHeader (vtable slot +0x48), field-filled, ISA-checked, and fwrite(buf, 1, 0x40, bin)-ed — the common skeleton of 2.1. The header word is (0x10<<8) | opcode, little-endian.

BIR inst / bundle roleGenerator (libwalrus)Opcode byte / wordISA mnemonicWire struct
BNStats (no transpose)CoreV2GenImpl::visitInstBNStats @0x123ab000x61 / 0x1061BatchNormStats2S4D2_BN
BNStats (transpose)(same body; opcode toggled)0x82 / 0x1082TransposeBatchNormStats2S4D2_BN
BNStatsAggregateCoreV2GenImpl::visitInstBNStatsAggregate @0x123c6a00x62 / 0x1062BatchNormAggregateS4D2_BN
BNGradients (AP/scalar coeff)CoreV2GenImpl::visitInstBNGradients @0x1241b100x63 / 0x1063BatchNormGradAccumulateS3S3D1_BN
BNGradients (register coeff)(same body; arch-gated)0x94 / 0x1094BatchNormGradAccumulate2S3S3D1_BN2
BNBackprop ctrlCoreV2GenImpl::visitInstBNBackprop @0x1235e800x64 / 0x1064BatchNormParamLoadS2_BN
BNBackprop2 ctrlCoreV2GenImpl::visitInstBNBackprop2 @0x12368e00x8E / 0x108EBatchNormParamLoad2S2_BNPL2
BNBackprop / BNBackprop2 data(both backprop bodies)0x65 / 0x1065BatchNormBackPropS3S3D3_TT
EnginePool — read-only EngineInfo at Inst+0x90 for the per-engine census; not chosen by the encoder (STRONG)
Bundle sizeexactly 64 bytes (std::array<u8,64> + fwrite(…,0x40,…)) — CONFIRMED
Header+0x00 opcode, +0x01 inst_word_len=0x10, +0x02..03 reserved=0 — CONFIRMED
Slot familyFamily C (BNStats/Aggregate, 4D-in at the +0x0C low anchor; 2D-out at the +0x30 high anchor — the A/B 3D-out slot, NOT the Family-C +0x2C 4D-out anchor, see note); Family B (Gradients/data, 3×3D +0x10/+0x20/+0x30) — CONFIRMED
dtype LUTsub_120E650byte_1DF5760 (BIR dtype ordinal → ISA wire tag) — STRONG
IMM_SRC encodersub_12051E0 (setupImmediate<NEURON_ISA_TPB_IMM_SRC>) — Mean / split-coeff scalar imm — CONFIRMED (call site)
CodeGenMode1 GENERATE (wire), 2 RUN_ISA_CHECKS (stack scratch, same layout), 0 COLLECT_OPCODES (census) — CONFIRMED all 5 bodies

CORRECTION — BNStats is a Family-C op whose output anchor is an exception. BNStats/Aggregate are Family-C instructions (input rides the Family-C low anchor +0x0C as a TENSOR4D), but their narrow 2-D Welford output does not land at the canonical Family-C destination anchor +0x2C (the 20-byte 4D-out slot). It lands at +0x30 — the high anchor that Families A and B use for their 3D output (2.1 the bundle, Family A dst +0x30). This is the documented mixed-width case: the anchor position is fixed but the descriptor width is per-operand, and a TENSOR2D stat pattern is written to the +0x30 high slot (lea 0x30(%r13) + assignAccess<TENSOR2D> @0x123adf6), not the +0x2C 4D slot. So an earlier shorthand reading "Family-C ⇒ dest +0x30" conflates two anchors: Family-C's real dest anchor is +0x2C; BNStats deliberately uses the +0x30 A/B anchor because its output is 2-D, not 4-D. The per-op field maps below pin +0x30 directly.


The common frame

All five bodies share the 2.1 prologue verbatim: emplace a 64-byte array, zero all 64 bytes (4×movups xmm0), then a virtual call through vtbl+0x48 to setupHeader, which writes only bundle[0]=opcode, bundle[1]=0x10 (inst_word_len = 16 dwords = 64 B), bundle[2:4]=0. Every byte the body does not explicitly write is a hard 0x00 on the wire.

CodeGenMode (this+156 dword) selects three arms in every body, exactly as for the matmul encoders:

// shared mode dispatch at the top of every visitInstBN* body
switch (self->codeGenMode) {            // *(int*)(self+156)
  case 1:  /* GENERATE_ISACODE   */     // emplace 64-B array, fill, ISA-check, fwrite 0x40
  case 2:  /* RUN_ISA_CHECKS     */     // fill a STACK temp (rbp-relative), ISA-check, no fwrite
  case 0:  /* COLLECT_OPCODES    */     // census only: insert opcode into set<u32>, no struct
  default: FATAL("Wrong CodeGenMode. It must be one of "
                 "GENERATE_ISACODE, RUN_ISA_CHECKS, or COLLECT_OPCODES");
}

The RUN_ISA_CHECKS arm builds the same field layout into an rbp-relative stack temp before running the ISA checker — which is itself a second, independent confirmation of every offset, since the two arms write the identical fields to two different bases. The COLLECT_OPCODES arm only inserts the opcode integer; its literal imm is the opcode-table proof.

GOTCHA — disassembling a visitInstBN* body you will see two descriptor writes for the same field: an [rbp+OFF] write (the RUN_ISA_CHECKS scratch) and an [r12/r13/r14/rbx+OFF] write (the live wire slot in GENERATE mode). Only the latter is the 64-byte bundle. Reading the rbp form as the wire layout is the classic mis-pinned-offset trap — though here the two layouts happen to agree, which is what makes the RUN arm a useful cross-check.

NOTE — engine. Every body reads EngineInfo at Inst+0x90 to key a per-engine census DenseMap; the physical engine is not chosen by the encoder and is not a wire field. The BN opcodes occupy the Pool-engine block of the _OPCODE table (grouped with Pool / Max8 / FindIndex8), so engine = POOL is STRONG by opcode grouping and partition-reduce semantics, not by a hard encoder assert.


visitInstBNStatsBatchNormStats2 (op 0x61 / 0x82, S4D2_BN)

Purpose

BNStats is the forward partial-statistics reduction: it folds the free dimension of the reduced ofmap into two per-lane Welford triples[n_A, mean_A, M2_A, n_B, mean_B, M2_B], six fp32 per partition. Lane A accumulates the even-indexed free elements (ceil(N/2) of them), lane B the odd (floor(N/2)); the silicon merges them at run time. The encoder's only numeric job is to seed the two counts; the mean/M2 words are produced by the hardware accumulator. M2 is the raw sum-of-squared-deviations — normalisation to population variance happens only in Aggregate.

Pre-guards and opcode select

// CoreV2GenImpl::visitInstBNStats @0x123ab00
function visitInstBNStats(self, Inst):
    if *(u8*)(Inst+0x110) != 0:                  // apply_transpose MaybeAffine<bool> variant TAG
        throw bad_variant_access;                 //   @0x123ab2e — the bool must be a constant
    // opcode from apply_transpose (Inst+0xF0): 0 -> 0x61, 1 -> 0x82
    opcode = (Inst[0xF0] == 1) ? 0x82 : 0x61;     // sbb/and 0xFFFFFFDF/sub 0x7E  @0x123ab5f
    ...
    // POST-fill: reject an indirect/symbolic input AP
    if isTensorIndirect(arg0_AP):                 // (vtbl+0x80)(arg0)
        reportError("TODO: support BATCH_NORM_STATS2 with TensorIndirect AP "
                    "once ISA supports it");      //   @0x123ad8e / 0x123b013

The transpose flag is folded into the opcode (0x61 vs 0x82 = TransposeBatchNormStats2); it is not a separate bundle field. The arithmetic is byte-verified: cmpb $0x1,0xf0(%rsi); sbb %r14d,%r14d; and $0xffffffdf,%r14d; sub $0x7e,%r14dtranspose=00x82-(-1&~0x20)-0 = 0x61, transpose=10x82. CONFIRMED.

Control-band field map (GENERATE path, base = r13)

OffWFieldpybind assertValue / sourceStore-siteConf
+0x001opcode 0x61/0x82s_batchnorm_stats_opcodesetupHeader(apply_transpose)hdrCONFIRMED
+0x011inst_word_len=16(header)setupHeader b1=0x10hdrCONFIRMED
+0x022format/reserved=0(header)setupHeaderhdrCONFIRMED
+0x0C16IN AP TENSOR4Ds_bnstats2_src_element_cntassignAccess<TENSOR4D>(arg0)lea 0xc(%r13) @0x123ade5CONFIRMED
+0x201IN dtype (ISA tag)bnstats2/s2_bn dtypesub_120E650([arg0+0x30])mov %al,0x20(%r13) @0x123acecCONFIRMED
+0x211OUT dtype (ISA tag)d2_bn_stats_out_[ty]sub_120E650([out+0x30])mov %al,0x21(%r13) @0x123ad07CONFIRMED
+0x221BASE partitions_bnstats2 src/base*(*(arg0+0x50)+8) (APPair0)mov %al,0x22(%r13) @0x123ad29CONFIRMED
+0x284EVEN count n_A (f32)s even_count must be finite / d2_bn_zero_counts(float)((InNEPP+1)>>1) = ceil(N/2)movss %xmm1,0x28(%r13) @0x123ad45CONFIRMED
+0x2C4ODD count n_B (f32)s odd_count must be finite / d2_bn_zero_counts(float)(InNEPP>>1) = floor(N/2)movss %xmm1,0x2c(%r13) @0x123ad75CONFIRMED
+0x3016OUT AP TENSOR2Ds_bnstats2_dst_e[lement_cnt]assignAccess<TENSOR2D>(out)lea 0x30(%r13) @0x123adf6CONFIRMED
+0x04..0B, +0x23..27RESERVED-0d2_bn_reserved_z[ero]bulk movups zero-initINFERRED

InNEPP = getNumElementsPerPartition(arg0). The count seeds are computed by add $1,%rax; shr $1,%rax; cvtsi2ss (ceil) and shr $1,%rax; cvtsi2ss (floor) — the same %rax reused — CONFIRMED @0x123ad39..0x123ad75. The output AP is assignAccess<TENSOR2D>: 12 bytes encoding the 6-fp32 destination, so the OUT operand uses the +0x30 high anchor (Family-C) with a narrow 2-D pattern, not a 4-D one — this is the mixed-width case 2.1 calls out. The verifier asserts out-NEPP == 6.

NOTE — base partition. +0x22 reads *(*(arg0+0x50)+8), the partition of the input AP's first APPair, guarded by [arg0+0x58]!=0 (else assert "idx < size()" @0x123b0b0) — the AP must carry at least one pair. The same (*(AP+0x50)+8) access recurs for the base byte of every BN op.

Welford seeding

// the ONLY numeric values the encoder writes into the 6-elt output's frame
n_A = ceil(N/2);   // +0x28, lane A = even-indexed free elements
n_B = floor(N/2);  // +0x2C, lane B = odd-indexed free elements
// mean_A, M2_A, mean_B, M2_B are produced by the silicon Welford accumulator at run time.
// The two "*_count must be finite" asserts validate the f32 seeds;
// d2_bn_zero_counts is the N==0 -> counts-must-be-zero guard.

visitInstBNStatsAggregateBatchNormAggregate (op 0x62, S4D2_BN)

Purpose

Aggregate is the cross-tile merge: it takes K partial Welford triples (K·3 fp32, the concatenated outputs of K BNStats runs) and reduces them to (mean, population-variance) — exactly two fp32 per partition. The Chan / parallel-variance merge (n_a+n_b, pooled mean, M2_a+M2_b + n_a·n_b/(n_a+n_b)·δ², then the divide-by-N to turn M2 into population variance) is the engine's; the encoder reads the counts from the packed input triples, not from the bundle.

It is the simplest of the family — a strict field subset of BNStats, minus the two count seeds (+0x28/+0x2C) and minus the transpose toggle (opcode is the constant 0x62).

Control-band field map (GENERATE path, base = r12)

OffWFieldpybind assertValue / sourceStore-siteConf
+0x001opcode 0x62 (const)s_batchnorm_aggregate_opcodesetupHeader(98)movb $0x62 @0x123c98eCONFIRMED
+0x0C16IN AP TENSOR4Ds_bnstats_agg_src_element_cntassignAccess<TENSOR4D>(arg0)lea 0xc(%r12) @0x123c825CONFIRMED
+0x201IN dtyped2_bn_agg_in_typ[e]sub_120E650([arg0+0x30])mov %al,0x20(%r12) @0x123c85fCONFIRMED
+0x211OUT dtyped2_bn_agg_out_ty[pe]sub_120E650([out+0x30])mov %al,0x21(%r12) @0x123c88dCONFIRMED
+0x221BASE partitions_bnstats_agg_dst / base*(*(arg0+0x50)+8)mov %al,0x22(%r12) @0x123c84dCONFIRMED
+0x3016OUT AP TENSOR2Ds_bnstats_agg_dst_element_cntassignAccess<TENSOR2D>(out)lea 0x30(%r12) @0x123c86bCONFIRMED
+0x28/+0x2C, +0x04..0B, +0x23..27RESERVED-0 (no count seeds)d2_bn_reserved_z[ero]zero-initCONFIRMED (absence of stores)

The verifier asserts in-NEPP % 3 == 0 (K triples) and out-NEPP == 2. The input is a TENSOR4D (assignAccess<TENSOR4D>, byte-verified) — the descriptor type is 4-D even though the data is a flat K·3 stat vector.


visitInstBNGradientsBatchNormGradAccumulate (op 0x63 / 0x94)

Purpose

BNGradients is the backward reduction: it folds x and dy against the forward statistics into a 3-element gradient vector — [Σdy, Σ(x−m)·dy·vf, Σ(x−m)·dy·vf²] — used to assemble the dx coefficients. One 64-byte bundle, but two encoding variants chosen by an arch-feature gate plus the operand kind: an AP/scalar-coefficient form (0x63, struct S3S3D1_BN) and a register-or-AP form (0x94, struct S3S3D1_BN2).

Operand and arch gate

// CoreV2GenImpl::visitInstBNGradients @0x1241b10
function visitInstBNGradients(self, Inst):
    in2 = getMeanAccessor(Inst);                  // sub_1240360 — "in2 must be AP" gate
    in3 = getArgument(Inst, 3);                    // Varfactor
    out = getOutput<AccessPattern>(Inst, 0);       // Gradients (3-elt)
    assert(isAP(in2) && isAP(in3),  "in2 and in3 must be AP argument type");         // line 1562
    assert(isLocationSB(in2) || isLocationPSUM(in2),
           "ISA immediate pointer should only be SB/PSUM accesses");                  // 1564/1566
    assert(in2.dtype == in3.dtype,  "in2.getType() == in3.getType()");               // NeuronAssertion 956

    if *(int*)(self+0x258) <= 0x31:                // arch level <= 49  (pre-v5)   @0x1242149
        if ModuleArchLevel >= ArchLevel::core_v5:
            FATAL("BatchNormGradAccumulate not supported in core_v5");                // NeuronAssertion 2101
        opcode = 0x63;  struct = S3S3D1_BN;        // VARIANT A — AP/scalar coeffs
    else:                                          // arch level > 49
        opcode = 0x94;  struct = S3S3D1_BN2;       // VARIANT B — register-or-AP coeffs

The gate is byte-verified: cmpl $0x31,0x258(%r12) @0x1242149. The core_v5 FATAL means the 0x63 AP-coeff form is illegal on core_v5; the register-coeff 0x94 path is its replacement.

CORRECTION (BN-1 to D-J09) — the 0x63 BatchNormGradAccumulate form FATAL-asserts on core_v5 (NeuronAssertion 2101, "BatchNormGradAccumulate not supported in core_v5"); 0x94 (BatchNormGradAccumulate2) is the core_v5+ register-coeff form. The self+0x258 <= 49 test is the pre-v5 selector. This gate was not in the earlier (D-J09) gradient map.

Shared AP-slot template (generateBNGradients<…> @0x133e550 / @0x133e630)

Both variants share a byte-identical AP layout, filled by the template helper before the variant-specific body. arg0 = Ofmap(x), arg1 = OfmapGrad(dy), out = Gradients(3-elt). S3S3D1 = Src-3D, Src-3D, Dst-1D:

OffWFieldValue / sourceStore-siteConf
+0x3C1packed dtype: low nibble = sub_120E650(Ofmap.D), high nibble = 16*sub_120E650(OfmapGrad.D)RMW &0xF0 then &0x0Fmov %al,0x3c(%rbx) ×2 @0x133e5a3/@0x133e5bdCONFIRMED
+0x3F1BASE partition*(*(Ofmap+0x50)+8) (guard [+0x58]!=0)mov %al,0x3f(%rbx) @0x133e5deCONFIRMED
+0x1016Ofmap AP TENSOR3DassignAccess<TENSOR3D>(arg0)lea 0x10(%rbx) @0x133e5ceCONFIRMED
+0x2016OfmapGrad AP TENSOR3DassignAccess<TENSOR3D>(arg1)lea 0x20(%rbx) @0x133e5e6CONFIRMED
+0x3016Gradients AP TENSOR1DassignAccess<TENSOR1D>(out)lea 0x30(%rbx) @0x133e5f7CONFIRMED

Variant A — op 0x63 (S3S3D1_BN, AP/scalar coefficients)

After the template, the body writes the scalar coefficient values and their dtypes:

OffWFieldpybind assertValue / sourceStore-siteConf
+0x0C4Mean (in2) scalar VALUEga_imm_check / ga_valid_immedia[te](self[76]→vtbl+0x20)(in2, 0) (4-byte wire imm)mov %eax,0xc(%r13) @0x12422abCONFIRMED
+0x384Varfactor (in3) scalar VALUEga_imm_check(self[76]→vtbl+0x20)(in3, 0)mov %eax,0x38(%r13) @0x12422e2CONFIRMED
+0x3D1in2 (Mean) dtypebnga_valid_input_t[ype]sub_120E650(in2.D)mov %al,0x3d(%r13) @0x12422fdCONFIRMED
+0x3E1out (Gradients) dtypega_valid_output_[type]sub_120E650(out.D)mov %al,0x3e(%r13) @0x1242299CONFIRMED

Full operands: AP0 Ofmap, AP1 OfmapGrad, in2 Mean(scalar), in3 Varfactor(scalar), OUT Gradients(3-elt).

Variant B — op 0x94 (S3S3D1_BN2, register-or-AP coefficients)

The dynamic / HWDGE form: Mean and Varfactor may each be a physical AP scalar or a register id, tagged by a 2-bit mode nibble at +0x3E:

OffWFieldpybind assertValue / sourceStore-siteConf
+0x3E1coeff MODE nibbles: low (in2) |0x1=AP / |0x2=reg; high (in3) |0x10=AP / |0x20=regga2_valid_outputRMW (&0xF0/&0x0F preserve)mov %al,0x3e(%r13) @0x124282d (reg)CONFIRMED
+0x0C4 or 1in2 value: AP → (self[76]→vtbl+0x20) [4B] OR reg → getRegId [1B]ga2_imm_checkmode-dependentmov [0xc(%r13)] @0x1242bb7 (AP) / @0x1242844 (reg)CONFIRMED
+0x384 or 1in3 value: AP-scalar [4B] OR reg getRegId [1B]ga2_imm_checkmode-dependentmov [0x38(%r13)] @0x1242bdb / @0x124291bCONFIRMED
+0x3D1packed dtype: low = sub_120E650(in2.D), high = 16*sub_120E650(out.D)ga2_valid_input_RMWmov %al,0x3d(%r13) @0x1242305CONFIRMED

A register coefficient with an offset register attached is rejected: "offset register cannot be applied to register immediate in ISA" (line 1584); a wrong kind throws "contains wrong argument type". 0x94 = BatchNormGradAccumulate2.

NOTE — self[76]. The scalar value encoder is the object at self+0x260 (= *((u64*)self+76)); its vtbl+0x20 slot encodes an AP-scalar immediate to a 4-byte wire value (the same IMM_SRC family the backprop setupImmediate uses). The register path bypasses it and calls bir::Register::getRegId for a 1-byte id.


visitInstBNBackpropdx apply, packed coeffs (op 0x64 + 0x65, two bundles)

Purpose

BNBackprop computes dx = (dy·N − ((x−m)·Bc + A)) · scale and emits two 64-byte bundles per BIR instruction, in wire order:

  1. bundle #1 = opcode 0x64 BatchNormParamLoad — the control bundle: the packed coefficient AP, total_elements (N), the coeff dtype, the base, and Mean as an IMM_SRC. Struct S2_BN.
  2. bundle #2 = opcode 0x65 BatchNormBackProp — the data bundle: the Ofmap(x) / OfmapGrad(dy) / Dst(dx) tensors. Struct S3S3D3_TT.

Order is byte-verified: in GENERATE mode the body stamps the 0x64 control header (movb $0x64,-0x112(%rbp) @0x12362b1) before the 0x65 data header (movb $0x65,-0x110(%rbp) @0x12360cd). Param-load first, back-prop second.

CORRECTION (BN-2 to D-J09) — D-J09 labelled 0x64 "BACK_PROP" and 0x65 "PARAM_LOAD"; the _OPCODE table and the pybind asserts invert this. 0x64 = BatchNormParamLoad (control: coeffs + N + Mean; asserts s_bn_param_load_*, grouped with s2_bn_dtype / s_valid_bn_param_load_imm0); 0x65 = BatchNormBackProp (data: x/dy/dx; assert s_batchnorm_backprop_opcode, grouped with s3d3_tt_*). Semantically: the param-load LOADS the coefficients/N/Mean, the back-prop APPLIES dx.

Bundle #1 — op 0x64 BatchNormParamLoad (S2_BN, base = r14)

Operands read: getArgument(0)=Ofmap (base only), getArgument(2)=Gradients (3-elt packed coeff [scale, A, Bc]), getArgument(3)=Mean.

OffWFieldpybind assertValue / sourceStore-siteConf
+0x001opcode 0x64s_bn_param_load_[opcode]setupHeader(100)movb $0x64 @0x12362b1CONFIRMED
+0x1016Gradients coeff AP TENSOR2Ds2_bn_src_element_cntassignAccess<TENSOR2D>(arg2)lea 0x10(%r14) @0x1236441CONFIRMED
+0x1C4total_elements N (f32)s_valid_bn_param_load_imm0*(float*)(Inst+0xF0)movss %xmm0,0x1c(%r14) @0x1236364CONFIRMED
+0x201coeff (arg2) dtypes2_bn_dtypesub_120E650([arg2+0x30])mov %al,0x20(%r14) @0x1236316CONFIRMED
+0x221BASE partitions2_bn / base*(*(arg0+0x50)+8)mov %al,0x22(%r14) @0x123634fCONFIRMED
+0x3016Mean (IMM_SRC)s_valid_bn_param_load_imm0setupImmediate(arg3) (sub_12051E0)lea 0x30(%r14) @0x12363e4CONFIRMED
+0x21, +0x23..2F, +0x02..0BRESERVED-0s2_bn_reserved_zer[o]zero-initINFERRED

total_elements at Inst+0xF0 is the dx N-multiplier; it is the BIR-JSON total_elements float. Mean is encoded by setupImmediate<NEURON_ISA_TPB_IMM_SRC> — a scalar immediate or an AP source, the same 4-function IMM_SRC encoder the matmul scalar path uses.

Bundle #2 — op 0x65 BatchNormBackProp (S3S3D3_TT, base = r14)

Operands: getArgument(0)=Ofmap(x), getArgument(1)=OfmapGrad(dy), getOutput(0)=Dst(dx). S3S3D3_TT = three rank-3 SBUF tensors laid out like a TensorTensor op.

OffWFieldpybind assertValue / sourceStore-siteConf
+0x001opcode 0x65s_batchnorm_backprop_opcodesetupHeader(101)movb $0x65 @0x12360cdCONFIRMED
+0x0C1packed dtype: low nibble = sub_120E650(Ofmap.D), high = 16*sub_120E650(OfmapGrad.D)s3d3_tt_dtypeRMW &0xF0/&0x0Fmov %al,0xc(%r14) @0x1236124..0x1236149CONFIRMED
+0x0D1OUTPUT (Dst) dtypes3d3_tt_dtypesub_120E650(out.D)mov %al,0xd(%r14) @0x123615dCONFIRMED
+0x0F1BASE partitions3d3_tt / base*(*(arg0+0x50)+8) (guard [+0x58]!=0)mov %al,0xf(%r14) @0x1236181CONFIRMED
+0x1016Ofmap AP TENSOR3Ds3d3_tt_valid_op[s]assignAccess<TENSOR3D>(arg0)lea 0x10(%r14) @0x1236179CONFIRMED
+0x2016OfmapGrad AP TENSOR3Ds3d3_tt_src_dst_[count]assignAccess<TENSOR3D>(arg1)lea 0x20(%r14) @0x123618fCONFIRMED
+0x3016Dst (dx) AP TENSOR3Ds3d3_tt_dst_elem[ent_cnt]assignAccess<TENSOR3D>(out)lea 0x30(%r14) @0x12361a7CONFIRMED

s3d3_tt_is_zero_[guard] is the empty / Welford-empty guard. Full operands across both bundles (the dx apply): AP0 Ofmap, AP1 OfmapGrad, AP2 Gradients(3-elt packed coeff), AP3 Mean, OUT Dst. The packed coeffs ride the param-load bundle's AP2; total_elements is its +0x1C float; Mean its +0x30 IMM_SRC.


visitInstBNBackprop2dx apply, split coeffs (op 0x8E + 0x65, two bundles)

Purpose

BNBackprop2 is structurally identical to BNBackprop (two bundles, control first), differing only in the control bundle. The dx numerics are identical; the difference is purely operand packing — the scalar scale is split into its own 1-element tensor and (A, Bc) into a 2-element tensor, encoded via two setupImmediate calls instead of one packed 3-element AP.

Order is byte-verified: control header movb $0x8E,-0x112(%rbp) @0x1236d11, then data header movb $0x65,-0x110(%rbp) @0x1236b2d.

CORRECTION (BN-3 to D-J09) — the BNBackprop2 control opcode is 0x8E (BatchNormParamLoad2), not 0x64. The disasm immediate at @0x1236d11 is 0x8E (movb $0x8e,-0x112(%rbp)), distinct from BNBackprop's 0x64 at @0x12362b1. CONFIRMED.

Bundle #1 — op 0x8E BatchNormParamLoad2 (S2_BNPL2, base = r14)

Five operands: getArgument(0)=Ofmap(base), getArgument(2)=GradientsA (1-elt [scale]), getArgument(3)=GradientsBc (2-elt [A, Bc]), getArgument(4)=Mean.

OffWFieldpybind assertValue / sourceStore-siteConf
+0x001opcode 0x8Es_s2_bnpl2_opcodesetupHeader(142)movb $0x8e @0x1236d11CONFIRMED
+0x1016GradientsA coeff AP TENSOR2Ds_valid_s2_bnpl2_imm_ptrassignAccess<TENSOR2D>(arg2)lea 0x10(%r14) (@0x123703f)CONFIRMED
+0x1C4total_elements N (f32)s_bnpl2_src_elem[ent]*(float*)(Inst+0xF0)movss %xmm0,0x1c(%r14) @0x1236db9CONFIRMED
+0x201GradientsA dtypes2_bnpl2 dtypesub_120E650([arg2+0x30])mov %al,0x20(%r14) @0x1236d7bCONFIRMED
+0x221BASE partitions2_bnpl2 / base*(*(arg0+0x50)+8)mov %al,0x22(%r14) @0x1236da4CONFIRMED
+0x3016split coeffs + Means_valid_s2_bnpl2_imm_ptrsetupImmediate(GradientsBc) THEN setupImmediate(Mean) — two sub_12051E0call 12051e0 @0x1236eda + @0x1236fc8CONFIRMED

The two IMM_SRC calls are the defining difference: call 12051e0 appears twice in the control body (@0x1236eda for GradientsBc, @0x1236fc8 for Mean), versus the single call in BNBackprop. Full operands: AP0 Ofmap, AP1 OfmapGrad, AP2 GradientsA(1-elt), AP3 GradientsBc(2-elt), AP4 Mean, OUT Dst.

NOTE — split rationale. The dx math is identical to BNBackprop; splitting scale into a 1-element tensor and (A,Bc) into a 2-element tensor with two IMM_SRC encodings (rather than one packed 3-element AP) lets the two coefficient groups be sourced separately. The exact "separate engines" rationale is INFERRED; the operand-packing difference is CONFIRMED.

Bundle #2 — op 0x65 BatchNormBackProp

Byte-identical to BNBackprop's data bundle: arg0/arg1/out as 3×TENSOR3D at +0x10/+0x20/+0x30; +0x0C dtype nibbles (Ofmap low / OfmapGrad high); +0x0D out dtype; +0x0F base. CONFIRMED — mov %al,0xc(%r14) @0x1236b8e/@0x1236ba9; [0xd(%r14)] @0x1236bbd; [0xf(%r14)] @0x1236be1; AP slots lea 0x10(%r14) @0x1236bd9 / lea 0x30(%r14) @0x1236c07.


The pybind assert roster → struct-offset join

The pybind module neuron_isa_tpb_pybind.cpython-310-*.so exposes each wire struct as a raw byte array; its .rodata assert table names the fields (16-char-truncated in the binding strings — full tails in brackets). Grouped by struct, every name below is confirmed present via strings on the module:

S4D2_BN  (BNStats / Aggregate) ──────────────────────────────────────────────
  s_batchnorm_stats_opcode        -> +0x00 opcode (Stats 0x61/0x82)
  s_batchnorm_aggregate_opcode    -> +0x00 opcode (Aggregate 0x62)
  s_bnstats2_src_element_cnt      -> +0x0C IN AP (TENSOR4D) element count
  s_bnstats2_dst_e[lement_cnt]    -> +0x30 OUT AP (== 6)
  s even_count must be finite     -> +0x28 f32 even-count (n_A)
  s odd_count must be finite      -> +0x2C f32 odd-count  (n_B)
  d2_bn_zero_counts               -> +0x28/+0x2C zero-when-N==0 guard
  d2_bn_stats_out_[ty]            -> +0x21 out dtype (Stats)
  d2_bn_agg_in_typ[e]             -> +0x20 in  dtype (Aggregate)
  d2_bn_agg_out_ty[pe]            -> +0x21 out dtype (Aggregate)
  s_bnstats_agg_src_element_cnt   -> in-AP NEPP (% 3 == 0)
  s_bnstats_agg_dst_element_cnt   -> out-AP NEPP (== 2)
  d2_bn_reserved_z[ero]           -> reserved-zero band

S3S3D1_BN / _BN2  (BNGradients) ─────────────────────────────────────────────
  s_s3s3d1_bn_opco[de]            -> +0x00 opcode 0x63
  s_s3s3d1_bn2_opc[ode]           -> +0x00 opcode 0x94
  s3d1_bn_channels                -> partition/channel count (+0x3F base + AP)
  bnga_valid_input_t[ype]         -> +0x3D in2(Mean) dtype (low nibble)
  ga_valid_output_[type]          -> +0x3E out dtype (A) / mode (B)
  ga_imm_check / ga_valid_immedia -> +0x0C in2 value / +0x38 in3 value
  ga_src_element_cnt_check        -> source (Ofmap/OfmapGrad) AP element count
  ga_dst_element_c[nt]            -> out Gradients AP NEPP (== 3, TENSOR1D)
  ga2_* (src_element_cnt_check, valid_output, valid_input_, imm_check)
                                  -> the op0x94 (BN2) parallels of ga_*

S2_BN / S2_BNPL2  (BNBackprop/2 control = param-load) ───────────────────────
  s_bn_param_load_[opcode]        -> +0x00 opcode 0x64 (S2_BN)
  s_s2_bnpl2_opcode               -> +0x00 opcode 0x8E (S2_BNPL2)
  s2_bn_dtype                     -> +0x20 coeff dtype
  s2_bn_reserved_zer[o]           -> +0x21/+0x23.. reserved band
  s_bnpl2_src_elem[ent]           -> +0x1C total_elements (BNPL2)
  s_valid_bn_param_load_imm0      -> +0x30 Mean IMM_SRC validity (BP1)
  s_valid_s2_bnpl2_imm_ptr        -> +0x30 split-coeff/Mean IMM_SRC validity (BP2)

S3S3D3_TT  (BNBackprop/2 data = back-prop) ──────────────────────────────────
  s_batchnorm_backprop_opcode     -> +0x00 opcode 0x65
  s3d3_tt_valid_op[s]             -> +0x10 Ofmap AP validity
  s3d3_tt_dtype                   -> +0x0C dtype-nibbles + +0x0D out dtype
  s3d3_tt_src_dst_[count]         -> +0x20 OfmapGrad / +0x30 Dst shape match
  s3d3_tt_src_element_cnt_check   -> source AP element count
  s3d3_tt_dst_elem[ent_cnt]       -> Dst AP element count
  s3d3_tt_is_zero_[guard]         -> empty/zero-count guard

GOTCHA — sub-byte widths are not pinned. The encoder writes whole bytes; the dtype tags at +0x20/+0x21/+0x0C/+0x3D may internally be 5-bit-tag-plus-3-reserved, but the binary stores the whole byte and the pybind exposes the structs as raw byte arrays with no per-field bit metadata. Only the BYTE field map and the explicit &0xF/&0xF0 nibble splits (the dtype-nibble and mode-nibble stores) are CONFIRMED to sub-byte granularity. Do not invent a bit layout below the byte.


Confidence ledger

ClaimConfidenceEvidence
Five CoreV2GenImpl::visitInstBN* bodies at the cited addressesCONFIRMEDnm -DC libwalrus.so (0x123ab00/0x123c6a0/0x1241b10/0x1235e80/0x12368e0)
BNStats mixed-width 4D-in +0x0C / 2D-out +0x30CONFIRMEDlea 0xc(%r13) @0x123ade5 + lea 0x30(%r13)+assignAccess<TENSOR2D> @0x123adf6
BNStats count seeds ceil/floor(N/2) @ +0x28/+0x2CCONFIRMEDadd$1;shr$1;cvtsi2ss;movss 0x28 / shr$1;…0x2c @0x123ad39..75
Opcode 0x61/0x82 folds transposeCONFIRMEDcmpb$1,0xf0;sbb;and$0xffffffdf;sub$0x7e @0x123ab5f
Aggregate = BNStats minus count seeds (op 0x62)CONFIRMEDmovb $0x62 @0x123c98e; +0x20/0x21/0x22/lea 0xc/0x30 on r12
Gradients arch gate <=490x63, else 0x94; core_v5 FATAL on 0x63CONFIRMEDcmpl $0x31,0x258(%r12) @0x1242149; NeuronAssertion 2101
Gradients template +0x3C dtype-nibble / +0x3F base / 3 AP slotsCONFIRMEDmov 0x3c(%rbx)×2 / mov 0x3f(%rbx) / lea 0x10/0x20/0x30(%rbx) @0x133e59a..f7
BNBackprop two bundles 0x640x65, control firstCONFIRMEDmovb $0x64 @0x12362b1 before movb $0x65 @0x12360cd
BNBackprop2 control opcode 0x8E (not 0x64)CONFIRMEDmovb $0x8e,-0x112(%rbp) @0x1236d11
BNBackprop2 split coeffs = two setupImmediateCONFIRMEDcall 12051e0 @0x1236eda + @0x1236fc8
S3S3D3_TT data bundle byte-identical across BP1/BP2CONFIRMED+0x0C/0x0D/0x0F + 3×TENSOR3D stores in both bodies
Every pybind assert name (d2_bn_*/s2_bn*/s3d3_tt_*/ga*/bnstats2_*)STRONGstrings neuron_isa_tpb_pybind.cpython-310-*.so
dtype LUT sub_120E650byte_1DF5760STRONGshared with the matmul dtype path; not re-traced this pass
Reserved-zero bandsINFERREDbulk movups zero-fill + *_reserved_zero asserts; no direct store
Engine = PoolSTRONGopcode-table grouping + partition-reduce semantics; Inst+0x90 is the scheduler's field

Cross-References