RankAndPermute 与基数排序顺序
地址适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id
89edbbe81c5b328a958fe628a9f2207d）。其他版本会有所不同。
摘要
RankAndPermute 是 SparseCore embedding 基数排序的逐位主体：这一趟会把一列已排序的 embedding id 转换为一个 gather permutation index vector，即从每个原始 id 到其去重代表所在槽位的映射。它由 RadixSortEmitterInternal 发出（SparseCore 方言 SortOp 降低到的 MLIR 层 emitter），也是 VectorExtended 页面记录的 bundle 层 Sort/Uniquify/DuplicateCount opcode 的 SSA 发射对应物。本页负责三件事：RankAndPermuteComputeFunction SSA 形状（UniqueOp / UniqueWithLaneIdsOp 去重加逆置换）、SparseMapRow 按行 sort+reduce（op-type 0x4），以及 radix digit 如何在 gather 前驱动排序。
结构思路是经典 LSD radix sort，但降低为 MLIR ops，而不是标量循环。每一位 pass 都从已排序 key 中提取一个 digit (key >> shift) & mask，在 (key, digit) 对上去重，用 indexed scatter-add 为每个 unique value 构建 rank，然后把 values 置换到 rank 顺序。compute function 是逐 chunk 回调；RankAndPermute 包装器用两个 CreateChunkIteratorLoop pass 将其平铺到 id 流上。整个结构围绕 SparseCore Unique/Permute/VectorLoadStoreIdxAdd/VectorStoreIdx 方言 ops，它们会降低为 TEC VEX reduction bundles。
对于重新实现者，契约是：
- 去重 key 是
(sorted_key, radix_digit)对。radix_digit = (sorted_key >> shift) & mask由GetDigits根据RadixSortEmitterInternal成员计算；UniqueOp和UniqueWithLaneIdsOp都把原始 sorted-key Value 和该 digit 作为两个操作数。 UniqueOp产生 3 个结果；UniqueWithLaneIdsOp产生 5 个结果。 capability/flag 查询会选择 with-lane-ids 形式，其两个额外结果是 lane-id / per-unique multiplicity，用来给 permute 加权。- rank 是 indexed scatter-add，permute 是其逆。
VectorLoadStoreIdxAddOp构建每个 unique id 的 prefix-position rank；PermuteOp(rank, lane_id)重新排序 values；VectorStoreIdxOp写出最终的 permutation index vector，并以原始 sorted key 为键。 SparseMapRow（op0x4）是降序排序加扫描。 它用"dscd"direction StringAttr 对一行的 keys 排序，使重复 token-id 相邻，然后"max"scan 确定逐行输出窗口大小；携带的ReduceDuplicatesreduce-fn 会把现在相邻的重复项折叠成一个 ELL row。
| Compute fn | RadixSortEmitterInternal::RankAndPermuteComputeFunction 0x134039c0 (0x1240 B) |
| Wrapper | RadixSortEmitterInternal::RankAndPermute 0x13404dc0 (0xE40 B) |
| IR 层级 | MLIR — SparseCore (sc_tpu.*) dialect ops，LLVM lowering 之前 |
| 去重原语 | UniqueOp::create 0x14622400 (3 results) · UniqueWithLaneIdsOp::create 0x146231a0 (5 results) |
| Rank / permute | VectorLoadStoreIdxAddOp::create 0x14634fc0 · PermuteOp::create 0x145f3920 |
| 输出存储 | VectorStoreIdxOp::create 0x14638460 (sc_tpu.vector_store_idx) |
| Digit 提取 | GetDigits 0x133fe480 · GetMappedKeysForRadixSort 0x133fe8c0 |
| 调用链 | EmitSort 0x131c5fa0 → SingleDigitRadixSort 0x1340c580 / …WithSoftwareCoalescing 0x1340c740 → RankAndPermute |
| SparseMapRow 窗口 | FusionEmitter::SetOutputWindowBoundsForSparseMapRow 0x13890d40 (0x9c0 B) |
| 源文件 | platforms/xla/sparse_core/lowering/internal/radix_sort_emitter_internal.cc (0x8783fe2) |
| 置信度 | CONFIRMED（decompile + mangled-symbol 锚定），除非某行或 callout 另有说明 |
RankAndPermute — 逐位 Rank+Permute Pass
目的
RankAndPermuteComputeFunction 是一个 radix digit pass 中一个 chunk 的 compute 回调。给定一个已排序 keys 的 chunk 以及随它们移动的 values，它会：(1) 提取每个 key 的 digit，(2) 在 (key, digit) 上去重，(3) 计算 rank，即每个 unique value 的 prefix position，(4) 发出将原始 values 重新排序到该 rank 顺序的 permutation，并把它存为 index vector。跨所有 digits 重复执行后，逐 chunk permutations 会组合成完整的 sorted+deduplicated ordering，供 embedding gather 和 sample combiner 消费。
Signature 和 ABI
const member function 来自 demangled symbol (0x134039c0)：
// RadixSortEmitterInternal::RankAndPermuteComputeFunction(...) const
StatusOr<SmallVector<Value,6>>
RankAndPermuteComputeFunction(
Target const& target, // rsi → [rbp-0x88]
OpBuilder builder, // by value [rbp+0x10] = r12
int chunk_start, // edx → [rbp-0x44] (per-chunk loop var)
Value key, // rcx → [rbp-0xc0] (the to-be-digitised key)
Value v5, // r8 → [rbp-0x108] (threaded to SliceBuffer)
KeyValueRangeData const& range, // r9 (drives GetMappedKeys/GetDigits)
Value sorted_key, // [rbp+0x68] — a19: the original sorted key column
ArrayRef<Value> values_in, // [rbp+0x38]/+0x40 (gather-index prologue)
ArrayRef<Value> values_out, // [rbp+0x48..] (the VectorStoreIdx loop)
Value, Value, Value, // alias-scope / loc helpers
AliasScopeAssignment* assignment, // a20
AliasScope* scope) const;QUIRK —
this指针（rdi → [rbp-0x30]）是RadixSortEmitterInternal，radix 参数存在它自己的成员中，而不是KeyValueRangeData中：+0x28= num-buckets exponent（digitmask），+0x34= digitshift，+0x38= 每个 chunk 的num_keys（循环上界）。KeyValueRangeData const&（r9参数）携带缩放 bucket mapping 的 key/value range bounds；其字段布局未解码（LOW）。
算法
function RankAndPermuteComputeFunction(...): // sub_134039c0
// --- PROLOGUE: 2 SmallVector<u32> (sized [this+0x38]=num_keys) + a gather-index chain
grow_pod+memset ×2 // ~line 0x13403ab1
for v in values_in: AddIOp(base, stride) // per-value offsets into values buffer
// --- KEY DERIVATION: digit = (key >> shift) & mask ----------------
et = operand(0).Shape::element_type() // 0x134040ad / b6 — key element type
mk = GetMappedKeysForRadixSort(et, …) // 0x133fe8c0 — bucket-map the key
digit = GetDigits([this+0x28], [this+0x34], // 0x133fe480 → [rbp-0xa8]
key, mk) // the two i32 scalars are the only
// ↑ decompile line 493: GetDigits(this, *(u32*)(this+0x28), *(u32*)(this+0x34), key, mk)
// GetDigits body: ConstantIntOp + MulIOp + BroadcastScalarToVector + ShRUIOp + AndIOp ( + opt SubIOp )
vecTy = VectorType::get(getI32Type(), …) // 0x1340411f/35 → [rbp-0xa0]
// --- DEDUP: select Unique vs UniqueWithLaneIds by a capability query
if target.caps[+0x88]() // 0x13404... line 497
and target.SupportsSparseCore() // line 499 (vtable slot 76)
and target.caps[+0x80]() : // line 501 → with-lane-ids path
u = UniqueWithLaneIdsOp::create(builder, loc, // 0x146231a0
sorted_key, digit) // (a19, digit) — SAME two operands
unique_vals = u.getNextResultAtOffset(0) // result0 → [rbp-0x98]
unique_idx = u.getNextResultAtOffset(1) // result1 → rbx
lane_id = u.getNextResultAtOffset(2) // result2 → [rbp-0x90] (multiplicity)
else:
u = UniqueOp::create(builder, loc, sorted_key, digit) // 0x14622400, op "sc_tpu.unique"
unique_vals = u.getNextResultAtOffset(0) // result0 → [rbp-0x98]
unique_idx = u.getNextResultAtOffset(1) // result1 → rbx
lane_id = NULL // [rbp-0x90] = 0 → no PermuteOp
// --- RANK: indexed scatter-add over a sliced values window --------
BroadcastI32ToVector + SubIOp // 0x134043aa / 1340444b — rank base
win = SliceBuffer(values, …) // 0x13404c00 — memref::SubViewOp window
rank = VectorLoadStoreIdxAddOp::create(builder, loc, // 0x14634fc0 → [rbp-0xb8]
vecTy, unique_idx, win, unique_vals, {digit}) // "sc_tpu.vector_load_store_idx_add"
assignment.AddToNewScope(rank.getDefiningOp()) // CHECK: "base_offsets" / "input_chunk"
// --- PERMUTE: re-order the values by rank × multiplicity ----------
if lane_id != NULL:
rank = PermuteOp::create(builder, loc, vecTy, // 0x145f3920, op "sc_tpu.permute"
rank, lane_id) // (rank vector, lane-id weight)
// --- STORE: write the permutation index vector per element --------
for i in 0 .. [this+0x38]: // cmp [rbp-0x44] < num_keys
s = VectorStoreIdxOp::create(builder, loc, /*bool*/0, // 0x14638460
permute_result, dest_slot, /*key=*/sorted_key, // "sc_tpu.vector_store_idx"
{gather_index}) // keyed by the sorted key (a19)
assignment.AddToScope(s, *scope) // CHECK: "output_scatter"
return ok // eax=1; free the SmallVectorsNOTE —
Unique操作数顺序是(sorted_key, digit)。 两个Uniquecreate 调用都按create(builder, loc, sorted_key, digit)的顺序接收操作数，即原始 sorted-key Value 在前，GetDigitsradix digit 在后（decompile0x134039c0第 529 / 569 行：UniqueWithLaneIdsOp::create(&a8, loc, a19, Digits)和UniqueOp::create(&a8, loc, a19, Digits)）。去重以(sorted_key, digit)对为键。相同的sorted_keyValue（a19）也会被复用为最终VectorStoreIdxOp的 keying operand（第 878 行）。
Unique SSA 形状
两个去重原语共享完全相同的两个操作数，只在结果数量上不同。getNextResultAtOffset(base, n)（0x1d8e9700）返回 op 的第 n 个结果；decompile 读取 offset 0、1、2，也就是前三个结果 result0/result1/result2，连续索引，不是隔一个取一个。
| Op | ::create | Results | 提取结果（本函数） | 降低到 |
|---|---|---|---|---|
UniqueOp | 0x14622400 | 3 | result0 = unique values，result1 = unique index/marker | tpu_uniquei/tpu_uniquef（3-field LLVM struct） |
UniqueWithLaneIdsOp | 0x146231a0 | 5 | result0 = unique values，result1 = index，result2 = lane-id / multiplicity | 相同 3-field struct + 2 个 ReplaceOpWithExtracts |
两者都通过共享的 DuplicateCountUniqueOpLowering<T> 模板降低（UniqueOp 位于 0x1359b280，UniqueWithLaneIdsOp 位于 0x1359bd20），该模板用 LLVMStructType::getLiteral(ctx, types, 3, 0) 构建 3-field literal struct（decompile 0x1359b280 第 46 行；0x600000003 字是 field-type list 的 size-3/capacity-6 SmallVector header）。lane-id 形式会为其两个新增结果多做两次 extract。每个字段的含义（unique-values vs write-mask vs segment-marker）是从 extract 顺序推断的，而不是来自字段名表（LOW）。result2 lane-id 是 per-unique multiplicity，即折叠进每个 unique entry 的原始 id 数量，并且是 PermuteOp 的第三个操作数（第二个 Value）。
GOTCHA —
UniqueOp（3 个结果）和UniqueWithLaneIdsOp（5 个结果）之间的选择是 target capability/flag 查询（builder.target.caps，decompile 第 497 行），不是带命名 accessor 的 config field（LOW）。在普通UniqueOp路径中，lane_id被设为0（第 573 行），并且整个PermuteOp步骤会被跳过，直接存储来自VectorLoadStoreIdxAddOp的 rank。总是发出PermuteOp的重新实现会在 no-lane-ids 路径上因 null operand 崩溃。
Rank 和逆置换
rank 构建是这一 pass 的核心，也是重新实现者最容易弄错数据流的地方。
// VectorLoadStoreIdxAddOp::create(builder, loc, vecTy, unique_idx, win, unique_vals, {digit})
// 0x1340456c — decompile: create(&a8, loc, v192, v83, v104, NextResultAtOffset, {Digits})
// vecTy = i32 VectorType [rbp-0xa0]
// unique_idx = Unique result1 (v83) — where each id writes
// win = SliceBuffer SubViewOp — the values gather window
// unique_vals = Unique result0 — the values being ranked
// {digit} = ValueRange{digit} — the ranked-by digit
// ⇒ result [rbp-0xb8] = the RANK: prefix position of each unique id in the deduped window.
// Registered as "base_offsets" in the alias scope (CHECK string @ line 685).VectorLoadStoreIdxAddOp 是 indexed scatter-add：它在 unique indices 处读取 sliced values window，执行 add，并产生每个 id 的 prefix position（rank）。随后 PermuteOp(rank, lane_id) 会把 values 重新排序到 rank 顺序，lane-id multiplicity 则给每个 unique entry 占据的槽位数量加权。结果是 inverse permutation：对于每个原始 id，给出其去重代表所在的槽位。VectorStoreIdxOp 会按元素写出该 vector，以原始 sorted_key 为键，写入目标缓冲区（output_scatter）。
函数映射
| Function | Address | 作用 |
|---|---|---|
RankAndPermuteComputeFunction | 0x134039c0 | 逐 digit、逐 chunk 的 rank+permute 主体 |
RankAndPermute (wrapper) | 0x13404dc0 | 将 compute fn 平铺到 chunks 上（两个 pass） |
GetDigits | 0x133fe480 | digit = (key>>shift)&mask（ShRUIOp+AndIOp+…） |
GetMappedKeysForRadixSort | 0x133fe8c0 | 在 digit extract 前对 key 做 bucket-map |
SliceBuffer (anon ns) | 0x13404c00 | values buffer 的 memref::SubViewOp 窗口 |
$_0 callback | 0x1341b7e0 | trampoline → RankAndPermuteComputeFunction（pass 1） |
$_1 callback | 0x1341b9e0 | 配套的 vector::BroadcastOp twin（pass 2） |
DuplicateCountUniqueOpLowering<UniqueOp> | 0x1359b280 | LLVM lowering → 3-field struct + extracts |
RankAndPermute 包装器与基数排序链 {#the-rankandpermute-wrapper-and-the-radix-sort-chain}
目的
RankAndPermute（0x13404dc0）驱动 compute function 遍历 id 流。它计算 chunk 数，加载 sorted keys，并在 CreateChunkIteratorLoop 中运行逐 chunk 回调。它由 SingleDigitRadixSort 对每个 digit 调用一次，而 RadixSortEmitter::EmitSort（0x131c5fa0）会对每个 digit 调用 SingleDigitRadixSort，以实现 SparseCore SortOp 降低到的多 digit LSD radix sort。
入口点
RadixSortEmitter::EmitSort 0x131c5fa0 the SortOp lowering target
└─ SingleDigitRadixSort 0x1340c580 one digit (call 0x1340c6cf)
│ └─ RankAndPermute 0x13404dc0 ── this wrapper
└─ SingleDigitRadixSortWithSoftwareCoalescing 0x1340c740 one digit, coalesced tails (call 0x1340c954)
└─ RankAndPermute 0x13404dc0算法
function RankAndPermute(target, builder, …, range, values_in, values_out): // sub_13404dc0
assignment.NewScope() // 0x13404e72 (line 181)
n_chunks = ceil(n / chunk_size) // RemUIOp 0x134050ab + SubIOp + DivUIOp 0x134052b4
keys = lowering_util::LoadChunk(sorted_keys) // 0x1340544c (line 439) — load the sorted keys
// PASS 1 — the rank+permute itself
CreateChunkIteratorLoop($_0) // 0x1340565d (line 513); $_0 0x1341b7e0
assignment.Materialize() // 0x134056e1 (line 554) — emit the alias scopes
// PASS 2 — the companion BroadcastOp pass
assignment.NewScope() // line 602
CreateChunkIteratorLoop($_1) // 0x13405a28 (line 662); $_1 0x1341b9e0
assignment.Materialize() // line 703
returnNOTE —
$_0和$_1是两个 chunk-iterator pass 的std::functiontrampolines。mangled names 确认二者都是RankAndPermute的 lambda closures（_functions.json中的…RankAndPermute…$_0/…$_1）；$_0调用RankAndPermuteComputeFunction，$_1是vector::BroadcastOpcompanion。包装器每个 pass 都打开一个新的NewScope，并在各自之后Materialize；这两个 pass 共享的 alias scoping 是 compute function 注册的input_chunk/base_offsets/output_scatter三元组。
SparseMapRow — 按行排序并规约
目的
SparseMapRow（SparseCore op-type 0x4，由 IsSparseMapRowHlo 0x13d7efe0 测试 SparseCoreOperationTypeFromString(...) == 4 确认）是 reduce_duplicates custom-call 发出的去重并折叠原语；SparseDenseMatmulOpDecomposer::ReduceDuplicates 0x136722e0 会将它与 DynamicBoundedSlice（op 0x11 = 17）一起构建。结构上，它是对一行 (token, sample) keys 的降序字典序排序，随后接一个 scan：排序会让重复 token-id 相邻，而携带的 reduce-fn（例如 add）会把每段相邻重复项折叠成一个 ELL row，即实现重复 multiplicity 的 CSR→ELL row collapse。multiplicity weighting 侧见 Dedup Multiplicity，ScanOp / SegmentedScanOp lowering 见 Scan Datapath。
算法 — Window Emitter
FusionEmitter::SetOutputWindowBoundsForSparseMapRow（0x13890d40）是 lowering 中确定逐行输出窗口大小的一半。它识别该 op，构建逐行 segment state，然后发出一个降序 SortOp 和一个 "max" scan：
function SetOutputWindowBoundsForSparseMapRow(target, hlo, builder, state, …): // sub_13890d40
if not IsSparseMapRowHlo(GetRootInstruction(hlo)): return // 0x13890d71 / d81
row = operand(0) // 0x13890da3 — the row input
BroadcastBoolToVector / BroadcastI32ToVector + GetCurrentState // per-row segment/window state
bound = ConstantIndexOp + AddIOp + AddIOp // window-bound arithmetic
dir = StringAttr::get(builder.ctx, "dscd") // 0x1389133d; "dscd" @0x8720761
s = SortOp::create(builder, loc, Ty, Ty, Ty, // 0x1389136a — 3 result types,
row, key1, key2, dir) // 3 key Values, sort-dir "dscd"
col = s.getNextResultAtOffset(1).getNextResultAtOffset(0) // 0x138913a6 — sorted column
ExtractVectorElement(col, 0) // pull the scalar window bound
redop = getStringAttr(builder, "max") // 0x138914a4; "max" @0x84c6977
sc = ScanOp::create(builder, loc, Ty, data, segid, redop) // 0x138914bf — window-extent scan
ExtractVectorElement(sc, …) ; IndexCastOp // 0x1389150d / a5 — finalise the bounddecompile（0x13890d40）确认第 468 行有字面量 "dscd"（Twine kind word 259 = 0x103，cstring-ptr+len）、第 474 行的 SortOp::create 有三个结果类型（v80, v81, v81）和三个 key Values，并且第 311/312 行有 getStringAttr("max") + ScanOp::create 对。
GOTCHA — 这里的
"max"scan 是 window-extent scan，用来计算逐行输出窗口大小，而不是 duplicate-value collapse。duplicate-value reduce 使用ReduceDuplicatesSparseMapRowcustom-call 携带的HloComputation（例如add），由一个单独的 FusionEmitter value-reduce 路径内联。两个ScanOp创建器是不同的；混淆它们会产生对 indices 求和而不是对 values 求和的 row collapse（value-collapse emitter 的主体这里未解码 — MEDIUM）。
"dscd" / "ascd" 排序方向
4 字符 sort-direction StringAttr 已通过 SortOpLowering::matchAndRewrite（0x13597700）中的 intrinsic split 按字节确认；该函数从 .rodata 池 "dscd\0ascd\0…"（0x8720761 / 0x8720766）读取两个字符串，并选择三个 comparator lambda 之一。direction × element-type 乘积映射到四个不同 intrinsics：
| Direction | StringAttr .rodata | Element type | Intrinsic ::create | 含义 |
|---|---|---|---|---|
"ascd" | 0x8720766 | integer | tpu_sort_ascdi 0x14739520 | ASCENDING int keys |
"ascd" | 0x8720766 | float | tpu_sort_ascdf 0x14738d00 | ASCENDING float keys |
"dscd" | 0x8720761 | integer | tpu_sort_dscdi 0x1473a560 | DESCENDING int keys |
"dscd" | 0x8720761 | float | tpu_sort_dscdf 0x14739d40 | DESCENDING float keys |
SparseMapRow 使用 "dscd"（降序），所以重复 token-id 会相邻，供携带的 reduce-fn 折叠。方言 SortOp op name 是 "sc_tpu.sort"（0x84de6a5）；它的 bundle 层 opcodes 和 EmitVectorSort emitter 归 VectorExtended 页面所有，这里不再重复。
SparseMapRow Op-Create 序列
| # | Op / call | Address | 作用 |
|---|---|---|---|
| 0 | IsSparseMapRowHlo (0x13d7efe0) + GetRootInstruction | 0x13890d71/d81 | 识别 op（...FromString == 4） |
| 1 | operand(0) | 0x13890da3 | row input |
| 2 | BroadcastBoolToVector / …I32… + GetCurrentState | 0x13890ede/ef4 | 逐行 segment state |
| 3 | ConstantIndexOp + AddIOp ×2 | 0x13891136+ | window-bound 算术 |
| 4 | StringAttr::get("dscd") | 0x1389133d | sort-dir attr（0x8720761） |
| 5 | SortOp::create(Ty,Ty,Ty, V,V,V, "dscd") | 0x1389136a | 3 个 result tys，3 个 keys |
| 6 | getNextResultAtOffset ×2 + ExtractVectorElement | 0x138913a6+ | 提取 sorted column + scalar |
| 7 | getStringAttr("max") | 0x138914a4 | reduction_op（0x84c6977） |
| 8 | ScanOp::create(Type, data, segid, "max") | 0x138914bf | window-extent scan |
| 9 | ExtractVectorElement + IndexCastOp | 0x1389150d/a5 | 完成 bound |
Radix Digit 如何排列索引
该排序是一个降低为 MLIR ops 的多 digit LSD radix sort。完整 radix-sort emitter（RadixSortEmitterInternal）是 count/scan/scatter 机制，其 rank pass 已在上文记录；count 和 bucket-scan 两半，即 HistogramKeysComputeFunction 0x133feca0、ScanBucketsComputeFunction 0x13400120、CalculateGatherIndices 0x13400400，在下文 Histogram 和 Bucket-Scan 两半 中解码。本页为把 digit histogram 转为逐 digit permutation 的 rank+permute 固定如下事实：
- Digit 提取。
GetMappedKeysForRadixSort（0x133fe8c0）对 sorted key 做 bucket-map；GetDigits（0x133fe480）用ShRUIOp+AndIOp+ConstantIntOp+MulIOp+BroadcastI32ToVector计算digit = (mapped_key >> shift) & mask。shift是RadixSortEmitterInternal+0x34，mask来自+0x28处的 num-buckets exponent。 - 在
(sorted_key, digit)上去重。UniqueOp/UniqueWithLaneIdsOp折叠相等的(key, digit)对，产生 unique values、unique index/marker，以及（带 lane ids 时）per-unique multiplicity。 - Rank。
VectorLoadStoreIdxAddOp将每个 unique value scatter-add 到其 bucket position，以产生 prefix-position rank，即每个 id 的去重槽位索引（base_offsets）。 - Permute。
PermuteOp(rank, lane_id)按 rank 顺序重新排序 values，并按 multiplicity 加权；VectorStoreIdxOp写出以原始 key 为键的 permutation index vector。
跨所有 digits 组合逐 digit permutations 会得到全局 sorted+deduplicated id ordering。由于每个 digit pass 都会去重，最终 pass 之后的 gather window 是 unique-id count，而不是 raw id count，这正是 stream gather/scatter 路径依赖的 redundant-gather elimination。
NOTE — digit
mask/shift和 bucket count 由RadixSortEmitterInternal成员（+0x28/+0x34/+0x38）与KeyValueRangeData参数共同驱动。送入 digit math 的成员 offset 已 CONFIRMED；KeyValueRangeData字段映射（key/value range bounds）未解码（LOW），所以确切 bits-per-pass parameterisation 是根据成员使用推断的。
Histogram 和 Bucket-Scan 两半 {#the-histogram-and-bucket-scan-halves}
上面的 rank+permute 是每个 digit pass 的第三阶段。它由前两个阶段供给：histogram（统计每个 bucket 的 keys 数量）和 bucket-scan（bucket counts 的 exclusive prefix-sum），再加上最后的 gather-index 算术，该算术把扫描后的 bucket base 转成逐元素 scatter offset。三者都是 RadixSortEmitterInternal 上的逐 chunk compute functions，主体在这里解码（CONFIRMED，decompile 锚定到 radix_sort_emitter_internal.cc 源码行）。
Histogram — HistogramKeysComputeFunction 0x133feca0
histogram 阶段按 digit value 统计有多少 keys 映射到每个 bucket。对于每个 chunk，它加载 masked keys 窗口，派生 digit（与 rank pass 使用相同的 GetMappedKeysForRadixSort + GetDigits 对），在 (key, digit) 对上去重，然后用 VectorStoreIdxOp 把结果 scatter 到 histogram buffer；这是 rank pass 用于其输出的同一个 indexed-store 原语。
function HistogramKeysComputeFunction(builder, …, key, …, range, target, opt_mask): // sub_133feca0
chunk = lowering_util::LoadChunkMasked(key, …) // masked chunk load (CHECK status line 244)
assignment.AddToNewScope(chunk.getDefiningOp()) // → "input_chunk" alias scope (line 246)
if opt_mask.has_value(): // a14 optional<int> upper bit set
m = BroadcastI32ToVector(opt_mask) // 0x133fef.. — the mask bound
cmp = arith::CmpIOp::create(builder, …) // line 253 — lane-active predicate
v = arith::AndIOp::create(builder, loc, chunk, cmp) // line 255 — gate inactive lanes
mk = GetMappedKeysForRadixSort(et, chunk, …) // 0x133fe8c0 — bucket-map the key
digit = GetDigits([this+0x28],[this+0x34], chunk, mk) // 0x133fe480 — (mapped>>shift)&mask (line 269)
u = UniqueOp::create(builder, loc, chunk, digit) // 0x14622400 — dedup (key,digit) (line 272)
uidx = u.getNextResultAtOffset(1).getNextResultAtOffset(0) // unique index (the bucket slot)
uval = u.getNextResultAtOffset(0) // unique value (the count contribution)
s = VectorStoreIdxOp::create(builder, loc, /*bool*/1, // 0x14638460 — scatter into histogram
uidx, ctx, uval, {digit}) // keyed by the digit bucket
assignment.AddToNewScope(s) // → "scatter" alias scope (CHECK line 274)NOTE — histogram 路径使用普通
UniqueOp（3 个结果），从不使用 lane-ids 形式；count 阶段只需要 unique value 和它的 bucket index，不需要 rank 阶段的PermuteOp消费的 multiplicity weight。可选 mask 参数（std::optional<int>，decompile 第 198 行的a14upper-bit test）是 software-coalescing tail mask：它会屏蔽部分 final chunk，使越界 lanes 不对 histogram 贡献计数。不存在该参数时，会统计完整 chunk。
Bucket-Scan — ScanBucketsComputeFunction 0x13400120
bucket-scan 把逐 bucket histogram counts 转为每个 bucket 的 scatter base offset，这是经典 exclusive prefix-sum。它是一个带 "sum" reduction 的 ScanOp，后接 inclusive→exclusive 修正（AddIOp − SubIOp）。
function ScanBucketsComputeFunction(builder, histogram, …): // sub_13400120
seg = lowering_util::BroadcastBoolToVector(…) // 0x13400... — all-active scan segment
sum = ScanOp::create(builder, loc, vecTy, seg, // 0x138... ScanOp, op "sc_tpu.scan"
histogram, StringAttr "sum") // "sum" @ ptr (Twine kind 259); line 411
last = BroadcastVectorElementToVector(sum, n-1) // line 420 — broadcast the inclusive total
add = arith::AddIOp::create(builder, …) // line 420 — inclusive prefix
arith::SubIOp::create(builder, …) // line 424 — subtract own count ⇒ exclusive
return add // the exclusive bucket base offsetsNOTE — bucket prefix-scan 是一个带
"sum"reduction 的sparse_core::ScanOp。 reduction StringAttr 是"sum"（decompile0x13400120：字面量"sum"在调用mlir::StringAttr::get前以ptr[0]="sum"和 Twine kind259/0x103存储，随后传给ScanOp::create）。这个ScanOp会完全按 Scan Datapath 所述通过ScanOpLowering::matchAndRewrite（0x1358ab00）降低；radix bucket-scan 是该 lowering 的一个使用者，不是单独的 scan 原语。inclusive scan 由末尾的AddIOp/SubIOp对转换为 exclusive base offset（源码第 420/424 行）。histogram 侧 prefix 是 vectorScanOp，不是tpu_mprefixi1-count 路径（该路径用于 boolean inputs；histogram 是i32count vector）。
Gather-Index — CalculateGatherIndices 0x13400400
最后的算术会把 exclusive bucket base（逐 bucket scalar）转为逐元素 destination index。它构建 lane sequence，加上 scanned base，然后用 modulo/divide 拆分每个 tile 的几何：
function CalculateGatherIndices(builder, a2, a3, base): // sub_13400400
seq = VlaneseqOp::create(builder, loc, idxVecTy) // 0x145... "sc_tpu.vlaneseq" (line 438)
s32 = arith::IndexCastOp(seq) // index→i32 (line 441)
b = vector::BroadcastOp(IndexCast(base)) // broadcast the scanned bucket base
pos = arith::AddIOp(s32, b) // line 448 — global position = lane + base
w = vector::BroadcastOp(ConstantIntOp a2) // tile-width constant (line 453)
lo = arith::RemUIOp(pos, w) // line 455 — within-tile offset
hi = arith::DivUIOp(pos, w) // line 459 — tile index
s = vector::BroadcastOp(ConstantIntOp a3) // tile-stride constant (line 463)
return arith::AddIOp(arith::MulIOp(hi, s), lo) // line 463/464 — tile*stride + offsetNOTE —
RemUIOp/DivUIOp拆分（第 455/459 行）将线性 scanned position 分解为(tile_index, within_tile_offset)对，然后重组为tile_index*stride + offset（第 463/464 行），使 scatter 落入 multi-tile scan 的正确 VMEM tile。两个arith::ConstantIntOp立即数（a2、a3，即 tile width 和 tile stride）由调用方传入；它们如何从RadixSortEmitterInternal成员派生未追踪（LOW）。CalculateGatherIndices是 bucket-scan base 与 rank pass 的VectorLoadStoreIdxAddOp窗口之间的桥梁。
三阶段链
| Phase | Compute fn | Key op(s) | Alias scope / output |
|---|---|---|---|
| 1 — Histogram | HistogramKeysComputeFunction 0x133feca0 | LoadChunkMasked → UniqueOp → VectorStoreIdxOp | input_chunk / scatter（逐 bucket counts） |
| 2 — Bucket-scan | ScanBucketsComputeFunction 0x13400120 | ScanOp("sum") → AddIOp/SubIOp | exclusive bucket base offsets |
| 3a — Gather-index | CalculateGatherIndices 0x13400400 | VlaneseqOp + RemUIOp/DivUIOp + MulIOp/AddIOp | 逐元素 scatter destination |
| 3b — Rank+permute | RankAndPermuteComputeFunction 0x134039c0 | UniqueWithLaneIdsOp → VectorLoadStoreIdxAddOp → PermuteOp | 逐 digit permutation |
平铺这些 compute functions 的 wrappers 是 HistogramKeys 0x133ff3a0、ScanLocalTileHistograms 0x13400c60、ScanHistogram 0x13401000、ScanBuckets 0x13401540 和 MultiTileScanBuckets 0x13402340（multi-VMEM-tile 变体）；每个 wrapper 都像 RankAndPermute 驱动 rank pass 一样，通过 CreateChunkIteratorLoop 驱动自己的 compute function。
PermuteOp ISA Lowering
rank pass 发出的 PermuteOp（0x145f3920，op name sc_tpu.permute）本身会降低为一条 SC TEC VectorAlu 指令。该 lowering 是一个很薄的单 pattern conversion；重活在 bundle encoder 中完成。
PermuteOpLowering::matchAndRewrite 0x135a1640 {#permuteoplowering-matchandrewrite-0x135a1640}
char PermuteOpLowering::matchAndRewrite(PermuteOp op, PermuteOpAdaptor a, ConversionPatternRewriter& r):
rty = a.getOperand(0).getType() // result type = type of operand[0]
p = mlir::sparse_core::tpu_sc_permute::create( // 0x14735ac0 — the lowered op
r, loc, TypeRange{rty}, ValueRange{op.operands}) // both PermuteOp operands flow through
r.replaceOp(op, p) // (*vtable+8)(rewriter, op, p)
return 1 // always matchesdecompile（0x135a1640）正是如此：从 operand[0] 派生 result type（ValueRange::dereference_iterator(…, 0) 与 0xF8 mask），构建 1 元素 TypeRange，调用 tpu_sc_permute::create，然后 replaceOp。它与 UniqueOp/DuplicateCountOp/SortOp lowerings 一起注册在 SparseCore→LLVM RewritePatternSet 中（0x13572820 处的 RewritePatternSet::add<…PermuteOpLowering…> 模板）。
tpu_sc_permute Op 及其 Masked Twin
tpu_sc_permute::build（0x147359a0）确认 op 形状：它 addOperands 两个输入 Values（第 12 行）和一个结果类型（第 33–45 行），即一个 2-operand、1-result op（NOperands<2>、OneTypedResult<Type>、MemoryEffectOpInterface）。还有一个 masked twin：tpu_sc_mask_permute，它具有相同的 2-operand/1-result 形状，是会门控 inactive output lanes 的 predicated permute（tpu_sc_mask_permute op name 在 .rodata 中出现 37 次，而 tpu_sc_permute 出现 39 次）。
| MLIR op | Op name | Operands | Result |
|---|---|---|---|
mlir::sparse_core::PermuteOp | sc_tpu.permute | 2 (data, index) | 1（= operand[0] 的类型） |
tpu_sc_permute (lowered) | tpu_sc_permute | 2 | 1 |
tpu_sc_mask_permute (predicated) | tpu_sc_mask_permute | 2 | 1 |
TEC VectorAlu Permute Opcode
tpu_sc_permute 由 ISA emitter 实现为一条 TEC VectorAlu VectorPermute 指令，并按元素宽度分派。opcode 位于 VectorAlu op word 的 bits 13..20（mask 0x1FE000）中，这一点已通过 …Opcode::Matches predicates 按字节确认：
| ISA op | Matches raw | Opcode field (word>>13)&0xFF | Element width | Address |
|---|---|---|---|---|
SparseCoreTecVectorAlu0VectorPermuteB32 | (word & 0x1FE000) == 0x10A000 | 0x85 | 32-bit | 0x1ea9f620 |
SparseCoreTecVectorAlu0VectorPermuteB16 | == 0x10C000 (1097728) | 0x86 | 16-bit | 0x1ea9f640 |
SparseCoreTecVectorAlu0VectorPermuteB8 | == 0x10E000 (1105920) | 0x87 | 8-bit | 0x1ea9f660 |
radix-sort permute 操作的是 i32 indices，因此它使用 VectorPermuteB32（opcode field 0x85）。它的 bundle emit 是 EmitVectorBinop<…SparseCoreTecVectorAlu_VectorPermuteB32…> / EmitVectorY<…> 模板族（gxc/glc 0x13a19ae0/0x13a21ec0，gxc/gfc 0x13aae200/0x13aae340），会通过 SparsecoreVregReadPort 集合路由两个操作数，并写入目标 VREG。viperfish（v5）generation 有自己的 SparseCoreTecVectorAlu_VectorPermute emitter（0x139ad760）；B8/B16/B32 宽度拆分是 gxc 时代（v6e/v7x）的细化。
NOTE — permute 是 VREG 内部的跨 lane gather：operand[0] 是 data vector（ranked values），operand[1] 是 per-lane source-index vector（rank）。
VectorPermuteB32会读取每个 lane 的 index，取出该 source lane 的元素，并写到 destination lane，即dst[i] = src[index[i]]。在 radix pass 中，这实现了 inverse permutation：每个原始 id 都从其去重代表的 rank slot 读取。tpu_sc_mask_permutetwin 会添加一个 M-register predicate，使 inactive output lanes 保持不变，与 Scan Datapath 上的 masked-scan datapath 匹配。
相关组件
| Name | Relationship |
|---|---|
RadixSortEmitter::EmitSort 0x131c5fa0 | SortOp lowering 入口；按 digit 调用 SingleDigitRadixSort |
HistogramKeysComputeFunction 0x133feca0 | radix sort 的 count 半边（上文已解码） |
ScanBucketsComputeFunction 0x13400120 | 供给 rank pass 的 bucket prefix-sum（ScanOp("sum")）（上文已解码） |
CalculateGatherIndices 0x13400400 | 将 scanned bucket base 转为逐元素 scatter index（上文已解码） |
PermuteOpLowering::matchAndRewrite 0x135a1640 | 降低 sc_tpu.permute → tpu_sc_permute（上文已解码） |
SparseCoreTecVectorAlu0VectorPermuteB32Opcode::Matches 0x1ea9f620 | TEC VectorAlu permute opcode（field 0x85，mask 0x1FE000） |
SparseDenseMatmulOpDecomposer::ReduceDuplicates 0x136722e0 | 发出 reduce_duplicates SparseMapRow + DynamicBoundedSlice（op 0x11）custom-calls 的 decomposer |
交叉引用
- VectorExtended (VEX) — 负责 bundle 层
Sort/Uniquify/DuplicateCountopcodes，以及这些 MLIR ops 降低到的EmitVectorSortemitter - Scan Datapath —
SparseMapRowwindow emitter 使用的ScanOp/SegmentedScanOplowering - Dedup Multiplicity — 给
PermuteOp和 ELL row collapse 加权的 lane-id /DuplicateCountmultiplicity - Emit Valency Loop — 消费去重后 gathered ids 的逐 sample valency loop
- Sample Combiner Emitter — 驱动 sort→dedup→gather→reduce datapath 的 embedding combiner
- Stream Gather / Scatter — 对 permutation 产生的 deduped unique-id window 执行 gather
- SparseCore Overview — radix-sort dedup 在 embedding pipeline 中的位置