TopK Legalization

All addresses on this page are virtual addresses (VMA) for neuronx-cc 2.24.5133.0+58f8de22 (hlo-opt, cp310); resolve via objdump --start-address or the VMA-keyed disasm/ sidecars. VA ≠ raw file offset: .text file_off = VA − 0x201000, .rodata file_off = VA − 0x200000 (section headers). Other builds will differ.

Abstract

legalize-topk is registration-order pass 8 in the hlo-opt --passes table, class xla::LegalizeTopK. Despite the name, it is not a sort+slice idiom recognizer. It is a narrow dtype specializer: it scans the entry computation for kCustomCall instructions whose custom_call_target() is exactly AwsNeuronTopK, gates on the element type (must be F16, BF16, or F32), parses the integer k out of the backend_config opaque string, builds an S64 index/iota constant over the operand's minor dimension, and rewrites the op into a per-dtype custom-call named TopK_f32, TopK_f16, or TopK_bf16 with operands {original_input, iota}. The string AwsNeuronTopK (0x25f115) is referenced by exactly one function in the binary — LegalizeTopK::Run (0x1f01260) — so the producer of AwsNeuronTopK lives upstream (driver-side Frontend / HLOToTensorizer) and is out of this binary.

The familiar reference frame is upstream XLA's TopkRewriter/HloTopKInstruction/TopkDecomposer, all statically linked into this same ELF but on a different code path: TopkRewriter::SortIsInTopK (0x29e24b0) recognizes the sort+slice idiom and emits a custom-call named plain "TopK" via CreateTopKCustomCall (0x29e4930). LegalizeTopK consumes the Neuron-specific AwsNeuronTopK form, never the stock "TopK" one. Think of this pass as the last HLO-level hop before the legalized TopK_* op enters NKI/KLR codegen.

The legalized TopK_* op resolves downstream to an internal NKI kernel (the router_topk / topk_reduce family), which lowers through KlirToBirCodegen into the BIR DVE reduce/index instructions (InstructionType 88–91: Max, MaxIndex, MatchReplace, MaxIndexAndMatchReplace), realized by the DVE Max-8 / Find-Index-8 / Match-Value-Load microcode primitives. That crosswalk is documented in §4.

For reimplementation, the contract is:

• The match predicate: kCustomCall + target == "AwsNeuronTopK" + element type ∈ {F16, BF16, F32}, scanned over the entry computation only, in post-order.
• The rewrite: append an S64 iota constant, look up the dtype-specialized target name, emit CreateCustomCall(shape, {input, iota}, name, /*opaque=*/"", /*api=*/API_VERSION_ORIGINAL), then ReplaceInstruction.
• The backend_config schema: a bare base-10 integer = k (no dim/sorted fields), validated to INT32 range; consumed for validation only and dropped (the new op carries an empty config).

LegalizeTopK::Run

Purpose

Consume an already-emitted AwsNeuronTopK custom-call and specialize it by dtype into one of three concretely-named TopK_<dtype> custom-calls, attaching the S64 index iota that the downstream NKI kernel consumes as its lane-index source. The pass performs no comparator construction, no CreateSort, and no reduce-window synthesis — it is purely a recognize-and-relabel rewrite over the entry computation.

Entry Point

xla::LegalizeTopK::Run  (0x1f01260, 2879 B)          ── pass body (HloPassInterface::Run)
  ├─ HloComputation::MakeInstructionPostOrder  (0x1f012c2 → 0x9634ab0)  ── scan order
  ├─ __dynamic_cast                            (0x1f0131c)              ── → HloCustomCallInstruction*
  ├─ std::string::compare("AwsNeuronTopK")     (0x1f01334)              ── target gate
  ├─ Shape::array_state / LiteralUtil::CreateR0<long>  (0x1f01387/0x1f01399)  ── iota length
  ├─ ShapeUtil::MakeValidatedShape(S64, {n})   (0x1f014d7)              ── index shape
  ├─ MutableLiteralBase::PopulateR1<long>      (0x1f0154b)              ── fill iota
  ├─ strtol(raw, &end, 10)                     (0x1f0159c)              ── parse k
  ├─ HloInstruction::CreateConstant            (0x1f01738)              ── materialize iota
  ├─ HloInstruction::CreateCustomCall          (0x1f01901)             ── emit TopK_<dtype>
  └─ HloComputation::ReplaceInstruction        (0x1f019a2 → 0x963fe50) ── splice in

Algorithm

Reconstructed instruction-by-instruction from the 0x1f01260 disasm, with anchors in «addr». The callees and immediates are read directly; the loop and skip control flow is reconstructed from jump targets, since this binary carries no Hex-Rays output.

StatusOr<bool> LegalizeTopK::Run(HloModule* m, const flat_hash_set<string_view>& threads) {
    // topkNames: a lazily-built, process-lifetime DenseMap<PrimitiveType,std::string>
    //   F16(0x0A)->"TopK_f16", F32(0x0B)->"TopK_f32", BF16(0x10)->"TopK_bf16"
    //   built under __cxa_guard in static-init region; see "topkNames table" below.

    HloComputation* entry = m->entry_computation();        // «[rdx+38h]»
    CHECK(entry != nullptr);                               // "nullptr != entry_computation_" hlo_module.h:164

    for (HloInstruction* inst : entry->MakeInstructionPostOrder()) {   // «0x1f012c2»
        if (inst->opcode() != kCustomCall) continue;       // opcode byte gate
        auto* cc = dynamic_cast<HloCustomCallInstruction*>(inst);      // «0x1f0131c»
        if (cc == nullptr) continue;
        if (cc->custom_call_target().compare("AwsNeuronTopK") != 0)    // «0x1f01321 / 0x1f01334»
            continue;

        // --- dtype gate: operand0 element type must be F16 / F32 / BF16 ---
        PrimitiveType et = cc->operand(0)->shape().element_type();     // ecx «0x1f01356»
        if ((uint32_t)(et - 0x0A) > 1 && et != 0x10)        // «lea eax,[rcx-0Ah]; jbe 0x1f01364; cmp ecx,10h 0x1f01366»
            continue;                                       //   accepts {0x0A=F16, 0x0B=F32, 0x10=BF16}

        // --- build S64 iota over the operand's minor (last) dimension ---
        const Shape& s = cc->operand(0)->shape();
        int64 n = s.array_state(); /* minor-dim element count */        // «array_state 0x1f01387/0x1f013ba»
        Shape idxShape = ShapeUtil::MakeValidatedShape(/*S64=*/5, {n}); // «mov esi,5 0x1f014c4; 0x1f014d7»
        Literal idxLit(idxShape);
        idxLit.PopulateR1<long>(/* 0..n-1 index vector */);            // «0x1f0154b»

        // --- parse k from backend_config (a bare base-10 integer) ---
        std::string raw = cc->backend_config().raw_string();
        errno = 0; char* end;
        long k = strtol(raw.c_str(), &end, 10);             // «0x1f0159c»
        if (end == raw.c_str())  throw invalid_argument("stoi");        // «0x1f01cdc»
        if (errno == ERANGE)     throw out_of_range("stoi");           // «[rbx]==0x22 0x1f01cd2»
        if (k < INT32_MIN || k > INT32_MAX) throw out_of_range("stoi"); // «0x1f015c5/0x1f015cf vs 0x7FFFFFFF»
        (void)k;   // validated, then DISCARDED — the new op carries NO backend_config

        // --- materialize iota and assemble {input, iota} ---
        HloInstruction* iota = entry->AddInstruction(
                                   HloInstruction::CreateConstant(std::move(idxLit)), "");  // «0x1f01738»
        HloInstruction* operands[2] = { cc->operand(0), iota };

        // --- look up dtype-specialized name and emit ---
        const std::string& newTarget = topkNames[et];       // "TopK_f32" / "TopK_f16" / "TopK_bf16"
        auto newCC = HloInstruction::CreateCustomCall(
                         cc->shape(),                        // output shape = original AwsNeuronTopK shape
                         absl::Span<HloInstruction* const>(operands, 2),
                         /*custom_call_target=*/ newTarget,
                         /*opaque(backend_config)=*/ "",     // EMPTY
                         /*api=*/ API_VERSION_ORIGINAL /*1*/);          // «push 1 0x1f018e3; 0x1f01901»
        HloInstruction* nn = entry->AddInstruction(std::move(newCC), "");
        CHECK_OK(entry->ReplaceInstruction(inst, nn));      // «0x1f019a2»
        // changed = true;
    }
    return /*StatusOr<bool>*/ changed;
}

GOTCHA — k is parsed and range-checked but then thrown away. The emitted TopK_<dtype> op has an empty backend_config (opaque=""). Downstream identity of the op is (target name, operand shapes, output shape) only. A reimplementation that tries to propagate k onto the new op's config will diverge — the value of k is recovered downstream from the output shape's minor dimension, not from a config field.

QUIRK — the second operand is an S64 iota of length = the input's minor dimension n (0..n-1), not the output's k. The NKI kernel uses this full-length index vector as the lane-index source for its argmax/find_index8 passes; it is sliced to k internally, not here.

NOTE — only entry_computation() is scanned (single [rdx+38h] read + null CHECK), in instruction post-order. This contrasts with passes such as NeuronHloInstCombine that iterate module.computations(). A TopK buried inside a fusion/while body is invisible to this pass.

Function Map

The topkNames table

Built once, lazily, under __cxa_guard in the static-init region; destroyed via __cxa_atexit. The three string payloads are stored as 8-byte little-endian inline immediates (SSO short strings) — decoded directly from the mov rax, imm64 instructions:

The same table is the single source of truth for both the dtype gate and the emitted name, so they cannot disagree: the jbe accepts et-0x0A ≤ 1 (F16/F32), the cmp accepts 0x10 (BF16), and the map is keyed on those exact ordinals.

NOTE — XLA PrimitiveType ordinals here are F16=10, F32=11, BF16=16 — these match the DenseMap keys exactly. The gate is two comparisons (lea eax,[rcx-0Ah] then jbe, plus cmp ecx,10h), not a switch; a fourth float type (e.g. FP8) is silently skipped, not legalized.

backend_config Schema

On the incoming AwsNeuronTopK, backend_config is a bare base-10 integer string equal to k. There is no structured proto, no dim field, and no sorted field. The reduction axis is implicit — the operand's minor (last) dimension, over which the iota is built. "Sorted" is implicit-true (top-k via the DVE Max/MaxIndex family is inherently ordered).

The wider Neuron-TopK constraint set surfaces through the hilo diagnostic catalog (lookup_cause 0x759d7d0 / lookup_resolution 0x759fc20), even though the enforcing checks for operand count and integer-type rejection are upstream of this pass:

Net contract: exactly 1 operand; element type F16/BF16/F32 (no 32/64-bit int); k ≤ INT32_MAX.

GOTCHA — two different TopK representations coexist in this ELF. The framed xla::HloTopKInstruction carries explicit k= and largest= attributes (printed by PrintExtraAttributesImpl 0x9686460), but that is the stock XLA form and never reaches this pass; the AwsNeuronTopK form consumed here encodes k as a bare integer string with no largest field at all. Note also that formatErrorMessage<PrimitiveType,PrimitiveType> (0x1eeb8d0) belongs to ArgMax — its sole caller is LegalizeAwsNeuronArgMax::Run — so it is not part of the TopK diagnostic path despite sitting beside it.

HLO-TopK → BIR-DVE Crosswalk

The legalized TopK_<dtype> custom-call is realized by an internal NKI kernel (registered in _INTERNAL_KERNEL_REGISTRY), which lowers through KLR/KlirToBirCodegen into the BIR DVE reduce/index family.

 upstream XLA (same ELF, other path)        hlo-opt (this pass)            NKI ISA (nisa.*)          BIR InstructionType (88-91)        DVE microcode op
 ───────────────────────────────────        ────────────────────           ────────────────         ────────────────────────────       ─────────────────
 sort+slice idiom
   │ TopkRewriter::SortIsInTopK (0x29e24b0)
   │   TransformPatternToCustomCall (0x29e5930)
   ▼   → CreateTopKCustomCall (0x29e4930)  target="TopK"  (NOT AwsNeuronTopK)
 HloTopKInstruction (k=, largest=)

 Neuron Frontend (out-of-binary) ─emits─► AwsNeuronTopK ─[LegalizeTopK::Run 0x1f01260]─► TopK_f32/f16/bf16 (+S64 iota)
                                                                                            │ NKI kernel: router_topk / topk_reduce /
                                                                                            │   cascaded_max / rotational_topk
                                                                                            ▼
                                                                                          nisa.max8                  → codegenMax8                → 88 Max                      → "Max-8"
                                                                                          nisa.nc_find_index8        → codegenFindIndex8          → 89 MaxIndex                 → "Find-Index-8"
                                                                                          nisa.nc_match_replace8     → codegenMatchReplace8       → 90 MatchReplace             → "Match-Value-Load"
                                                                                          nc_match_replace_indices8  → codegenMaxIndexAndMatchReplace → 91 MaxIndexAndMatchReplace

BIR InstructionType enum (InstructionType2string @ 0x2d5bf0, 110 values):

KLR/Walrus codegen hooks: BirCodeGenLoop exposes codegenMax8, codegenFindIndex8, codegenMatchReplace8, codegenMaxIndexAndMatchReplace; neuronxcc::backend::KlirToBirCodegen::codegenMatchReplace8(Ptr<klr::MatchReplace8>) returns a bir::InstMaxIndexAndMatchReplace. NKI kernels using the family: cascaded_max.py, router_topk.py, topk_reduce.py, rotational_topk_utils.py, plus compiled _private_kernels/router_topk.*.so.

NOTE — the DVE microcode integers are not pinned. The BIR enum ordinals and the NKI ISA names above are read directly, but the exact DVE-engine integer opcode behind each microcode name (Max-8, Find-Index-8, Match-Value-Load) is not; that column of the mapping rests on name correspondence alone and remains an open gap. The internal TopK_* kernel body — its per-k, per-dtype tiling of max8 / find_index8 / match_replace8 — lives in a compiled .so that was not decoded here.

End-to-end: AwsNeuronTopK → (legalize-topk) → TopK_f32/f16/bf16 (+ S64 iota) → (NKI router_topk/topk_reduce) → nisa.max8 / nc_find_index8 / nc_match_replace8 → (KlirToBirCodegen) → BIR Inst 88 Max / 89 MaxIndex / 90 MatchReplace / 91 MaxIndexAndMatchReplace → DVE Max-8 / Find-Index-8 / Match-Value-Load.

Evidence summary

• AwsNeuronTopK is consumed here and produced nowhere in this binary. In strings.json, 0x25f115 lists referenced_by_functions = [LegalizeTopK::Run] with a single xref from 0x1f01321; no other function touches it.
• The dtype gate admits exactly three types. 0x1f01358 lea eax,[rcx-0Ah] with the jbe at 0x1f01364 admits et ∈ {0x0A, 0x0B}, and 0x1f01366 cmp ecx,10h admits 0x10 — matching the three topkNames keys.
• The emitted names decode from inline immediates. 0x1f01bcf mov rax,3233665F4B706F54h is TopK_f32; 0x1f01bee is TopK_f16; 0x1f01c0d is TopK_bf1 plus a trailing '6', giving TopK_bf16. These are the payload movs; the key-store sites in the 0x1f01c17 region are a separate set of instructions and decode to nothing meaningful on their own.
• k is parsed, range-checked, then dropped. strtol @ 0x1f0159c, the range check against 0x7FFFFFFF @ 0x1f015c5/cf, and CreateCustomCall @ 0x1f01901 preceded by push 1 (API_VERSION_ORIGINAL) @ 0x1f018e3 — with the opaque argument slot holding the empty string.
• The second operand is an S64 iota over the input's minor dimension. mov esi,5 (S64) @ 0x1f014c4 feeds MakeValidatedShape @ 0x1f014d7; PopulateR1<long> sits @ 0x1f0154b, the array_state minor-dim read @ 0x1f01387, and CreateConstant @ 0x1f01738.

Limits of this reading

Two details are reconstructed rather than read. The 0-to-n−1 values the iota is populated with are [INFERRED] from the semantics of PopulateR1<long> and LiteralUtil::CreateR0<long> — the per-element fill body was never opened — though the S64 R1 shape and the operand position are exact. And the DVE microcode integers behind Max-8 / Find-Index-8 / Match-Value-Load remain unpinned, as noted in the crosswalk above.

Related Components

Cross-References

• The hlo-opt Pass Registry (the --passes Table) — resolves legalize-topk to order 8, vtable 0x40e350, Run 0x1f01260
• CC-Op Decompose & Legalize Family — sibling custom-call legalizers (ArgMax and the broader family)
• Scan / Reduce / Top-K Primitives — router_topk / topk_reduce kernels and the max8/find_index8/match_replace8 ISA surface
• The Bitonic SORT / TOPK Builtin Algorithm — the alternative GPSIMD-engine top-k path
• DVE Search & Datamove Encoding — the DVE microcode primitives BIR Inst 88–91 select