Custom-Call 降低与目标注册表

本页中的所有地址、符号、目标字符串和 proto 名称都适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，构建 libtpu_lts_20260413_b_RC00）。其他版本会有所不同。

摘要

HloOpcode::kCustomCall 是 JAX、PyTorch/XLA 和 Pallas 注入核心 HLO 集无法表达的行为时所经过的唯一 HLO opcode：Mosaic/Pallas 设备 kernel、分片和内存放置标记、线性代数块原语、host-offload、MegaScale collective 元数据以及 SDC instrumentation。TPU 编译器不会统一处理这些内容，而是按指令的 custom_call_target 字符串分派；该字符串是普通的 std::string，通过 Swiss-table 相等性与进程级注册表匹配。本页是该分派层的规范。

最重要的结构事实：xla::jellyfish::CustomCallRegistration 不是多态基类。它是五个相互独立的全局注册表的持有者，每个注册表都是针对不同 facet 回调类型（lowering emitter、can-fuse 谓词、编译属性、HLO 成本分析、SPMD 分区 visitor）的 util_registration::FunctionRegistry<std::string, T>。五者都以同一个目标字符串为键；handler 在模块初始化时选择加入它所需的 facet。这里没有不透明编码，也没有虚表；分派就是 target_string → callback 查找，针对编译器对每个 custom-call 提出的五个正交问题重复五次。

对于重新实现，本页固定的契约是：

五 facet 注册表 — 它的五个 Register* 入口点、它们的精确回调签名、CompilationProperties 声明式元数据结构体，以及通过 .init_array 进行的 google_init_module_* 播种机制。
目标目录 — 跨五个类别约 52 个高置信度目标字符串，每个字符串都列出注册它的 handler/namespace 常量以及触发的降低动作。
六种降低动作 — pre-pass HLO 重写、标记剥离、emit-helper → LLO、Mosaic 缓存 MLIR body、runtime/host-action 透传，以及仅 SPMD 分区，并说明哪些目标采用哪种动作。
Mosaic 导入接缝 — tpu_custom_call（逃逸口）如何在 backend_config.custom_call_config.mlir_module 中携带序列化的 tpu dialect 模块，如何被解析/缓存为 MosaicMlirCacheEntry，并通过 CustomCallEmitter::Emit 路由。（下游 tpu dialect pipeline 记录在 MHLO → XTile → tpu 降低和 Mosaic 概览中 — 此处只链接，不重复。）
验证层与错误路径 — HloVerifier → TpuHloSupportChecker（其中拒绝 __cudnn$…/__triton$… 目标）→ TpuCustomCallLegalizer，以及每层发出的已恢复 DCHECK/status 字符串。


分派层	`xla::jellyfish::CustomCallRegistration`（5 个 facet 注册表，来源 `…/jellyfish/custom_call_registration.h`）
键类型	`custom_call_target` : `std::string`（Swiss-table，默认 hash）
Facet 注册表	`RegisterLoweringEmitter`（3 个重载） · `RegisterCanFuse` · `RegisterCompilationProperties` · `RegisterHloCostAnalysis` · `RegisterSpmdPartitioningVisitor`
注册入口点	`RegisterLoweringEmitter` @ `0x10e8bf40`/`0x10c9de40`/`0x10e8f1c0` · `RegisterCanFuse` @ `0x10eb5680` · `RegisterCompilationProperties` @ `0x10f940e0` · `RegisterHloCostAnalysis` @ `0x111eee40` · `RegisterSpmdPartitioningVisitor` @ `0x14ba8000`
播种	每个 handler 通过 `.init_array` 调用 `google_init_module_*()`，由 BSS 保护
Mosaic 逃逸目标	`"tpu_custom_call"` → `CustomCallEmitter::Emit` @ `0x111ef740`，body 通过 `GetCachedCustomCallBody` @ `0x13e31860`
Mosaic 缓存	`MosaicMlirCacheEntry`，以 `tsl::Fprint128` 为键，存储在 `HloModule` 上
验证	`HloVerifier` → `xla::TpuHloSupportChecker` → `TpuCustomCallLegalizer::RunImpl` @ `0x11036080`
目录规模	约 52 个不同的高置信度目标字符串
置信度	高（符号/字符串锚定），除非某行或标注另有说明

CustomCallRegistration 回答关于 custom-call 的五个正交问题，每个问题都有自己的注册表。反编译符号表确认了五个不同的 FunctionRegistry<std::string, T> 实例化（每个 facet 一个 T），而逐字的 Register*("Name", callback) 源码引用以宏字符串化 DCHECK 参数的形式保留在 .rodata 中，因此精确的目标字符串和回调体就是事实依据。

Facet（注册表）	注册入口点	回调签名（已恢复）	回答的问题
Lowering emitter	`RegisterLoweringEmitter` ×3	`StatusOr<OperandData>(HloInstruction, LoweredGetter const&, LloRegion, LloValue, BackendConfigMap, … context)`	该目标如何发射到 LLO？
Can-fuse	`RegisterCanFuse` `0x10eb5680`	`bool(HloInstruction* producer, HloInstruction* consumer, Target const&, optional<FusionOptions>, optional<HloReachabilityMap const*>)`	该 op 是否可以与相邻 op 融合？
编译属性	`RegisterCompilationProperties` `0x10f940e0`	`CompilationProperties(HloInstruction*)`	该 op 携带哪些声明式标志？
HLO 成本分析	`RegisterHloCostAnalysis` `0x111eee40`	`Status(HloInstruction*, ShapeSizeFunction const&, HloCostAnalysis::Properties&)`	它的 FLOPs/bytes 是多少？
SPMD 分区	`RegisterSpmdPartitioningVisitor` `0x14ba8000`	`Status(SpmdPartitioningVisitor, HloInstruction)`	它如何跨 shard 分区？

emitter 签名中的 BackendConfigMap 已从 demangled __policy_func thunk 中完整恢复：它是 absl::flat_hash_map<HloInstruction const*, std::unique_ptr<BackendConfig>>。完整 emitter 原型（所有三个重载的可选 context 尾部的并集）是：

cpp

using LoweringEmitter = std::function<absl::StatusOr<OperandData>(
    const HloInstruction*        hlo,
    const LoweredGetter&         get_lowered,
    LloRegion*                   region,
    LloValue*                    output,
    BackendConfigMap*            backend_config_map,
    // optional context tail — present only in the wider overloads:
    const ProgramSharedRegistry*       registry        = nullptr,
    const LogicalTopologyInfo*         topology_info   = nullptr,
    sdc_reporter::SdcRegistrator*      sdc_registrator = nullptr,
    const llo_log::LogRecorder*        log_recorder    = nullptr,
    const HloInstruction*              original_hlo    = nullptr)>;

三个 RegisterLoweringEmitter 重载的差异仅在于它们接收多少 context 尾部；三者最终汇入同一个以 std::string 为键的 map。最窄形式（无 context）由 PartialReduce 等简单 op 使用；最宽形式（带 SDC/log/topology）由触及 SDC checker 或发射 log record 的 op 使用。

`CompilationProperties` — 声明式元数据结构体

每个 handler 都可以注册一个 RegisterCompilationProperties 回调，返回 CompilationProperties 聚合体。已恢复的 designated-initializer 源码引用固定了字段集合：

cpp

struct CompilationProperties {
  bool has_communication            = false;  // carries ICI/DCN traffic
  bool supports_hlo_dedup           = false;  // HLO-level CSE-safe
  bool instruction_can_change_layout = true;  // layout assignment may relayout
  bool supports_internal_checksums  = false;  // SDC checker integration
  bool requires_mxu_assigner        = false;  // forces MXU register assignment
  bool check_fifos_are_empty        = false;  // must drain pipelines first
};

已恢复的注册（逐字来自 .rodata）显示这些设置通常非常稀疏 — 大多只设置一个字段：

目标	已恢复的 `CompilationProperties` body
`kSharding`, `kPartialReduce`, `kDevicePlacement`, `kTopk`, `kTopkBatchMajorSmallK`, generic `target`	`{ .instruction_can_change_layout = false }`
`WindowPrefetchEmitter::kWindowPrefetch`	`{ .supports_hlo_dedup = true }`
`IciSdcTestEmitter::kKey`	`{ .has_communication = true, .supports_hlo_dedup = true }`
`SdcCheckerGetStatsEmitter::kKey`, `SdcCheckerReportSdcEventEmitter::kKey`, `SdcCheckerStartWithAlternativeCoresEmitter::kKey`	`{ .has_communication = false, .supports_hlo_dedup = true }`

注意 — tpu_custom_call 不能使用静态 CompilationProperties。 与上面的固定目标 handler 不同，Mosaic 逃逸口没有编译期已知的属性 — 它的行为是用户序列化的 MLIR。因此它的 RegisterCompilationProperties 回调会在查找时查询缓存 body，以动态填充该结构体。已恢复的回调体（.rodata 中的逐字源码引用）从 proto 中读取 has_communication，并查询 MosaicMlirCacheEntry::EmitsSdcChecksums（→ supports_internal_checksums）和 MosaicMlirCacheEntry::RequiresMxuAssigner（→ 同时设置 requires_mxu_assigner 与 check_fifos_are_empty），同时固定 supports_hlo_dedup = true 和 instruction_can_change_layout = true。参见 Mosaic 导入接缝。

注册表播种 — `google_init_module_*`

每个 handler 都携带一个通过链接器 .init_array（不是 __attribute__((constructor))）连线的 google_init_module_<handler>() 函数，每个函数都由 google_initializer_module_<handler> BSS 标志保护。首次触达时，这些函数会把它们的 (target_string, callback) 对插入到各 facet 注册表中。已恢复的初始化器（列出提取到的地址）：

Handler init 函数	地址	播种的目标
`google_init_module_custom_call_emitter`	`0x213ec9e0`	`tpu_custom_call`（Mosaic 逃逸口）
`google_init_module_alloc_handler`	`0x213ed4e0`	`AllocateBuffer`
`google_init_module_assume_handler`	`0x213ed5e0`	`AssumeGatherIndicesInBound`
`google_init_module_qr_handler`	`0x213edd80`	`QrDecompositionBlock`, `CompactWyHelper`, `InvertDiagBlocks*`
`google_init_module_resize_handler`	`0x213edee0`	`ResizeBilinear[Grad]`, `ResizeNearest[Grad]`
`google_init_module_topk_handler`	`0x213ee300`	`TopK`, `TopKWithUnique`
`google_init_module_x64_handler`	`0x213ee500`	`X64Combine`, `X64SplitLow/High`
`google_init_module_sliceid_handler`	`0x213ee260`	`SliceId`
`google_init_module_xla_llo_log_emitter`	`0x213ed920`	`kTpuLogCustomCallTarget`
`google_init_module_xla_sdc_checker_emitters`	`0x213ed220`	`xla-sdc-checker-get-stats`, `xla-sdc-checker-report-sdc-event`, `xla-sdc-checker-ici-sdc-test`, `xla-sdc-checker-start-with-alt-cores`
`megascale_custom_call_handler::google_init_module_…`	`0x213ed440`	`xla.megascale.provide_metadata`

另有每个 emitter 的 google_init_module_*_emitter 初始化器（partial_reduce_emitter, window_prefetch_emitter, mosaic_broadcast_emitter, cholesky_emitter, qr_emitter, eigh_emitter, lu_emitter, invert_diag_blocks_emitters, compact_wy_emitter, padding_emitters, barrier_start_emitter, async_collective_{start,done}_emitter, barna_core_address_handler_emitter, topk_batch_major_small_k_emitter）。重新实现必须保证在 HLO→LLO 循环运行前注册表已完全填充 — .init_array 顺序会在进程启动时、任何编译之前完成这一点。

分派字符串格式

HloCustomCallInstruction 上的 custom_call_target 字段是普通的 std::string；查找是针对注册表的字符串相等性匹配。三种命名约定并存：

约定	含义	示例
裸 CamelCase	内建 TPU 原语 lowering	`Cholesky`, `QrDecompositionBlock`, `EighTpu`, `LuDecompositionBlock`, `TopK`, `TopKWithUnique`, `ResizeBilinear`, `Pin`, `Unpin`, `WindowPrefetch`, `AllocateBuffer`, `PadToStatic`, `SliceToDynamic`, `X128Combine`, `X64Combine`, `MaskAggregatorBlock`, `AssumeGatherIndicesInBound`, `MoveToHost`, `MoveToDevice`, `Sharding`, `SPMDFullToShardShape`, `BarrierStart`, `InspectSharding`
lowercase_underscore	通用分派容器	`tpu_custom_call`（`backend_config` 中的 Mosaic body）、`single_tpu_custom_call`（legalizer sentinel）、`recover_custom_call`（debug-replay）、`annotate_device_placement`（设备放置标记）
点分 namespace	SDY / MegaScale 逃逸口	`xla.sdy.Sharding`, `xla.sdy.ShardingGroup`, `xla.sdy.FuncResultSharding`, `xla.sdy.GlobalToLocalShape`, `xla.sdy.LocalToGlobalShape`, `xla.sdy.PropagationBarrier`, `xla.megascale.provide_metadata`

有两条解析规则已经锚定：

xla.sdy. 前缀检测使用 hlo->custom_call_target().rfind("xla.sdy", 0)（前缀测试）。该族目标会通过 Shardy importer 往返。
以 $ 为前缀的名称是保留的。 已恢复的 DCHECK（反编译输出逐字）：Invalid custom_call_target "%s": Call targets that start with '$' are reserved for internal use. 这是拒绝 __cudnn$convForward 和 __triton$… 等跨后端逃逸口的门禁 — 它们会通过开源 HloVerifier，但会死在 TpuHloSupportChecker 中（参见验证层与错误路径）。

Custom-Call 目标目录

约 52 个不同的高置信度目标字符串。已分类；每行列出注册它的 namespace 常量（恢复为 k… rodata 锚点）以及降低动作。动作代码 (A)–(F) 在六种降低动作中定义。

分片标记

目标	通过何者注册	动作 / lowering
`Sharding`	`sharding_handler::kSharding`	(B) 被 `ShardingPropagation` 消费；括号由 `HloDomainRemover("sharding", …)` 移除
`SPMDFullToShardShape`	（无 handler；SpmdPartitioner）	(F) global→per-shard 形状边界；在分区后的图中变成 `Reshape`
`SPMDShardToFullShape`	（无 handler；SpmdPartitioner）	(F) 反向；如果不存在 sharding annotation 则报错

SDY 往返（Shardy）

目标	通过何者注册	动作 / lowering
`xla.sdy.Sharding`	`SdyCustomCallPattern`	由 `ImportSdyCustomCallsPass` 重写为 `mlir::sdy::ShardingConstraintOp`
`xla.sdy.ShardingGroup`	`SdyCustomCallPattern`	→ `mlir::sdy::ShardingGroupOp`；导入后必须没有 use
`xla.sdy.FuncResultSharding`	`getFuncResultSharding`	输出 sharding 载体；剥离到 FuncOp result attr 上
`xla.sdy.GlobalToLocalShape`	`kGlobalToLocalShapeCallTargetName`	`sdy::ManualComputationOp` 往返的 reshape 边界
`xla.sdy.LocalToGlobalShape`	`kLocalToGlobalShapeCallTargetName`	上述动作的反向
`xla.sdy.PropagationBarrier`	PropagationBarrier importer	→ `mlir::sdy::PropagationBarrierOp`；需要 `allowed_direction` attr
`InspectSharding`	`RemoveInspectShardingCustomCall`	(B) 无条件移除（JAX `inspect_array_sharding`）；通过 `RegisterCustomCallPartitioner("InspectSharding")` 注册为 partitioner

内存放置

目标	通过何者注册	动作 / lowering
`MoveToHost`	`memory_annotations::kMoveToHostCustomCallTarget`	(B) `HostOffloader::HandleMoveToHostCustomCall` 重置 `memory_space`，删除该 call
`MoveToDevice`	`memory_annotations::kMoveToDeviceCustomCallTarget`	(B) 反向（`HandleMoveToDeviceCustomCall`）
`Pin`	`memory_annotations::kPinToDeviceCustomCallTarget`	(C) tensor→memref pin；"Pin custom_call should have a memref output"
`Pin` (SRAM)	`memory_annotations::kPinToDeviceSramCustomCallTarget`	(C) 强制 SRAM/VMEM 放置
`Unpin`	memory_annotations_handler	(C) memref→tensor
`annotate_device_placement`	`memory_annotations::kDevicePlacement` / `device_placement_handler::kDevicePlacement`	强制设备放置；`instruction_can_change_layout=false`（`kDevicePlacement` 常量解析为字符串 `"annotate_device_placement"`）

线性代数块原语

目标	通过何者注册	动作 / lowering
`Cholesky`	`cholesky_handler::kCholesky`	(A) 如果已预展开则为 `TpuCholeskyExpander`（dot+TRSM）；否则 (C) emit-helper
`QrDecompositionBlock`	`qr_handler::kQrDecompositionBlock`	(C) per-block QR（Givens），`emit_helper`
`CompactWyHelper`	`qr_handler::kCompactWyHelper`	(C) blocked QR 的 compact-WY builder
`InvertDiagBlocksLowerTriangular`	`qr_handler::k…LowerTriangular`	(C) triangular-solve expander helper
`InvertDiagBlocksUpperTriangular`	`qr_handler::k…UpperTriangular`	(C) upper-triangle 变体
`EighTpu`	eigh_emitter	(C) block-Jacobi 特征分解，`emit_helper`
`LuDecompositionBlock`	lu_emitter	(C) per-block LU，`emit_helper`
`MaskAggregatorBlock`	anonymous emit_helper	(C) per-block mask aggregator（attention masking）

排序 / 选择 / 规约

目标	通过何者注册	动作 / lowering
`TopK`	`topk_handler::kTopk`	(C) emit + (CanFuse=`true`) + cost + properties
`TopKWithUnique`	`topk_handler::kTopkWithUnique`	(C) 同族
`TopKBatchMajorSmallK`	`topk_batch_major_small_k_handler::kTopkBatchMajorSmallK`	(C) 专用；CanFuse=`false`（逐字恢复）
`ApproxTopK`	open-source XLA	verifier：operand 数为偶数，恰有 1 个 called_computation
`PartialReduce`	`partial_reduce_handler::kPartialReduce`	(C)+(F) `PartialReduceEmitter`；CanFuse=`true`；已注册 cost analysis

内存 / 动态形状 / RNG / 图像

目标	通过何者注册	动作 / lowering
`AllocateBuffer`	`alloc_handler::kAllocateBuffer`	(C)+(F) 静态 buffer 预留；memref result，无输入
`WindowPrefetch`	`WindowPrefetchEmitter::kWindowPrefetch`	(C) `WindowPrefetchEmitter::Emit` @ `0x10f93f60`；`supports_hlo_dedup=true`
`PadToStatic`	`dynamic_padding_handler::kPadToStatic`	(C) `static_padding_emit_helper`
`SliceToDynamic`	`dynamic_padding_handler::kSliceToDynamic`	(C) `dynamic_padding_emit_helper`
`ResizeBilinear` / `Grad`	`resize_handler::kResizeBilinear[Grad]`	(C)+(F) 每个变体都有自定义 HLO cost analysis
`ResizeNearest` / `Grad`	`resize_handler::kResizeNearest[Grad]`	(C)+(F) `ResizeNearestHloCostAnalysis`
`RngBitGenerator`（HLO opcode，不是 custom-call target）	`TpuRngBitGeneratorExpander`	(A) `kRngBitGenerator` opcode 由 expander 重写为 Philox/ThreeFry HLO；列在这里仅因为它有时会被误认为 custom-call target — 它永远不会到达注册表

Runtime intrinsic / 精度 / async / host / SDC

目标	通过何者注册	动作 / lowering
`DeviceId`	`deviceid_handler::kDeviceId`	(C) → `tpu.device_id`
`SliceId`	`sliceid_handler::kSliceId`	(C) 返回 slice id（MegaScale-aware）
`AssumeGatherIndicesInBound`	`assume_handler::k…`	(B) propagation 后由 `TpuAlgebraicSimplifier` 移除；已注册 cost analysis
`X64Combine`	`x64_handler::kX64Combine`	(C) 合并 64-bit 值的 hi+lo
`X64SplitLow` / `X64SplitHigh`	`x64_handler::kX64Split{Low,High}`	(C)+(F) SPMD visitor，因此每个 shard 拥有自己的一半
`X128Combine`	X128 emit helper	(C) `XPrecisionRewriter(kX128Precision)` 的 128-bit combine
`PrepareAsyncCallStart`	`PrepareAsyncCallInserter::kPrepareAsyncCallStartTarget`	(C) → `tpu.async_start`
`PrepareAsyncCallDone`	`PrepareAsyncCallInserter::kPrepareAsyncCallDoneTarget`	(C) async completion marker
`kAsyncCollectiveStart` / `Done`	`async_collective_fusion_util::k…`	(C) async-collective 对；`emit_helper`
`kBarrierStart`	`async_barrier_util::kBarrierStart`	(C) barrier kickoff；`emit_helper`
`kDcnAllReduceStart`	DCN all-reduce kickoff	通过 `IsCustomCall("kDcnAllReduceStart")` 匹配
`HostExecute`	host_callback flow	(E) host execute；必须恰有一个 called computation
`xla.megascale.provide_metadata`	`runtime::kXlaMegaScaleCustomCallMetadataName`	(E) MegaScale collective 元数据；`instruction_can_change_layout=false`
`xla-sdc-checker-start-with-alt-cores`	`SdcCheckerStartWithAlternativeCoresEmitter::kKey`	(C)/(E) SDC checker start
`xla-sdc-checker-ici-sdc-test` / `xla-sdc-checker-get-stats` / `xla-sdc-checker-report-sdc-event`	`IciSdcTestEmitter`/`SdcCheckerGetStatsEmitter`/`SdcCheckerReportSdcEventEmitter::kKey`	(C)/(E) SDC instrumentation
`kTpuLogCustomCallTarget`	`LogEmitter::kTpuLogCustomCallTarget`	(C) `OpEmitter::Emit<LogEmitter>`（已恢复 lambda）

Mosaic 逃逸口

目标	通过何者注册	动作 / lowering
`tpu_custom_call`	`CustomCallEmitter` + `MosaicMlirCacheEntry`	(D) 逃逸口 — `backend_config` 携带序列化的 `tpu` dialect MLIR 模块；见下文

已识别但被静默消费的目标（由另一个 pass 重写为 CustomCall 内部标记，然后被移除/展开）：tf.XlaCallModule（MLIR-op-backed module-call 载体）、single_tpu_custom_call（legalizer sentinel）、recover_custom_call（debug-replay provenance）。键 xla.sdy.in_shardings / xla.sdy.out_shardings 是前端属性名，不是目标字符串。

明确拒绝（永远不会到达注册表）：任何 __<vendor>$… 目标 — 由 TpuHloSupportChecker 按保留 $ 前缀过滤。

六种降低动作

在整个目标目录中，分派会解析为六种动作类别之一。单个目标可以采用多个类别（例如 PartialReduce 同时是 (C) 和 (F)）。

代码	动作	何时应用	示例目标
(A)	在 pre-pass 阶段重写为其他 HLO	HLO PreOptimization	`Cholesky`→`TpuCholeskyExpander`, `RngBitGenerator`, Qr/Eigh/Lu/TriangularSolve/Fft opcodes
(B)	消费后剥离标记	HLO mid-pipeline	`Sharding`→`ShardingPropagation`, `InspectSharding`→`RemoveInspectShardingCustomCall`, `MoveTo*`→`HostOffloader`
(C)	通过 emit-helper 降低到 MLIR/LLO	HLO→LLO 降低循环	`TopK`, `PartialReduce`, `Resize`, `AllocateBuffer`, `WindowPrefetch`, `PadToStatic`, `SliceToDynamic`, `DeviceId`, `SliceId`, `X64`, linalg blocks
(D)	通过缓存的 MLIR module body 降低	Mosaic 路径	`tpu_custom_call`（唯一目标）
(E)	作为 host-action 标记保留到 runtime	runtime emission	`xla.megascale.provide_metadata`, `HostExecute`, `xla-sdc-checker-*`
(F)	仅 SPMD 分区	SpmdPartitioner	`SPMDFullToShardShape`, `SPMDShardToFullShape`, `X64Split`, `PartialReduce`, `Resize`, `AllocateBuffer`

(A) Pre-pass 重写

此类目标永远不会到达注册表的 emit-helper — HLO pre-pass 会先把它们展开为普通 HLO。（参见 HLO Pre-Passes。）

目标 / opcode	Pre-pass 类
`Cholesky`	`xla::TpuCholeskyExpander`
QR opcode	`xla::TpuQrExpander`
`Eigh` opcode	`xla::TpuEighExpander`
`Lu*`	`xla::LuDecompositionExpander`
`TriangularSolve` opcode	`xla::TpuTriangularSolveExpander`
`Fft` opcode	`xla::FftExpander`
`RngBitGenerator`	`xla::jellyfish::TpuRngBitGeneratorExpander`
`Gather` / `Scatter`	`xla::TpuGatherExpander` / `xla::TpuScatterExpander`
`RaggedAllToAll`	`xla::jellyfish::RaggedAllToAllExpander`

xla::jellyfish::MosaicFusion 也在此阶段运行，但它不是由 custom-call target 触发的 — 它把普通 HLO 子图提升为符合 Mosaic 条件的 kernel（发射新的 tpu_custom_call 指令），因此在注册表目录中没有条目。

(B) 标记剥离

ShardingPropagation 消费 Sharding/SPMDFullToShardShape/SPMDShardToFullShape，发射 kDomain 括号，然后 HloDomainRemover("sharding", ApplyDomainSharding) 在保留 sharding 属性的同时删除括号（Sharding Propagation）。
RemoveInspectShardingCustomCall::RunImpl（0x1278a040）擦除每个 InspectSharding，转发 operand。
HostOffloader::HandleMoveToHostCustomCall（0x110778a0）/ HandleMoveToDeviceCustomCall（0x11078b00）重置 producer memory_space 并删除该 call。

(C) Emit-helper 路径 — 注册表默认路径

这是常见情况：HLO→LLO visitor 看到一个 kCustomCall，在 lowering-emitter 注册表中查找其目标并调用回调；该回调向 LloRegion 中发射 mlir::llo::*（以及辅助性的 mlir::tpu::*）op，并返回结果 OperandData。多数回调会包装一个领域特定的 OpEmitter 子类。两个已恢复的 lambda body 展示了规范形状：

cpp

// PartialReduce — narrow overload (no SDC/log context)
RegisterLoweringEmitter(partial_reduce_handler::kPartialReduce,
  [](const HloInstruction* hlo, const LoweredGetter& get_lowered,
     LloRegion* region, LloValue* output, BackendConfigMap* backend_config_map)
       -> absl::StatusOr<OperandData> {
    auto emitter = PartialReduceEmitter(hlo, get_lowered, region, output);
    TF_ASSIGN_OR_RETURN(auto return_value, emitter.Emit());
    return OperandData::CreateStaticOperandData(return_value, hlo->shape(), region);
  });

// LogEmitter — wide overload (uses log_recorder from the context tail)
RegisterLoweringEmitter(LogEmitter::kTpuLogCustomCallTarget,
  [](const HloInstruction* hlo, LoweredGetter get_lowered, LloRegion* llo_region,
     LloValue* output, BackendConfigMap* backend_config_map,
     const ProgramSharedRegistry* registry, const LogicalTopologyInfo* topology_info,
     sdc_reporter::SdcRegistrator* sdc_registrator, const llo_log::LogRecorder* log_recorder)
       -> absl::StatusOr<OperandData> {
    TF_ASSIGN_OR_RETURN(auto ret_val,
        OpEmitter::Emit<LogEmitter>(hlo, get_lowered, llo_region, log_recorder));
    return OperandData::CreateStaticOperandData(ret_val, hlo->shape(), llo_region);
  });

二者结尾完全相同（OperandData::CreateStaticOperandData(value, hlo->shape(), region)）；差异仅在于 emitter 消费 context 尾部的哪一段。

(E) Runtime / host-action 目标

这些目标会作为 sentinel op 留过 HLO→LLO，只在 executable-emission 时才重写为 HostCommand（或 MegaScaleAction）。它们通过共享的上游 xla::cpu/gpu::CustomCallThunkProto 序列化为 Thunk::Kind::kCustomCall 条目（已确认：存在 CustomCallThunkProto::_InternalSerialize/MergeImpl/ByteSizeLong）。TPU 侧复用该 proto，但通过 xla::megascale::runtime::CustomCallHandler（HandleCollectiveInput 0x1cbd3a00，HandleCollectiveOutput 0x1cbd4c20，AddPendingChain/GetPendingChain/ActiveGraph）执行，而不是使用上游 thunk。

(F) 仅 SPMD 分区

对于通过 RegisterSpmdPartitioningVisitor 注册的目标，partitioner 会调用已注册的 visitor，而不是默认 handler。已恢复的注册（逐字）：kAllocateBuffer, kPartialReduce, kResizeBilinear[Grad], kResizeNearest[Grad]，以及一个 generic target（uniform partitioner）。两个 SPMD helper 也会消费 custom-call：TpuCustomCallShardingHelper（实现 xla::CustomCallShardingHelper：InferShardingFromOperands 0x1278bf80、PropagateUserSharding、IsCustomCallShardable、CanPropagateShardingToOperands）和 TpuLogCustomCallPartitioner（replicated-only xla-log sharding）。参见 Auto-Sharding 与 SPMD。

Mosaic 导入接缝

tpu_custom_call 是唯一采用动作 (D) 的目标，也是整个 custom-call 机制的主要用途：它是 Pallas/Mosaic kernel — 由 JAX 前端在 libtpu 之外写成 tpu dialect MLIR — 进入编译器的方式。tpu dialect 从不由通用 MHLO lowering 产生；它只会通过这个接缝被导入。（下游 tpu→LLO pipeline 记录在 MHLO → XTile → tpu 降低和 Mosaic 概览中；本节只覆盖 custom-call 入口与缓存。）

一个 tpu_custom_call 指令携带：

custom_call_target = "tpu_custom_call"
backend_config = 序列化的 xla.jellyfish.BackendConfig，其 custom_call_config（xla.jellyfish.CustomCallConfig）持有 kernel。

已确认存在 CustomCallConfig proto（CustomCallConfig::{ByteSizeLong,Clear,CopyFrom,GetClassData}、arena ctors，以及重复嵌套 submessage CustomCallConfig_InputMemorySpaceColor / CustomCallConfig_OutputMemorySpaceColor 与它们的 RepeatedPtrFieldBase::Add / Arena::CopyConstruct 实例化）。字段集合（名称已恢复；tag number 未提取 — 参见置信度总结）：

protobuf

message CustomCallConfig {
  bytes  mlir_module;          // serialized tpu-dialect module (bytecode or textual)
  string host_mlir_module;     // host computation used for shape inference; "" if shapes static
  string kernel_name;          // recovered from the ",kernel_name=" descriptor fragment
  repeated InputMemorySpaceColor  input_memory_space_colors;   // per operand
  repeated OutputMemorySpaceColor output_memory_colors;        // per result (accessor confirmed; tuple-arity must match)
  bool   has_communication;    // survives async-fusion as a comm boundary if true
  int64  collective_id;        // required for tpu.get_barrier_semaphore in the body
  CustomCallCost cost_estimate;        // {flops, transcendentals, bytes_accessed} — proto type confirmed
}

input_memory_space_colors / output_memory_colors 字段名对是从 descriptor 字符串和 custom_call_config.output_memory_colors() accessor 调用点恢复的；这种非对称命名（只有 input 带 _space_）与二进制中观察到的一致，不是转录错误。submessage 类型是 CustomCallConfig_InputMemorySpaceColor / CustomCallConfig_OutputMemorySpaceColor。独立的 metadata 字段未确认；metadata 是通过指令上的 SetCustomCallMetadata 附加的，而不是通过命名 proto 字段附加的。

Pipeline：HLO 入口 → 缓存 body → LLO

HLO 入口。 前端为每个 pl.kernel 发射一个 kCustomCall("tpu_custom_call")，并填充 custom_call_config.mlir_module。
MosaicFusion（RunImpl 0x10f12500，包在 HloPassFix 中）把周围 HLO 子图提升为符合 Mosaic 条件的 kernel，直到固定点。
TpuCustomCallLegalizer::RunImpl（0x11036080）通过 ConfigureSparseCoreConfig（0x110355e0）/ ConfigureMegachipParallelism（0x11035a20）将每个 call 分类为 TensorCore、SparseCore 或 Megachip；SparseCore kernel 会通过 OffloadOneSparseCoreCustomCall（0x11035700）就地 offload。参见降低到 SparseCore LLVM。
TpuCustomCallMemorySpacePolicy::RunImpl（0x110364a0）通过 RunHbmPolicy（0x11038120）或 RunMsaReservationPolicy（0x110367c0）分配 input_memory_space_colors/output_memory_colors，由 --xla_tpu_tpu_custom_call_memory_space_spec proto 驱动。参见 MSA 预留 / HBM 策略。
TpuCustomCallScopedVmemAdjuster::RunImpl（0x1104de40）运行试验性 BufferAssignment（通过构造时捕获的 absl::AnyInvocable<StatusOr<…BufferAssignment…>(HloModule*, Target const&)>），并把 scoped-VMEM byte count 重写为实际需求。
HLO→LLO lowering 通过 CustomCallEmitter::Emit（0x111ef740）进行：
- GetCustomCallAndConfig(HloInstruction const*)（0x111fe020）提取 (HloCustomCallInstruction*, CustomCallConfig) 对。
- GetCachedCustomCallBody(HloCustomCallInstruction const*, CustomCallConfig const&)（0x13e31860）返回 MosaicMlirCacheEntry；cache miss 时通过 GetMlirModule(CustomCallConfig const&, MLIRContext&, bool)（0x13e31220）→ ParseMlirModuleString（0x0f908580）解析 config.mlir_module()。该 entry 存储在 HloModule 上（HloModule::SetCacheEntry<MosaicMlirCacheEntry>），以 tsl::Fprint128 为键，因此相同 kernel 只解析一次。
- SetupAsyncCollectiveLowering（0x111f7520）在 has_communication=true 时运行，把该 call 融合进 async collective。
- GetDeviceAssignment（0x111f4ee0）从父 computation 解析设备列表。
MosaicMlirCacheEntry introspection — 缓存 entry 暴露 HasAnyCoreType(Span<mlir::tpu::CoreType const>)、EmitsSdcChecksums、IsSparseCoreKernel、RequiresMxuAssigner、GetCacheKey（tsl::Fprint128）。这些会供给 tpu_custom_call 的动态 RegisterCompilationProperties 回调。
导入的 tpu module 随后运行标准 tpu dialect pipeline（serde version upgrade、layout inference、createLowerToLLOPass）— 记录在 Mosaic 概览中。

MosaicMlirCacheEntry 构造函数参数已精确恢复：(std::string&, std::unique_ptr<mlir::MLIRContext>, mlir::OwningOpRef<mlir::ModuleOp>, MosaicMlirCacheEntry*) — 即解析后的 module 及其 owning context，传入缓存。

陷阱 — cache key 是 kernel bytecode，而不是 HLO 指令。 因为 key 是基于 MLIR body 的 tsl::Fprint128，两个共享同一 kernel 的不同 tpu_custom_call 指令会折叠为一次解析和一个已编译 body。若重新实现按 HLO 指令指针作为 key，则会重复解析相同的 Pallas kernel，并错失 dedup。

验证层与错误路径

三个层次会按 pipeline 顺序验证 custom-call：

HloVerifier — 针对 kCustomCall 的通用 shape/operand-count invariant（operand 数与 operand_shapes_with_layout.size()，layout 存在性）。每个 pre-pass 后运行。
xla::TpuHloSupportChecker — TPU 接受性测试。拒绝任何不在注册表中的目标；__cudnn$convForward、__triton$… 和其他以 $ 为前缀的跨后端逃逸口就是在这里死亡的。
TpuCustomCallLegalizer::RunImpl — 语义合法化（memory-space coloring sanity、SparseCore offload feasibility、megachip parallelism）。

代表性的已恢复 status/DCHECK 字符串（逐字来自 .rodata / 反编译输出），按类别列出：

类别	锚定字符串
未知目标	`Attempting to lower unimplemented custom call %s` · `Custom call target %s is not implemented.` · `Unimplemented custom-call: %s`
保留名称	`Invalid custom_call_target "%s": Call targets that start with '$' are reserved for internal use.`
错误 arity	`Custom calls with target %s must have exactly one operand. %s has %d.` · `%s custom call must have no operands.` · `%s custom call must have an S32 scalar shape.`
Mosaic body	`Failed to parse custom call kernel: …` · `Custom call does not have a custom_call_config field.` · `Cannot lower tpu.get_barrier_semaphore op because a barrier config was not provided. Perhaps, collective_id was not set in the TPU custom call.` · `Custom call body is empty.`
Pallas 输出	`Pallas custom calls must have array or tuple output. Found: …` · `Pallas custom calls with tuple outputs must have exactly one output memory space color per tuple element.`
Scoped VMEM	`Failed to determine the scoped vmem requirement for one or more custom calls in module …` · `… is an unsupported memory space in TpuCustomCallScopedVmemAdjuster. …`
Pin / Unpin	`Pin custom_call should have a memref output` · `Unpin custom_call should have a tensor output` · `custom-call to Pin must have one output-to-operand aliasing`
AllocateBuffer	`CreateBuffer custom_call should have a memref result` · `custom-call to CreateBuffer can't have an operand`
SDY 往返	`expected CustomCallOp with xla.sdy target name.` · `xla.sdy.ShardingGroup CustomCallOp should have no uses.` · `An SPMDShardToFullShape custom call found without a sharding annotation.`
ApproxTopK	`ApproxTopK takes an even number of operands.` · `ApproxTopK takes exactly 1 called_computation.`
Host execute	`Host execute custom call must have exactly one called computation.`
MegaScale	`MXLA custom call does not support collectives on multiple channels.` · `Host offloaded collective custom call '$0' only works in multi slice environment.`
Backend-config parse	`Unable to parse backend config for custom call: …` · `Failed to parse WindowPrefetch backend config.` · `Failed to register custom call partitioner for …`

编译期标志

标志（`FLAGS_xla_tpu_…`）	作用
`tpu_custom_call_memory_space_spec`	驱动第 4 步的序列化 `TpuCustomCallMemorySpaceSpec` proto
`enable_tpu_custom_call_scoped_vmem_adjustments`	控制 `TpuCustomCallScopedVmemAdjuster`
`custom_call_nop_return_token_vdelay`	nop-return-token call 的延迟周期
`mosaic_fusion`	控制 `MosaicFusion` pre-pass
`enable_mosaic_emitters`	控制 Mosaic emitter 注册表
`enable_async_collective_fusion_with_mosaic_custom_call`	控制 `SetupAsyncCollectiveLowering`
`sdc_checker_{checksum,enable,timestamp,zero}_custom_call*`	`tpu_custom_call` 的 SDC instrumentation
`layer_scheduler_min_custom_call_flops`	scheduler FLOPs 阈值
`llo_race_analysis_analyze_tpu_custom_calls`	custom-call body 上的 LLO race analysis

置信度总结

断言	证据
五 facet 注册表及列出的 `Register*` 入口点/签名	`FunctionRegistry<std::string, CompilationProperties(HloInstruction const)>::Register` + `RegisterCanFuse`/`RegisterSpmdPartitioningVisitor` demangled signatures；`.rodata` 中的 `Register` 源码引用
`LoweringEmitter` 回调签名，包括 `BackendConfigMap = flat_hash_map<HloInstruction const*, unique_ptr<BackendConfig>>`	从 `__policy_func` thunk + 逐字 lambda body 恢复
`CompilationProperties` 六字段结构体 + 每目标取值	designated-initializer 源码引用（`kSharding`, `WindowPrefetch`, SDC keys, …）
约 52 个目标的目录及 handler 常量	`k…` rodata 锚点 + 每个 handler 的 `google_init_module_*` 符号；每个逐字字符串均确认存在于二进制中
Mosaic 导入接缝（`GetCachedCustomCallBody`/`GetMlirModule`/`MosaicMlirCacheEntry`/`CustomCallEmitter::Emit`）	所有符号都存在于列出的地址；以 `Fprint128` 为键的 `HloModule::{Set,Get}CacheEntry<MosaicMlirCacheEntry>`；ctor 参数列表已恢复
三个合法化 pass + 签名	`TpuCustomCallLegalizer`/`MemorySpacePolicy`/`ScopedVmemAdjuster` `RunImpl` 和 helper signature 已 demangle；ctor `AnyInvocable<…BufferAssignment…>` 已恢复
Error/DCHECK 字符串	reserved-`$` 和 unimplemented 字符串已在反编译输出中确认；其余来自 `.rodata` quoted source
`CustomCallConfig` proto 字段 tag number	descriptor 存在（submessage 名称、accessor），但 tag number 未从 FileDescriptorProto 提取
每个 handler 的 emit-helper LLO body（精确 op 序列）	handler→target 映射为高；`OpEmitter` 子类 LLO 重写未按 op 反编译
`MosaicMlirCacheEntry::GetCacheKey` 字段组成	返回基于 body 的 `tsl::Fprint128`；`kernel_name`/`cost_estimate` 是否参与未恢复
`xla.sdy.GlobalToLocalShape` per-axis reshape 协议	target-name 常量为高；flatten/pack/relabel 语义未反汇编

交叉引用

TPU 编译器（概览） — 五阶段主干；custom-call 分派位于设备路径的 HLO→LLO lowering 循环中。
编译阶段 0–3 — (A) pre-pass expander 和 (B) marker-strip pass 相对于 HLO→LLO emit 循环的运行位置。
HLO Pre-Passes — 在注册表看到 action-(A) 目标前消费它们的 expander 家族（TpuCholeskyExpander, TpuRngBitGeneratorExpander, MosaicFusion, …）。
MHLO → XTile → tpu 降低 — 证明 tpu dialect 通过 tpu_custom_call 被导入、而绝不会由 MHLO lowering 产生的接缝；GetMlirModuleOpFromCustomCall / RunMLIRPasses。
Mosaic 概览 — 在 (D) cached-body 提取后运行的 imported-kernel tpu dialect pipeline（serde、layout inference、lower-to-LLO）。
降低到 SparseCore LLVM — TpuCustomCallLegalizer 的 SparseCore offload（OffloadOneSparseCoreCustomCall）目的地。
tpu MLIR Dialect — 由上游编写并由 CustomCallEmitter 内联的 dialect。
Sharding Propagation — (B) sharding 标记和 SDY 往返目标的 consumer。
Auto-Sharding 与 SPMD — 通过 RegisterSpmdPartitioningVisitor 和 TpuCustomCallShardingHelper 消费 action-(F) 目标的 SpmdPartitioner。
MSA 预留 / HBM 策略 — TpuCustomCallMemorySpacePolicy 调用的 memory-space coloring 策略。
Ragged-Dot 与 Convolution 降低 — 与 custom-call linalg block 相邻的 op 的相关 jellyfish emitter。
二进制： extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d）
索引条目： Part V — Compiler: Lowering & Optimization Passes / Front-end and pipeline — 返回索引

Custom-Call 降低与目标注册表 ​

摘要 ​

五 Facet 注册表 ​

CompilationProperties — 声明式元数据结构体 ​

注册表播种 — google_init_module_* ​

分派字符串格式 ​

Custom-Call 目标目录 ​

分片标记 ​

SDY 往返（Shardy） ​

内存放置 ​

线性代数块原语 ​

排序 / 选择 / 规约 ​

内存 / 动态形状 / RNG / 图像 ​

Runtime intrinsic / 精度 / async / host / SDC ​

Mosaic 逃逸口 ​

六种降低动作 ​

(A) Pre-pass 重写 ​

(B) 标记剥离 ​

(C) Emit-helper 路径 — 注册表默认路径 ​

(E) Runtime / host-action 目标 ​

(F) 仅 SPMD 分区 ​

Mosaic 导入接缝 ​

Pipeline：HLO 入口 → 缓存 body → LLO ​

验证层与错误路径 ​

编译期标志 ​

置信度总结 ​

交叉引用 ​