TpuHloCostAnalysis

地址适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（BuildID md5 89edbbe81c5b328a958fe628a9f2207d）。该二进制文件未剥离符号；下面的每个符号都是已 demangle 的 C++ 名称。.text/.rodata VMA == 文件偏移；.data.rel.ro VMA − 0x200000 == 文件偏移。其他版本会有所不同。

摘要

xla::jellyfish::TpuHloCostAnalysis 是 HLO 层级的 flop/字节/超越函数成本模型。它是通用 XLA xla::HloCostAnalysis 的轻量子类（vtable @ 0x218fb618，ctor @ 0x130a1620）：它原样继承每个按 opcode 的 emitter，只重写正好三个方法：HandleGather、HandleScatter、HandleCustomCall。每个 Handle* 遍历一条 HLO 指令，并将三个 float 成本属性写入按指令索引的 Properties flat-hash-map：flop 数（+0x50）、超越函数操作数（+0x54）和访问字节数（+0x58）。熟悉的参照框架完全是上游 XLA 的 HloCostAnalysis：相同的 kFlopsKey / kTranscendentalsKey / kBytesAccessedKey 属性，相同的 2·M·N·K dot 模型；只在通用模型会错误估计硅片成本的三个操作上加入 TPU 专用补丁：gather 和 scatter（按 chunk 粒度 DMA 计价，而不是按 flops），以及 custom-call（由按 target key 索引的注册表计价）。

这些 flop 属性是 TPU 编译器在同一份 HLO 上构建的两个成本面之一，本文不能把它们混为一谈。这里的 Properties flops 是融合优先级输入：flop_count(inst)（@ 0x1e4841e0）读取 Properties[inst].float[+0x50]，并供给 NormalizedComputationCost 按 unique_id 缓存的卷积周期估计。另一个成本面是bundle 占用率成本：CostModel::RecordHloCycles（@ 0x130bbfe0）通过 opcode 跳转表路由，将按 op 的吞吐周期存入 23 槽 ResourceVector，入口是 GetHloResourcesImpl（@ 0x130aa580，见 GetHloResources 路由）。两者共享同一份 HLO，也同意哪些 op 成本较高，但它们由不同函数计算，存入不同结构，并且从不共享数值。本文记录 flop/字节 Properties 模型，以及与之并行运行的 RecordHloCycles opcode→slot 路由。

第三条线索是路由决策：它选择由哪个成本子 emitter 为某个 op 计价。GetHloResourcesImpl 是五路分派：collective → 网络模型；conv-lowerable / reduce-window → MXU 输出融合模型；collective-compute fusion → 自有模型；否则 → 最终调用 RecordHloCycles 的 loop/elementwise 路径。本文覆盖重写表、RecordHloCycles opcode 跳转表，以及这些重写未改动的基础 flop 公式。

对重新实现而言，契约是：

• 子类形状：哪三个 Handle* 被重写，+0x50/+0x54/+0x58 Properties 布局，以及开启这些重写的 model_tpu_specific_overheads gate（this+456）。
• TPU 模型逐字继承的基础 HloCostAnalysis flop 公式：dot、convolution（同时除以 feature-group 和 batch-group 计数）、reduce、reduce-window，以及 elementwise 的 flop 与 transcendental 分类。
• 三个 TPU 重写：gather/scatter chunk-ratio 字节模型、scatter 额外加入的 combiner flops，以及按 custom_call_target 索引的 custom-call FunctionRegistry。
• RecordHloCycles opcode→ResourceVector 槽位跳转表：发出的按 op CT::Instruction、落入的具名槽、周期数量，以及把 add/subtract 分散到专用 vector-ALU lane 的浮点类型 gate。

类	xla::jellyfish::TpuHloCostAnalysis : xla::HloCostAnalysis
Vtable / ctor	0x218fb618 / 0x130a1620
重写方法	HandleGather 0x130a2de0, HandleScatter 0x130a3160, HandleCustomCall 0x130a35c0
Properties 字段	+0x50 flops, +0x54 transcendentals, +0x58 bytes-accessed（全为 float）
Flop 读取器	HloCostAnalysis::flop_count 0x1e4841e0 → Properties[inst].float[+0x50]
重写 gate	*(byte*)(this+456) = model_tpu_specific_overheads；off → tail-call base
路由分派	CostModel::GetHloResourcesImpl 0x130aa580（五路）
Bundle 写入	CostModel::RecordHloCycles 0x130bbfe0 + jump table @ .rodata 0xae0ebbc
源文件	platforms/xla/service/jellyfish/tpu_hlo_cost_analysis.cc

---

GetHloResourcesImpl — 五路路由分派

目的

GetHloResourcesImpl 是前门，它决定某个 HLO op 在写入 bundle 的 ResourceVector 之前由哪个成本子模型计价。它不是 flop 路径（那是 Properties/Handle* 一侧）；它是资源路由路径，在 collective-network 模型、MXU conv/output-fusion 模型、collective-compute-fusion 模型，以及到达 RecordHloCycles 的 loop/elementwise 默认路径之间做选择。

入口点

GetHloResourcesImpl (0x130aa580)            ── StatusOr<ResourceVector>
  ├─ IsSupportedCollectiveHlo (0x130aeda0)  ── collective?  → GetCollectiveCycles (0x130abfc0)
  ├─ IsFusionSupportedHlo     (0x130abee0)  ── numeric-type + fusion gate
  ├─ IsConvLowerable          (0x14553620)  ┐
  ├─ ExtractConvLikeHlo       (0x1d6aa140)  ├ conv / reduce-window selection
  ├─ GetReduceWindowType      (0x1454d4a0)  ┘  (type −1/2 max-pool → not conv)
  │     → GetOutputFusionOrConvolutionCycles (0x130aede0)
  ├─ IsCollectiveComputeFusion(0x13e028c0)  → GetCollectiveComputeFusionCycles (0x130b13a0)
  └─ (default)                              → GetLoopFusionOrUnfusedHloCycles  (0x130b2bc0)
```text

### 算法

```c
StatusOr<ResourceVector> GetHloResourcesImpl(inst, opts, fs, isFused):  // sub_130AA580
    op = inst.opcode;                              // byte [inst+0xc]

    // (1) COLLECTIVE — all-reduce / all-gather / … priced by the network path.
    if IsSupportedCollectiveHlo(inst):             // sub_130AEDA0
        // (a fusion of opcode 0x3d that wraps a collective is also peeled here)
        return GetCollectiveCycles(inst, &rv);     // sub_130ABFC0 — rv[+8] = scalar cycles

    // (2) NUMERIC-TYPE + FUSION-SUPPORT GATE
    et = inst.shape.element_type;                  // [inst+0x58]->[+0]
    if et <= 0x21 and (0x2FFF91FFE >> et)&1 ...    // numeric/packed-type mask (bittest64)
       and IsFusionSupportedHlo(inst, opts.target): // sub_130ABEE0
        // (3) CONV / REDUCE-WINDOW MAIN-OP SELECTION
        convlowerable = IsConvLowerable(inst);     // sub_14553620
        conv = ExtractConvLikeHlo(inst);           // sub_1D6AA140
        if conv and conv.opcode == 0x5e:           // reduce-window (0x5e == 94)
            t = GetReduceWindowType(conv);         // sub_1454D4A0
            if t == -1 or t == 2:                  // max-pool / unknown — NOT conv-priced
                goto default_path;
        if convlowerable:
            return GetOutputFusionOrConvolutionCycles(inst, opts, fs);  // sub_130AEDE0

    // (4) COLLECTIVE-COMPUTE FUSION (collective fused with compute)
    if IsCollectiveComputeFusion(inst):            // sub_13E028C0
        return GetCollectiveComputeFusionCycles(inst, opts, fs);        // sub_130B13A0

default_path:
    // (5) DEFAULT: loop / elementwise / unfused op — the per-op deposit path
    return GetLoopFusionOrUnfusedHloCycles(inst, opts, fs);             // sub_130B2BC0

0x130aa9dc 处的 numeric gate 是 cmp et, 0x21，随后对 element-type ordinal 使用 bt 0x2FFF91FFE：这是 bundle 模型可计价的 numeric/packed element types 位掩码（非 numeric 的 tuple/token/opaque type 会直接落入后续路径）。reduce-window sentinel 是一个不直观的分支：如果某个 reduce-window 的 GetReduceWindowType 为 -1（unknown）或 2（max-pool），它就不是 lowerable convolution，因此按普通 loop op 计价，而不是经过 MXU output-fusion estimator。

注意 — GetLoopFusionOrUnfusedHloCycles 是唯一会到达 RecordHloCycles 的分支。对 fused region，它经过 RecordCyclesIfFused（@ 0x130cc720），后者先剥离 convolution（opcode 0x2b → RecordConvolutionCycles @ 0x130ca6c0）、reduce-window（0x5e → RecordReduceWindowCycles @ 0x130c94e0）、loop-fusion 和 output-fusion，然后才按 leaf op 落入 RecordHloCycles。IsProducerUse（@ 0x130ab0c0）会在写入前丢弃 producer→consumer edge 的输入 DMA。conv / output-fusion emitter 的 window-iteration 主体不在本文范围内（见 Bundle 感知成本）；本文覆盖 leaf opcode→slot 路由。

函数映射

函数	地址	作用
CostModel::GetHloResourcesImpl	0x130aa580	五路路由分派
IsSupportedCollectiveHlo	0x130aeda0	Collective predicate（分支 1）
IsFusionSupportedHlo	0x130abee0	Numeric + fusion-support gate
IsConvLowerable	0x14553620	Conv-lowerable predicate
ExtractConvLikeHlo	0x1d6aa140	拉取 conv/reduce-window root
GetReduceWindowType	0x1454d4a0	−1/2 max-pool sentinel
GetCollectiveComputeFusionCycles	0x130b13a0	分支 4
GetLoopFusionOrUnfusedHloCycles	0x130b2bc0	默认 → RecordHloCycles

---

RecordHloCycles — Opcode→Slot 跳转表

目的

RecordHloCycles 是 leaf 的按 op 写入逻辑：给定一条 HLO 指令、一个输出元素数，以及 bundle 的 ResourceVector，它把该 op 的吞吐周期写入这个 op 占用的具名 Resource 槽。op→slot 决策是对 opcode 的 switch，编译成 .rodata 0xae0ebbc 处 0x7f 项的自相对跳转表（index = opcode − 3，ja > 0x7e → default）。

每次写入会发出一个或多个 CT::Instruction（约 33 个 bucket 的 LLO 折叠），通过 CycleTable::GetResource(k)（@ 0x1c89ce20，gen-invariant table @ 0xb438aec；见 Resource Enum）解析槽位，通过 cycle-table vtable+0x10 读取按代际的吞吐量（GetCyclesForThroughput(k)，per-gen，见 Per-Opcode Cycle Constants），并累加 Acc(slot, element_count × throughput × W)（ResourceVector::Acc @ 0x1c89adc0），其中 W 是从 .rodata 读取的固定 FP 乘数。

算法

RecordHloCycles(inst, window, rv, fs, nesting):       // sub_130BBFE0
    elems     = Product(output_dims);                 // xla::Product → [rbp-0x30]
    isFloat   = ShapeUtil::ElementIsFloating(inst.shape);  // sub @0x130bc... line 139
    if RecordHloCyclesIfTopLevel(...) != 1: return;   // sub_130CDD80 (fusion-root guard)

    switch (inst.opcode):                             // jump table @ .rodata 0xae0ebbc
      case 0x03 add:        slot = isFloat ? GetResource(CT 0x12) : R5;  // R4 / R5
                            Acc(slot, elems × thru(CT 0x12));
      case 0x4b multiply:   Acc(GetResource(CT 0x14), elems × thru(CT 0x14));   // R3
      case 0x7b subtract:   slot = isFloat ? GetResource(CT 0x13) : R5;  // R4 / R5
                            Acc(slot, elems × thru(CT 0x13));
      case 0x32 divide:     /* 4-deposit reciprocal+mul micro-sequence */
            Acc(R6, elems × thru(CT 0x18));                       // VectorEup
            Acc(GetResource(CT 0x14)=R3, elems × thru(CT 0x14) × 3.0);
            Acc(GetResource(CT 0x12)=R4, (elems×2) × thru(CT 0x12));
            Acc(R5, elems × 9.0);                                 // VectorAluAny
      case 0x47 logistic:   /* sigmoid micro-sequence */
            Acc(GetResource(CT 0x12)=R4, elems × thru(CT 0x12));
            Acc(GetResource(CT 0x14)=R3, (elems×2) × thru(CT 0x14));  // VectorAlu0, same slot as multiply
            Acc(R6, elems × thru(CT 0x1a));                       // VectorEup, lane-cmp
      case 0x38 erf:        if isExtPrecPath(inst):               // vtable+312(11)
                                Acc(R6, elems × thru(CT 0x11));   // single VectorEup deposit
                            else:                                 // 4-deposit polynomial
                                Acc(R6, elems × thru(CT 0x18));
                                Acc(GetResource(CT 0x14)=R3, elems × thru(CT 0x14) × 16.0);
                                Acc(GetResource(CT 0x12)=R4, (elems×2) × thru(CT 0x12));
                                Acc(R5, elems × 4.0);
      case 0x2a convert:    if ElementHasBitWidth(inst.shape, 1): // 1-bit PRED
                                Acc(R5, elems × 2.0);             // packed-bool repack
                            else: /* no deposit — free */
      case 0x6c select:     Acc(R5, elems × 2.0);                 // read both branches + pred
      case 0x52 parameter:  if IsFused(inst): /* RecordFusionInputCycles → MemXfer */
                            else: /* no deposit — non-fused param is free */
      case 0x5b reduce:     if IsFused(inst): Acc(R5, elems);
                            else: w = GetInputWindow(inst); Acc(R5, Product(w));  // operand window
      case 0x18,0x1a,0x27,0x29,0x43,0x61,0x81: /* ZERO — no deposit */
      default:              Acc(R5, elems × 1.0);                 // VectorAluAny, 1 cy/elem
```text

> **易错点 —** 这个表中有两个 slot-routing 陷阱。(1) `add` 和 `subtract` 受**浮点类型 gate** 控制：`slot = ElementIsFloating(shape) ? GetResource(CT) : R5`。整数 add/subtract 像 default 一样写入 `VectorAluAny`（`R5`）；只有浮点 add/subtract 会到达专用 lane（通过 `CT 0x12`/`0x13` 到 `R4`）。忽略该 gate 会让整数 fusion 过度占用专用 lane。(2) `convert` **只**对 1-bit PRED 写入（`R5`，`elems × 2.0`，packed-bool repack）；所有更宽的 numeric width-conversion 都落入 zero path，成本为零，这与“窄转换免费、宽转换计费”的直觉正好相反。

### 解码后的跳转表

`element_count` 是 `Product(output dims)`；`thru(k)` 是按代际的 `GetCyclesForThroughput(CT::Instruction k)`；`W` 是固定 FP 乘数。CT-bucket→slot 通过 [`CycleTable::GetResource`](resource-enum.md#opslot-mapping--cycletablegetresource)。

| Opcode（名称） | Block @ | 发出的 CT → slot | 周期数量 |
|---|---|---|---|
| `0x03` add | `0x130bc0f7` | float: `CT 0x12`→`R4`; int: `R5` | `elems × thru(0x12)` |
| `0x4b` multiply | `0x130bc4dd` | `CT 0x14`→`R3` | `elems × thru(0x14)` |
| `0x7b` subtract | `0x130bc4af` | float: `CT 0x13`→`R4`; int: `R5` | `elems × thru(0x13)` |
| `0x32` divide | `0x130bc3cf` | `R6` + `CT 0x14`→`R3` + `CT 0x12`→`R4` + `R5` | 4-deposit；`×3.0`(R3), `×2`(R4 elems), `×9.0`(R5) |
| `0x47` logistic | `0x130bc17b` | `CT 0x12`→`R4`, `CT 0x14`→`R3`, `CT 0x1a`→`R6` | sigmoid seq；`×2`(R3 elems) |
| `0x38` erf | `0x130bc245` | fast: `CT 0x11`→`R6`; slow: `R6`+`R3`+`R4`+`R5` | gated；slow `×16.0`(R3), `×2`(R4), `×4.0`(R5) |
| `0x2a` convert | `0x130bc293` | 1-bit: `R5`; else none | 1-bit: `elems × 2.0`；wider: 0 |
| `0x6c` select | `0x130bc2a9` | `R5`（`LABEL_26` ×2 path） | `elems × 2.0` |
| `0x52` parameter | `0x130bc2b8` | fused→`RecordFusionInputCycles`; else none | fused: input-DMA into MemXfer；non-fused: `0` |
| `0x5b` reduce | `0x130bc2f9` | `R5` over operand window | 按被 reduce 的输入窗口计价 |
| DEFAULT（大多数 op） | `0x130bc3c0` | `R5` `VectorAluAny` | `elems × 1.0` |
| ZERO-cost | `0x130bc8c2` | （无写入） | `0` |

DEFAULT 分支覆盖上面未具名的每个 numeric elementwise / structural op（按 decompiler 的 jumptable annotation：`cases 4-23, 25, 27-38, 40, 43-49, 51-55, 57-66, 68-70, 72-74, 76-81, 83-90, 92-96, 98-107, 109-122, 124-128`）。ZERO 分支覆盖 data-layout op：`bitcast`（`0x18`）、`broadcast`（`0x1a`）、`concatenate`（`0x27`）、`constant`（`0x29`）、`iota`（`0x43`）、`reshape`（`0x61`）、`tuple`（`0x81`）。

### 固定 FP 乘数（`.rodata`，按字节验证）

| 地址 | 值 | 使用处 |
|---|---|---|
| `0xa2df230` | `1.0` | default per-op multiplier |
| `0xa2df930` | `3.0` | divide `VectorAlu0` (R3) scale |
| `0xa2deb40` | `9.0` | divide `VectorAluAny` (R5) scale |
| `0xa2df040` | `16.0` | erf slow-path `VectorAlu0` (R3) scale |
| `0xa2de830` | `4.0` | erf slow-path `VectorAluAny` (R5) scale |
| `0xa2df5c8` | `0.5` | `MaxResourceCycles` `VectorAlu` port-balance blend |

### 考量

三个 `VectorAlu` 槽之间的路由，是 `MaxResourceCycles` 0.5-blend 组 `{R3, R4, R5}` 的 port-balancing 输入（见 [Resource Enum](resource-enum.md#the-maxresourcecycles-reduction)）。Multiply 固定到 `VectorAlu0`（`R3`），floating add/subtract 固定到 `VectorAlu1`（`R4`），catch-all DEFAULT 加上 `×2` op（`select`、1-bit `convert`）以及 integer add/subtract 落在 `VectorAluAny`（`R5`），即 blend 会重新分配的 lane。因此，混合 multiply、floating add 和 “any” op 的 fusion 会填满两个专用 lane，bundle 成本是平衡后的最大值，而不是串行总和。

> **怪癖 —** bundle 模型在这里把 `broadcast`（`0x1a`）定价为**零**，即便 [NormalizedComputationCost](normalized-computation-cost.md) 可能在 fusion-priority 一侧对同一个 op 收取 cross-lane-movement 权重。两个模型对 layout op 的分歧是有意设计的：bundle 路径把 broadcast 视为免费的 register splat（无功能单元占用），而 priority 模型会计入数据移动。不要把它们统一起来。

`parameter` 和 `reduce` 会路由到 ALU lane **之外**。Fused `parameter` 是由 `RecordFusionInputCycles`（@ `0x130ce940`）计入 `MemXfer` 槽 `R9..12` 的输入 DMA；VMEM-resident params 是否按 memory space 改路由到 `R7` `VectorLoad` 尚未按字节确认（MEDIUM）。**非 fused** `parameter` 不写入任何内容：`case 0x52` 直接落到完成路径（`if (!IsFused) goto LABEL_77`），因此顶层 argument 的功能单元成本为零。非 fused `reduce` 通过 `GetInputWindow` 按其*operand* 窗口计价（因此成本随大的被 reduce 输入张量缩放，而不是随小的输出缩放），这是下面 `HandleReduce` 的 `ExtentProduct(operand)` flop 公式在 bundle 一侧的镜像。

---

## 基础 Flop 公式（继承自 `HloCostAnalysis`）

`TpuHloCostAnalysis` 只重写 gather/scatter/custom-call；其他每个 op 都原样使用基础 `xla::HloCostAnalysis` emitter。这些是缓存在 `Properties[inst].float[+0x50]`（flops）/ `[+0x54]`（transcendentals）中，并由 `flop_count`（@ `0x1e4841e0`）读回的 flop 属性。

### Dot — `2·M·N·K`

`HandleDot`（@ `0x1e47c9c0`）→ `GetDotFlops`（@ `0x1e47c7a0`）：

```c
GetDotFlops(out_shape, lhs_shape, dnums):           // sub_1E47C7A0
    contract_batch = 1;
    for d in dnums.contracting_and_batch_dim_indices:
        contract_batch *= lhs_shape.dims[d];         // imul over the index list
    return 2 * contract_batch * ExtentProduct(out_shape);  // ×2 for multiply-add

反编译结果是一串展开的 imul，遍历 DotDimensionNumbers 索引列表，然后 imul ExtentProduct(out)，再 add rax, rax（即 ×2）。这是经典的 2·M·N·K。

易错点 — raw dot 会在 TPU 成本路径运行前被 lowering 成 convolution（dot/conv MXU lowering）；在未 lowered 的 dot 上到达 GetDotFlops 会触发 CHECK-fatal（buffer != nullptr，shape.h:843）。实践中，为 fusion priority 提供输入的是下面的 convolution 公式。

Convolution — 同时除以两个 group count

HandleConvolution（@ 0x1e480be0）调用 GetConvolutionFlops（vtable+1128）并将结果存到 +0x50。GetConvolutionFlops（@ 0x1e480060）：

GetConvolutionFlops(inst, out, lhs, rhs):           // sub_1E480060
    out_spatial   = ExtentProduct(window-iterated output spatial+batch);  // v114
    window_vol    = Product(kernel spatial window sizes);                  // v108
    in_feature    = GetDimension(lhs, input_feature_dim);                  // Dimension
    out_feature   = GetDimension(out, output_feature_dim);                 // v115
    fgc = inst.feature_group_count(); bgc = inst.batch_group_count();
    return 2 * (out_feature / bgc) * (in_feature / fgc) * window_vol * out_spatial;
```text

> **易错点 —** convolution flop 同时除以 `feature_group_count` **和** `batch_group_count`，而不只是 feature-group：有两段独立的 `idiv`/`div` 序列（`v109 = Dimension/fgc`，`v110 = v115/bgc`）。任何丢掉 batch-group 除法的模型，都会错误估计 grouped convolution（depthwise、batch-grouped）的成本。

这是 [NormalizedComputationCost](normalized-computation-cost.md) conv 路径按 `unique_id` 缓存的 flop，然后除以 per-LHS-format peak 以得到 MXU cycles。

### Reduce / Reduce-Window — combiner 成本 × extent

`HandleReduce`（@ `0x1e47d6a0`）：`flops = to_apply().per_element_cost × ExtentProduct(operand_being_reduced)`。combiner subcomputation 的 per-element cost（flops、transcendentals、bytes，每个都通过 `vmulss` 缩放）乘以被 reduce 的**输入**元素数。成本随大的被 reduce 张量缩放，而不是随小的输出缩放。

`HandleReduceWindow`（@ `0x1e47e1c0`）：`flops = to_apply().per_element_cost × Product(window dim sizes) × ExtentProduct(output_shape)`；combiner 对每个输出元素的每个窗口元素运行一次。

两者都通过基础 `ProcessSubcomputation` 递归供给 combiner cost（此处不重新遍历）。

### Elementwise — flop/transcendental 分类器

`HandleElementwiseOp`（@ `0x1e47b320`）计算 `cost = ExtentProduct(output_shape)`（每元素 1 个 op），并根据 opcode 的 `switch` 选择存储偏移：

```c
HandleElementwiseOp(inst):                          // sub_1E47B320
    cost = ExtentProduct(inst.shape);
    off  = is_transcendental(opcode) ? 0x54 : 0x50; // ecx = 84 : 80
    Properties[inst].float[off] = (float)cost;

22 个 opcode 的 transcendental 集合（存到 +0x54），来自 switch case 的逐字节精确结果：

Opcode	名称	Opcode	名称	Opcode	名称
0x01	acos	0x2f	cosine	0x55	power
0x02	acosh	0x30	cosh	0x68	rsqrt
0x0e	asin	0x38	erf	0x75	sine
0x0f	asinh	0x39	exponential	0x76	sinh
0x13	atan2	0x3a	exponential-minus-one	0x79	sqrt
0x14	atanh	0x45	log	0x7c	tan
0x1c	cbrt	0x46	log-plus-one	0x7d	tanh
		0x47	logistic

其他每个 elementwise op（add、multiply、subtract、and/or/xor、compare、clamp、negate、abs、floor、ceil、sign、round、shift、min、max、remainder 等）都存到 +0x50（flops），每元素 1 个。flop 模型路由到 +0x54 的 transcendental 集合，与 bundle 模型按 multi-µop sequence 计价的集合相同（上面的 divide、logistic、erf），因此两个成本面对哪些 op 昂贵的判断是一致的。

---

TPU 专用重写

全部三个重写都受 *(byte*)(this + 456) gate 控制，即 model_tpu_specific_overheads 选项。未设置时，每个方法 tail-call 对应的基础 HloCostAnalysis 方法，TPU 模型不贡献额外内容。设置后如果 Target 指针（this + 56）为 null，gather/scatter 会抛出 FailedPrecondition("target_ not specified but model_tpu_specific_overheads ... enabled")（tpu_hlo_cost_analysis.cc）。

HandleGather — chunk-ratio 字节，无 flops

HandleGather（@ 0x130a2de0）：

HandleGather(inst):                                 // sub_130A2DE0
    if !this[456]: return base::HandleGather(inst);  // overrides off
    if !this.target: FailedPrecondition(...);
    out_sz   = GetShapeSize(inst.output_shape);                 // bytes
    ratio    = GetGatherSizeInChunkRatio(target, inst);         // sub_14A8E420
    op1_sz   = GetShapeSize(inst.operand(1).shape);             // index operand
    // four float fields stored at this+0x58/+0x6c/+0x70/+0x74 (bytes-accessed family)
    store {out_sz, ratio-amplified bytes, op1_sz, ...} to Properties bytes fields;
```text

Gather 被建模为**纯 memory/DMA op**，不增加 flops。成本是粒度“chunk ratio”（`GetGatherSizeInChunkRatio` @ `0x14a8e420`）：这是作用于输出和索引大小的 chunk-aligned read-amplification factor。反编译显示通过 `vcvtsi2ss`/`vmovss` 存储四个 `float` 结果（`+0x58`、`+0x6c`、`+0x70`、`+0x74`），也就是写入 bytes-accessed 属性族，而不是单个字段。`GetGatherSizeInChunkRatio` 的内部粒度数学未枚举（amplification factor 已确认；公式可信度 LOW）。

### HandleScatter — chunk-ratio 字节**加上** combiner flops

`HandleScatter`（@ `0x130a3160`）：

```c
HandleScatter(inst):                                // sub_130A3160
    if !this[456]: return base::HandleScatter(inst);
    if !this.target: FailedPrecondition(...);
    op2_sz = GetShapeSize(inst.operand(2).shape);               // updates
    ratio  = GetScatterSizeInChunkRatio(target, inst);          // sub_14A90CE0
    op1_sz = GetShapeSize(inst.operand(1).shape);               // indices
    store ratio-amplified bytes into Properties bytes fields (this+0x58/0x5c/0x60/...);
    // PLUS the scatter-combiner compute:
    upd_extent = ExtentProduct(inst.operand(2).shape);          // updates element count
    per_elem   = ProcessSubcomputation(inst.to_apply());        // vtable slot 150
    for each non-zero combiner property p:
        Properties[inst].float[p_offset] += per_elem[p] * upd_extent;  // incl. +0x54 transcend.

Scatter = gather 风格的 chunk-ratio 字节项加上 update-combine compute。反编译代码会将每个非零 combiner per-element property（flops +0x50、transcendentals +0x54，以及 bytes 字段）乘以 ExtentProduct(updates) 并折入 Properties；它通过按标准 cost-analysis property name 索引的 Properties::operator[] 迭代 combiner 的 property map。因此，带有 transcendental combiner 的 scatter 会正确增加 transcendental count，而不只是 flops。GetScatterSizeInChunkRatio 位于 0x14a90ce0（粒度公式可信度 LOW）。

HandleCustomCall — 按 target 索引的 registry，否则 base

HandleCustomCall（@ 0x130a35c0）：

HandleCustomCall(inst):                             // sub_130A35C0
    reg = CustomCallRegistration::GetGlobalRegistry();  // lazy __cxa_guard; FunctionRegistry
    cb  = reg.Get(inst.custom_call_target());           // string-keyed lookup
    if cb found:
        return cb(inst, &this.shape_size_fn /*this+240*/, &this.Properties /*this+80*/);
    VLOG(1) "Custom hlo cost analysis not found: " << target;
    return base::HandleCustomCall(inst);                // sub_1E482A20
```text

Custom-call 是开放的扩展 hook：`util_registration::FunctionRegistry<string, Status(HloInstruction*, function<long(Shape const&)> const&, Properties&)>` 按 custom-call target 字符串索引。注册的 callback 会计算自定义 flop/字节模型（传入 instruction、`this+240` 处的 `GetShapeSize` functor，以及 `this+80` 处的 `Properties` map）；未注册的 target 会记录一次 `VLOG(1)` miss，并回退到基础 emitter。这与 fusion 一侧的 custom-fusion registry 相呼应（[fusion cost model](../compiler/fusion-cost-model.md)）。

### 函数映射

| 函数 | 地址 | 作用 |
|---|---|---|
| `TpuHloCostAnalysis::HandleGather` | `0x130a2de0` | chunk-ratio bytes；受 `this+456` gate 控制 |
| `TpuHloCostAnalysis::HandleScatter` | `0x130a3160` | chunk-ratio bytes + combiner flops/transcend. |
| `TpuHloCostAnalysis::HandleCustomCall` | `0x130a35c0` | target 上的 `FunctionRegistry`；否则 base |
| `GetGatherSizeInChunkRatio` | `0x14a8e420` | gather read-amplification factor |
| `GetScatterSizeInChunkRatio` | `0x14a90ce0` | scatter write-amplification factor |
| `HloCostAnalysis::GetShapeSize` | `0x1e47a6e0` | shape 的字节大小 |
| base `HloCostAnalysis::HandleCustomCall` | `0x1e482a20` | fallback emitter |

---

## 示例 — elementwise fusion bundle 写入

一个 loop-fusion body `{ multiply [256,128] (f32), add [256,128] (f32), tanh [256,128] }`。`element_count = 256·128 = 32768`。每个 op 都通过 `RecordHloCycles` 跳转表路由：

```text
multiply (0x4b)  → CT 0x14 Shuffle → GetResource → R3 VectorAlu0:   Acc(R3, 32768 × thru(0x14))
add (0x03), f32  → CT 0x12 RotIn   → GetResource → R4 VectorAlu1:   Acc(R4, 32768 × thru(0x12))
tanh (0x7d)      → DEFAULT block   →               R5 VectorAluAny: Acc(R5, 32768 × 1.0)

flop 一侧将 tanh 归类为 transcendental（Properties+0x54），将 multiply/add 归类为 flops（Properties+0x50）；bundle 一侧则通过上表为三者计价。MaxResourceCycles 在 {R3, R4, R5} 上的 0.5-blend 会并行运行 multiply（VectorAlu0）和 add（VectorAlu1）lane，把 tanh 的 “any” 工作负载均衡到较空闲的 lane，并以 50% 重叠剩余部分，因此 bundle 成本约为 max(R3, R4) + half the tanh residual，而不是串行总和。如果 add 是整数类型，它会按 HCA-1 与 tanh 一起落到 R5，而不是固定到 R4。

---

相关组件

组件	关系
GetHloResources 路由	本文总结的完整 GetHloResourcesImpl → sub-emitter 路由
NormalizedComputationCost	flop_count / conv flop 的 fusion-priority 消费者
Resource Enum（23 槽）	写入流向的 ResourceVector 槽和 MaxResourceCycles reduction
Bundle 感知成本	基于 per-op 写入的 loop composition
Per-Opcode Cycle Constants	乘入每次写入的 per-gen GetCyclesForThroughput 整数

交叉引用

• 成本模型概览 — 构建写入查询所用 per-gen CycleTable / Performance 的三类 class family
• GetHloResources 路由 — 五路分派，以及把 fusion 剥离到 RecordHloCycles 的 router
• NormalizedComputationCost — 消费 flop_count 和 conv flop 的 opcode→weight fusion-priority 模型
• Resource Enum（23 槽） — Resource 槽名、Acc 和 MaxResourceCycles overlap 模型
• Per-Opcode Cycle Constants — per-gen GetCyclesForThroughput(CT::Instruction) 整数
• CycleTable Family — GetResource（op→slot）和 Instruction bucket enum
• Bundle 感知成本 — 基于 per-op 写入的 software-pipelined loop cost
• dot/conv MXU Lowering — raw dot 为何会在成本路径之前 lowered 成 convolution
• Fusion Cost Model — 与 custom-call hook 对应的 custom-fusion registry
• Compiler Overview — HLO cost analysis 在编译流水线中的位置