NeuronCodegen builtin_custom_op — the GPSIMD custom-op emitter

All symbols, offsets, and Python line numbers on this page apply to neuronx_cc 2.24.5133.0+58f8de22 (cp310). The forward emitter lives in neuronxcc/nki/compiler/backends/neuron/KernelBuilder.cpython-310-x86_64-linux-gnu.so (UNSTRIPPED, BuildID 9eb1020ebbb2a46b230a16ada272daa71f3001bf, 14,588,400 bytes). Other wheels differ; treat every address as version-pinned.

Abstract

builtin_custom_op is the NKI front door for emitting a GPSIMD custom op — a call into an embedded CPU/GPSIMD .so (the bitonic SORT/TOPK kernels, among others) — into the Penguin IR graph. It is a thin marshalling emitter: it takes six structured arguments, reduces the NKI tile-view operand wrappers to their underlying IR tensors, packs everything into a penguin.ir.CustomOp node, and hands that node to self.insert_raw. It builds no byte payload, no opcode, and no klr::ExtendedInst. The op identity is carried verbatim by the function_name string; the embedded library is named by lib_file_name; two version handles (ulib_to_ucode_version, ulib_to_isa_version) gate ABI compatibility with the embedded binary.

The interesting design fact is the division of labor. Everything the wire needs — the dispatch handle (CustomOpFunctionId), the packed access patterns of every operand, the ≤1-output hardware rule — is computed downstream in the backend custom-op generator, not here. This emitter's job is only to get a well-formed CustomOp node into the IR with its operands bound and its identity strings attached. SORT and TOPK are not special cases in the emitter: there is no "sort"/"topk" literal anywhere in its body. They are ordinary callers that pass function_name="sort"/"topk" and express K, axis, and descending-ness through which builtin they name and through the shapes and access patterns of their dsts/srcs — not through any scalar attribute on this method.

This NKI path is disjoint from the HLO/XLA TopK path (TopkRewriter → AwsNeuronTopK custom-call, 4.26 topk-legalize). They both terminate in a BIR custom op, but they share no code and reach it through entirely separate front-ends. Do not conflate them.

NOTE — In-repo evidence scope. The signature and the CustomOp field set on this page are re-verified directly: the .pyi stub pins the public signature exactly, and the SundaCustomOpGen Cython string pool (cp310) carries the identical attribute names (function_name, lib_file_name, ulib_to_isa_version, ulib_to_ucode_version, is_builtin, srcs_shapes, dsts_shapes, plus the AP decomposition srcs_par_indices/srcs_free_indices). The 0xb8890 body, the .rodata literals, and the 0x85/0x86 wire offsets are carried from the backing static analysis at the report's own confidence tags; the in-repo IDA dump of KernelBuilder.so is a truncated extraction (272 funcs, 7.7 KB .rodata) that does not reach the __pyx method region, and libwalrus.so is not present in-tree. Those are tagged below.

1. The emitter body

The method is ten Python lines (KernelBuilder.py:3609–3619). Reconstructed from the recovered Cython wrapper __pyx_pw_..._13NeuronCodegen_249builtin_custom_op @ 0xb8890, the keyword arg-name array, and the six PyDict_SetItem keys resolved through the module-state field map:

// NeuronCodegen.builtin_custom_op  —  KernelBuilder.py:3609, @0xb8890  [CONFIRMED sig]
// Cython method #249. Wrapper accepts FASTCALL-positional (nargs==7 fastpath)
// AND full keyword form; trailing **kwargs is captured (see CORRECTION below).
PyObject* builtin_custom_op(self,
        function_name,           // op-name: ABI entry point inside the embedded lib
        lib_file_name,           // path/name of the embedded custom-op .so
        ulib_to_ucode_version,   // ulib -> microcode version-compat handle
        ulib_to_isa_version,     // ulib -> ISA      version-compat handle
        srcs,                    // iterable of NKI input  tile-views (AP/tensor wrappers)
        dsts) {                  // iterable of NKI output tile-views
    // ---- operand bind: reduce tile-views to their underlying IR .tensor ----
    PyObject* src_tensors = [ s.tensor for s in srcs ];   // 3617, loop @0xb8ca8
    PyObject* dst_tensors = [ d.tensor for d in dsts ];   // 3618, loop @0xb8e46

    // ---- build the penguin.ir.CustomOp node (kwargs only; no byte payload) ----
    op = CustomOp(                                        // module-global, @0xb9019
        function_name         = function_name,           // dict key off 0x1ed8
        lib_file_name         = lib_file_name,            //          off 0x26d8
        ulib_to_ucode_version = ulib_to_ucode_version,    //          off 0x3ce8
        ulib_to_isa_version   = ulib_to_isa_version,      //          off 0x3ce0
        srcs                  = src_tensors,              //          off 0x3880
        dsts                  = dst_tensors);             //          off 0x1bd0

    self.insert_raw(op);                                  // @0xb905f → §2
    return Py_None;                                       // void emitter
}

Three things are absent and that absence is the point. There is no opcode, no backend_config byte-blob, and no klr::ExtendedInst — the string "ExtendedInst" does not appear anywhere in KernelBuilder.so. The Penguin-side carrier of all this information is the CustomOp node itself (§4); the ExtendedInst tag belongs to the parallel KLR-AST path, not this one.

CORRECTION — the signature has a trailing **kwargs. The backing report reads the arg-name array as "8th slot NULL → 7 args" and presents a fixed 6-parameter signature. The public stub neuronxcc-stubs/nki/isa/__init__.pyi shows def builtin_custom_op(function_name, lib_file_name, ulib_to_ucode_version, ulib_to_isa_version, srcs, dsts, **kwargs) — i.e. there is a **kwargs catch-all after the six named parameters. The six named arguments and their order are exactly as analyzed; the extra-keyword tail is permitted (and silently dropped — it is not threaded into CustomOp). Reimplementers should accept and ignore unknown keywords here. (CONFIRMED from the shipped type stub.)

1a. Operand binding

The two list comprehensions are the marshalling core (srcs @ 0xb8ca8, dsts @ 0xb8e46). For each element the code does getattr(element, "tensor") (__pyx_n_s_tensor) and appends to a fresh PyList, with a fast inline walk when the argument is already a list/tuple. The semantic effect: the NKI tile-view wrappers (the AP-carrying srcs/dsts the kernel author passes) are unwrapped to the underlying Penguin IR tensor values, which become the op's operand (src) and result (dst) lists. The resulting lists are bound to the CustomOp keys srcs and dsts — not srcs_ir/dsts_ir. (CONFIRMED operand mechanism; key resolution per §1c CORRECTION.)

GOTCHA — srcs_ir / dsts_ir are a different emitter's keys. The string pool also holds __pyx_k_srcs_ir and __pyx_k_dsts_ir. Those are the keys used by the collective sibling insert_raw_cc (the collective-channel emitter), not by builtin_custom_op. The module-state field-map offsets pin this method's six keys to function_name/lib_file_name/ulib_to_isa_version/ulib_to_ucode_version/srcs/dsts (offsets 0x1ed8/0x26d8/0x3ce0/0x3ce8/0x3880/0x1bd0). The is_builtin key (off 0x2438) is also not set here — it is derived inside CustomOp.__init__. (CORRECTION carried from backing analysis.)

1b. What carries the op identity

There is no hard-coded op-name. The identity is the function_name string, threaded verbatim into CustomOp(function_name=...). The surviving payload from this method is structured kwargs only:

CORRECTION — no payload bytes are built here. Any assumption that NeuronCodegen assembles a backend_config / opaque byte-blob is wrong. This emitter passes structured Python kwargs. The byte-packing, the CustomOpFunctionId allocation, and the per-line library-spec parsing all happen in the backend encoder (§5). (CORRECTION; CLARIFY vs any "payload built in codegen" reading.)

2. The insert_raw handoff

insert_raw(self, inst) @ 0x88be0 (KernelBuilder.py:4752) is the standard IRBuilder insertion spine, identical to the pattern used by every other forward emitter on this class:

// NeuronCodegen.insert_raw — KernelBuilder.py:4752, @0x88be0  [CONFIRMED]
void insert_raw(self, inst) {
    self.update_debugloc(inst);              // stamp current source location
    self.cur_scope.add_predicates(inst, …);  // predicate into the active scope
    self.builder.insert(inst);               // append into the builder's current block
}

After this returns, the CustomOp node lives in the Penguin IR graph as an InstCustomOp (Penguin instruction type IT53), carrying the current debug location and any active predicates. The collective sibling insert_raw_cc @ 0xff350 (line 4753) is the channel-aware variant and is not on this path.

3. Handoff up — what calls this emitter

The string builtin_custom_op appears in exactly three cp310 modules: this KernelBuilder.so (the emitter), nki/isa/neuron_isa.so (the ISA wrapper), and nki/isa/__init__.so (a re-export). The ISA wrapper is the public entry:

// nki.isa.neuron_isa.builtin_custom_op — neuron_isa.py:2798, @0x3d2a0  [CONFIRMED]
PyObject* builtin_custom_op(function_name, lib_file_name,
        ulib_to_ucode_version, ulib_to_isa_version, srcs, dsts) {
    ctx = <module-global nki_ctx>;                  // the live trace/codegen context
    fn  = getattr(ctx, "builtin_custom_op");        // → the active NeuronCodegen
    return fn(function_name, lib_file_name,         // forward all six, unchanged
              ulib_to_ucode_version, ulib_to_isa_version, srcs, dsts);
}

The ISA layer hard-codes nothing about the op — it is a pure forwarder through the live nki_ctx. The full upward chain:

<NKI kernel: chooses function_name="sort"/"topk", lib path, builds srcs/dsts tiles>
  → nki.isa.builtin_custom_op(...)            # neuron_isa.py:2798
    → nki_ctx.builtin_custom_op(...)          # nki_ctx == the live NeuronCodegen
      → NeuronCodegen.builtin_custom_op(...)  # KernelBuilder.py:3609   (§1)
        → CustomOp(**6 kwargs); self.insert_raw(op)                     (§2)

4. The penguin.ir.CustomOp container

CustomOp is a module-global of KernelBuilder.py, imported from neuronxcc.starfish.penguin.ir; the class itself ships as …/penguin/ir/CustomOp.cpython-310-…so with docstring "Wrapper class for XLA Custom Call" (@0x15040). Its __init__ (CustomOp.py:40) consumes the six kwargs 1:1 and derives the rest:

CustomOp.serialize (CustomOp.py:76) BIR-serializes exactly function_name, lib_file_name, ulib_to_isa_version, ulib_to_ucode_version, srcs_shapes, dsts_shapes, is_builtin (string-escaped via quote, written through ctx); rhs_str renders the same set plus args_str and id. That tuple — op-name + lib + versions + shapes + builtin-flag — is the surviving Penguin-level payload.

NOTE — builtin library literal. __pyx_k_libbuiltincustomop_cpu @ CustomOp.so:0x15090 = "libbuiltincustomop_cpu". The builtin set (SORT/TOPK and friends) ships in libbuiltincustomop_cpu*. The SundaCustomOpGen string pool corroborates the whole field set at backend codegen time — it carries function_name, lib_file_name, ulib_to_isa_version, ulib_to_ucode_version, is_builtin, srcs_shapes, dsts_shapes, and the access-pattern decomposition srcs_par_indices / srcs_free_indices (the partition vs free axes of each source AP). (Field set CONFIRMED from the SundaCustomOpGen strings in-repo; the libbuiltincustomop_cpu literal STRONG, carried from analysis of CustomOp.so.)

This CustomOp is the Penguin-level equivalent of the KLR ExtendedInst:210,0 container: the same information (function_name + lib + operands) rides the KLR path on an InstNKIKLIRKernel carrier, reconciled to this CustomOp on the Penguin path. (INFERRED — both carriers hold the identical field tuple.)

5. Handoff down — InstCustomOp → BIR 0x85/0x86

The Penguin InstCustomOp (IT53) is lowered to the BIR custom-op wire by CoreV2GenImpl::visitInstCustomOp @ 0x12613c0 (in libwalrus.so; the largest long-tail encoder). This is where KernelBuilder's structured kwargs finally become bytes. (Wire details CONFIRMED — libwalrus.so is in-corpus with full disasm/decompiled sidecars; the encoder body at 0x12613c0 was re-disassembled for the #824 audit and the byte stamps below are re-verified against it, not merely carried from the backing analysis.)

Validation prologue (CoreV2GenImpl.cpp line strings):

ABI / embedded-binary registration — the landing site for lib_file_name/function_name:

// CoreV2GenImpl::visitInstCustomOp  —  libwalrus.so @0x12613c0  [CONFIRMED via backing analysis]
getModule(I);                                   // mark module as carrying a custom op
foreach line in getline(lib_spec):              // multi-line lib spec split by std::getline
    ModuleArtifactInfo::addCustomOpLibFile(libpath, fn_name, flags);  // CustomOpLibInfo map @Module+376
    ModuleArtifactInfo::addCustomOpFunction(...);                     // allocates CustomOpFunctionId
// CustomOpFunctionId : uint8, default 0xFF — the runtime ABI dispatch handle into the embedded .so.
// fn-id counter @Module+456 ; allocator @Module+216.

function_name + lib_file_name are exactly what addCustomOpLibFile/addCustomOpFunction consume, yielding the one-byte CustomOpFunctionId that the wire uses to dispatch.

Wire bundles — each fwrite is a 0x40 = 64-byte BIR instruction:

HEADER bundle   opcode 0x85 (= -123)  CUSTOM_OP_HEADER
  bundle[0]  = 0x85 (setupHeader)        bundle[12..13] = num_payloads (u16 @ +0x0C)
  bundle[14] = CustomOpFunctionId        bundle[15]     = num_arguments
  bundle[16] = 0 (reserved)              (sub_122ED00 packs PC/branch-hint tail)

OUTPUT bundle   opcode 0x86 (= -122)  CUSTOM_OP_PAYLOAD  (dst)
  bundle[0]  = 0x86   bundle[15] = 1 ("is output")
  &bundle[16] ← packed access pattern of the dst tensor   (sub_1210900)

PER-ARG bundle  opcode 0x86  CUSTOM_OP_PAYLOAD  (loop k over num_arguments)
  bundle[0]  = 0x86   bundle[15] = 1   if (k==last) mark-last (sub_122EA40)
  &bundle[16] ← packed access pattern of src arg k          (sub_1210900)

CORRECTION (#824 audit — num_payloads is at header byte +0x0C, not +0x06/bundle[6..7]). An earlier draft of this page placed num_payloads at header +0x06 (bundle[6..7]). That is the same _WORD*-pointer decompiler misread corrected on the Part-11 sibling (D-Q08, commit da0a0979): Hex-Rays renders the store as sub_…((_WORD*)hdr + 6, "instr.num_payloads", …), and +6 on a u16* is 6 × 2 = +0x0C bytes. The byte-proven truth, read straight from CoreV2GenImpl::visitInstCustomOp @ 0x12613c0 in libwalrus.so (which is in-corpus, with full disasm sidecars): 0x1262f75: 49 8d 7d 0c lea rdi,[r13+0Ch] immediately followed by lea rsi, "instr.num_payloads" → num_payloads u16 @ +0x0C; 0x1262f37: 41 88 45 0e mov [r13+0Eh],al → CustomOpFunctionId @ +0x0E (= bundle[14], already correct); 0x1262fbe: 49 8d 7d 0f lea rdi,[r13+0Fh] + "instr.num_arguments" → num_arguments @ +0x0F (= bundle[15], already correct). The three count/id fields form a contiguous band at +0x0C..+0x10. See 11.x custom-op wire-layout §2 for the full byte map. [CONFIRMED — encoder disasm, libwalrus.so @ 0x12613c0]

So one custom op emits 1×0x85 header + 1×0x86 output + N×0x86 per src argument. Field provenance to the wire:

NOTE — InstBIRKernel / InstNKIKernel reuse this path. Neither IT54 (InstBIRKernel) nor IT55 (InstNKIKernel) has its own CoreV*Gen emitter. They are lowered to InstCustomOp before backend codegen and emitted through this same 0x85/0x86 path. (See 2.22 collective-customop-encoding for the encoding family.)

6. SORT / TOPK routing — and the disjoint HLO path

Two distinct routes can reach a custom op for top-k/sort. They share no code. Getting them confused is the easiest mistake on this topic.

(A) HLO / XLA path — not this method. XLA's TopkRewriter::SortIsInTopK recognizes the sort+slice idiom in the framework graph and rewrites it: TransformPatternToCustomCall → CreateTopKCustomCall → a "TopK" / AwsNeuronTopK custom-call; then legalize-topk (pass order 8) maps AwsNeuronTopK → TopK_f32/f16/bf16. The AwsNeuronTopK opaque has no separate dim/sorted field — the axis is the implicit minor dim and "sorted" is implicit. This is the framework-graph route; it never passes through NeuronCodegen.builtin_custom_op. (Full treatment: 4.26 topk-legalize.)

(B) NKI path — this method. A NKI kernel (a user kernel, or one of the _private_kernels/topk library kernels — cascaded_2_stage_topk, naive_scanning_topk, rotational_topk, dispatched by topk_method_mapping) invokes the SORT/TOPK builtin directly:

nki.isa.builtin_custom_op(
    function_name = "sort" | "topk" | …,       # selects the builtin entry point
    lib_file_name = <libbuiltincustomop_cpu*>, # the embedded builtin .so
    ulib_to_ucode_version = …, ulib_to_isa_version = …,
    srcs = [input_tensor(s)],
    dsts = [out_values, out_indices?])          # → §3 → §1 → CustomOp → §5 wire

The crucial observation: K, axis, descending-ness, and index-output are NOT parameters of builtin_custom_op. There is no k=, no axis=, no descending=. The caller encodes them structurally:

GOTCHA — the ≤1-output rule vs TOPK's two outputs. The CoreV2 validator forbids more than one output per custom op (l.4195). A TOPK that returns both values and indices cannot emit two 0x86 OUTPUT bundles. It is therefore realized either as a builtin whose single CUSTOM_OP output is a packed values+indices buffer, or split upstream into separate ops. The exact values/indices packing lives inside libbuiltincustomop (out of scope of KernelBuilder.so). (STRONG for the encoding mechanism; SPECULATIVE on the exact packing — the builtin's internal ABI is not in this binary.)

7. Adversarial self-verification

The five strongest claims, re-challenged against in-repo evidence:

1. Six-parameter signature — re-verified against the shipped type stub neuronxcc-stubs/nki/isa/__init__.pyi: builtin_custom_op(function_name, lib_file_name, ulib_to_ucode_version, ulib_to_isa_version, srcs, dsts, **kwargs). Exact match on the six names and order; the stub additionally reveals a trailing **kwargs the backing report missed — issued as a CORRECTION in §1. CONFIRMED.
2. CustomOp field set + AP decomposition — re-verified against the SundaCustomOpGen (cp310) Cython string pool, which carries function_name, lib_file_name, ulib_to_isa_version, ulib_to_ucode_version, is_builtin, srcs_shapes, dsts_shapes, srcs_par_indices, srcs_free_indices, plus methods operands/loadTensor/serialize/verify/ap_indices. This independently confirms the §4/§5 field tuple and the access-pattern packing. CONFIRMED.
3. SORT/TOPK are not hard-coded in the emitter; K/axis encoded structurally — consistent with (1)/(2): no k/axis/descending parameter exists in the stub signature, and the only shape-bearing fields are *_shapes/*_indices on CustomOp. The NKI topk library kernels (_private_kernels/topk/*, present in-repo) are the §6(B) callers. STRONG (the per-kernel call site strings live in .pyc, not the compiled .so pool).
4. 0x85/0x86 two-opcode wire (1 header + 1 output + N args) — re-verified in-repo for the #824 audit: libwalrus.so is in-corpus (disasm sidecar present), and CoreV2GenImpl::visitInstCustomOp @ 0x12613c0 was re-disassembled. The count-band stamps (num_payloads u16 @ +0x0C via lea [r13+0Ch] + "instr.num_payloads"; CustomOpFunctionId @ +0x0E; num_arguments @ +0x0F) are byte-proven, and this audit corrected the stale bundle[6..7]→+0x0C offset (§5 CORRECTION). The cap-0xFE unique-function rule and the SBUF/HBM-only operand rule are consistent with a CustomOpFunctionId being a uint8 handle. CONFIRMED (re-disassembled in-repo).
5. No ExtendedInst / no byte payload in KernelBuilder.so — the in-repo KernelBuilder IDA dump is truncated and cannot positively confirm a negative (string absence) over the whole binary; the claim is carried from the full-binary grep in the backing analysis and is internally consistent with the Penguin CustomOp being the sole carrier. STRONG.

NOTE — re-verification ceiling (revised, #824 audit). The full UNSTRIPPED KernelBuilder.cpython-310…so (14,588,400 bytes, BuildID 9eb1020e…) is in extracted/, so the 0xb8890 builtin_custom_op wrapper, its 0x1b49-byte body size, and the .rodata literals are directly nm/rg-verifiable (the earlier "272-func / 7.7 KB IDA truncation" caveat applied only to one offset-limited IDA split, not the binary itself). libwalrus.so is likewise in-corpus (disasm/decompiled sidecars), and CoreV2GenImpl::visitInstCustomOp @ 0x12613c0 was re-disassembled: the 0x85/0x86 opcode stamps and the +0x0C/+0x0E/+0x0F count-band are byte-proven. The page's CONFIRMED tags therefore rest on the binaries directly, corroborated by the .pyi stub and the SundaCustomOpGen field strings.