# TC 编译来源映射来源映射为程序映像中每个 (PC, 物理槽) 关联三类信息：Pallas 源码位置、所属函数区间，以及编译器原生的注释。这里的 PC 是 bundle 在程序映像中的序号，不是设备 IMEM 地址。支持 `tpu-v4-tc` 与 `tpu-v6e-tc`，两者共用补丁和恢复流程，只有槽名和 overlay 坐标的单位不同。功能分为两个阶段： - **捕获**：在 `compiler_source_mapping()` 上下文中编译时进行，需要对 libtpu 打进程级补丁，补丁与 libtpu 版本绑定。 - **恢复**：`executable_source_maps()` 只读取 serialized executable。除匹配的编解码后端外，不需要 TPU、补丁或任何 dump。本功能不做以下事情：不解析 final bundles 或 LLO dump；不根据 `vtrace` 的 PC 差或指令类别猜测来源；来源未知时保持未知，不能把“没有来源”解读为“由编译器生成”。 ## 为什么借用 libtpu 的逐槽注释 libtpu 本身就会为每条 LLO 指令携带一段注释字符串，经过调度、组包和链接，最后按 bundle 和物理槽写进 `CompilerMetadata`。tpuasm 的做法是在指令发射前，把一条自包含的记录追加到这段注释里。此后由 libtpu 自己的 emitter 决定记录落到哪个槽，由它的 linker 重定位 PC。这样，tpuasm 不需要重新实现槽分配和重定位，而这两部分正是最容易随编译器版本变化出错的地方。补丁有一条不变式：**不改变生成的机器码。** [tests/reproduce_tpu_v4_tc_source_mapping.py](../../tests/reproduce_tpu_v4_tc_source_mapping.py) 对每个案例分别在打补丁和不打补丁的情况下编译，要求两份程序映像相同，且数值结果一致。 ### 记录格式 v1 追加的文本为 `[[tpuasm:v1:]]`，hex 编码一条独立的 protobuf wire 消息： | 字段 | 类型与含义 | |---|---| | 1 | varint，最终的 LLO ordinal | | 2 | bytes，筛选后的 `SourceMapProto` | | 3 | string，HLO instruction 名 | | 4 | string，HLO module 名 | | 5 | varint，HLO module ID | 字段 2 只保留 ordinals 中含有本指令 ordinal 的 location，并把这些 location 的 ordinals 缩减为本指令的 ordinal 一项，但保留完整的字符串表。因此一条记录可以单独解析，不依赖已经销毁的 LLO module。ordinals 必须缩减：一个 location 列出由它 lower 出的全部指令，完全展开的大运算（例如 512 行的 `jnp.dot`）中这张表有上万项，原样保留时每条记录约 31 KB，注释表随指令数平方增长；一个 17782 个 bundle 的 kernel 的 `annotation_metadata` 达 2.1 GB，再大一些就超过 protobuf 的 2 GiB 上限，executable 无法序列化。缩减后同一程序为 37 MB。所用的 SourceMap 字段为：SourceMap 的 locations=1、strings=2；SourceInfo 的 frames=1、primitive=2、scope_stack=3、ordinals=4；Frame 的 path=1、line_start=2、line_end=3、col_start=4、col_end=5、function_name=6。同一条注释中已有 v1 记录时不再追加。修改记录的字段需要使用新的版本标签，恢复端继续解析 v1。 ## 捕获 ### JAX lowering Pallas 的 `jaxpr_subcomp` 用 `ctx.name_stack + eqn.source_info.name_stack` 生成每个方程的 MLIR location。lowering 规则递归调用 `jaxpr_subcomp` 处理子 jaxpr 时，传入的 `ctx.name_stack` 不一定包含外层方程的 scope，于是位置信息丢失了静态的父 scope。[tc_source_lowering.py](../../src/tpuasm/tc_source_lowering.py) 用一个 `ContextVar` 保存外层方程的完整 name stack，作为嵌套调用的父栈；外层方程的 traceback 也以同样方式接在内层 traceback 之后，使内层 op 的位置带上外层调用点。改动只影响 location 参数，不改变 name stack 的 trace 操作。第二处改动针对完全静态展开的循环（例如 `pl.loop(..., unroll=True)`）。JAX 的 `_lower_jaxpr_to_for_loop` 会把循环体 lower 成一个游离的 `func.func`，再用 `jax_mlir_ext.inlined_func_call` 逐次克隆到调用处。克隆时 JAX 把每个 op 的 location 改写为以循环方程为调用方的 `CallSiteLoc`，op 类型一律换成循环方程的类型，名称则拼接为调用方名称加原名称。libtpu 据此记录的 primitive 全部变成 `scan`，原来的 `max`、`min` 只剩 scope 的最后一段；原名称中已含 tpuasm 父栈，于是 scope 还会重复一次外层前缀。补丁在上下文内把完全展开的循环改回逐次调用 `jaxpr_subcomp`，每条展开出的 op 直接使用自己方程的 location。两条路径生成的是同一组 op；复现脚本中的 `unrolled` 案例在补丁前后编译出的程序映像逐字节相同，数值也一致。第三处改动针对 Pallas 的逐方程 lowering 缓存（已报告为 [jax-ml/jax#41153](https://github.com/jax-ml/jax/issues/41153)）。`jaxpr_subcomp` 把第一次遇到的方程 lower 成一个游离的 `func.func` 并缓存，之后键相同的方程都用 `inlined_func_call` 克隆它，克隆出的 op 位置被改写为 `callsite(第一个方程 at 当前方程)`。于是相同运算第二次出现时带上了第一次的行号，scope 也会重复。影响不限于来源映射：编译错误和 traceback 同样会指向第一次出现的位置。这个缓存与 jit 无关，`pltpu.roll`、DMA 和取模运算都会命中。补丁在上下文内关闭该缓存，让每个方程用自己的位置调用 lowering 规则。缓存只影响编译时间，关闭前后编译出的机器码相同。兼容性以函数源码为准：先读取已安装 JAX 中 `jaxpr_subcomp` 和 `_lower_jaxpr_to_for_loop` 的源码字节，分别计算 SHA-256 并与 `tc_source_lowering.py` 中的 `_SOURCE_SHA256` 比较，不同就拒绝；然后做必须恰好匹配一次的文本替换（前者两处，后者一处）；最后以原文件名和原行号编译，使 traceback 仍指向真实位置。摘要覆盖 `inspect.getsourcelines()` 确定的函数行范围内的原始文件字节，包括空白、注释和换行，不做格式归一化。之所以不检查 JAX 版本号，是因为 nightly 的版本号不能标识函数内容，真正需要保证的是函数源码没有变化。当前摘要来自 JAX commit `886d2370c1c959d210f522e352e3ddc6bcff7d6c` 的 [jaxpr_subcomp](https://github.com/jax-ml/jax/blob/886d2370c1c959d210f522e352e3ddc6bcff7d6c/jax/_src/pallas/mosaic/lowering.py#L1746) 和 [_lower_jaxpr_to_for_loop](https://github.com/jax-ml/jax/blob/886d2370c1c959d210f522e352e3ddc6bcff7d6c/jax/_src/pallas/mosaic/lowering.py#L4413)。仓库和安装包只保存摘要与固定版本链接；运行时直接读取已安装的 JAX，不下载或分发完整函数副本。JAX 修改其中任一函数后，需要先核对上游 diff、重新验证适配逻辑，再更新摘要和链接。只核对这两个函数并不能覆盖它们调用的其他 JAX 代码：`inlined_func_call` 的问题正是在 `jaxpr_subcomp` 源码不变时出现的，因此升级 JAX 后仍须运行复现脚本。 ### JAX jit 追踪缓存 `jnp.maximum` 和 `+` 对应的 `jnp.add` 等函数本身是 `jax.jit` 函数。JAX 用 `pjit._infer_params_cached` 按函数和参数的 aval（含 sharding）缓存追踪出的 jaxpr；kernel 中以相同的 aval 再次调用时，直接复用第一次调用的 jaxpr，其中方程的 traceback 仍指向第一次调用。于是后一次调用生成的指令，来源同时带有自己的行号和第一次调用的行号。例如 [examples/pallas/reduction.py](../../examples/pallas/reduction.py) 第 25 至 27 行连续三次调用 `jnp.maximum(maximum, pltpu.roll(...))`。不设置 abstract mesh 编译时，第 26 行的 `max` 显示为 `reduction.py:25:14-25:72 <- 26:14-26:72`，libtpu 的 `loc(...)` 也写成第 25 行。同一行 `roll` 的指令也带上第 25 行，原因则是上文的 Pallas lowering 缓存，补丁关闭该缓存后只剩本行。缓存是否命中还取决于有没有 abstract mesh。从 Ref 读出的值，其 aval 的 sharding 使用空 mesh；在 `jax.sharding.use_abstract_mesh(...)` 中，jit 函数输出的 sharding 带有当前的 abstract mesh。因此第 25 行的输入来自 Ref，第 26 行的输入来自上一个 jit 输出，两者的缓存键不同，第 26 行被重新追踪；第 27 行的输入与第 26 行相同，仍然命中缓存。追踪整个 kernel 时，缓存未命中从 5 次变为 7 次，多出的是第 26 行的 `jnp.maximum` 和第 31 行的 `jnp.add`。这时第 26 行的来源只剩本行，机器码不变。离线编译必须设置 abstract mesh，Pallas lowering 才能读出 TPU 代际，所以 [examples/pallas/common.py](../../examples/pallas/common.py) 在设备上编译时也设置它，使两种方式导出的清单逐字节相同。这是 JAX 追踪层的行为，tpuasm 如实记录 JAX 给出的位置，不做修正。 ### 编译器 flag 需要打开 `xla_jf_emit_annotations`、`xla_mosaic_enable_llo_source_annotations` 和 `xla_xprof_register_llo_debug_info` 三个 flag。其中最后一个决定是否创建 LLO SourceMap；只打开注释相关的 flag，并不会产生来源信息。libtpu 可能已经初始化，也可能还没有，所以这三个 flag 既通过原生 flag setter 设置，也追加到 `LIBTPU_INIT_ARGS`。退出上下文时两处都恢复原值。 ### 原生 hook [tc_source_native.cc](../../src/tpuasm/tc_source_native.cc) 替换了 libtpu 中若干处调用。每个 hook 都围绕原函数工作，只增加、保留或合并来源元数据（SourceMap ordinal、MLIR 位置、注释文本），不改变指令、调度或机器字节。选择 hook 点的原则是：来源必须沿编译器实际执行的替换、合并、展开或发射关系传递，不从操作数或相邻指令推断。 - **发射 hook**：替换 bundle 发射过程中对单条指令发射函数的调用。它沿指令→region→module 找到 SourceMap，直接遍历原生的 `locations`，只序列化 ordinals 含有本指令的 `SourceInfoProto`，再从原生的 `strings` 读出字符串表和 HLO 名，把记录追加到指令注释中。不对每条指令序列化整个 SourceMap：那样每次发射的开销与程序大小成正比。上述 kernel 不捕获来源时编译需要 3 秒；捕获来源时，逐条序列化整个 SourceMap 的做法需要 53 秒，只序列化匹配的 location 需要 9 秒。追加通过原生 setter 完成，setter 会复制字符串，所以字符串的生命周期由 LLO 指令负责，能覆盖延迟发射和 bundle finalization。 - **替换 hook**（两处指令替换、一处 region 替换）：LLO 优化把一个值替换为新值时，把新值的 ordinal 加入所有含有被替换子图 ordinal 的 SourceInfo。子图的范围是旧值的操作数图中，不经过新旧表达式共同边界就能到达的节点。如果新值本来就在旧值的图中，说明优化是用已有的值消去了一个运算，这时不传播，否则该值的其他用途也会错误地继承这个来源。 - **合并 hook**：BF16 load/store 合并会把两个候选合成一个新值。hook 先调用原来的注释 setter，再把两个候选的来源传播给新值。第二个候选的指针位于调用方的寄存器中，由 trampoline 前缀代码放进第三个参数。 - **store 注释 hook**：store 的外层 annotator 已经分配了 store 槽，内层 `StoreCommon` 的 annotator 于是看不到新增的槽；但它析构时仍会清空 emitter 的当前注释，导致外层 annotator 也失去来源。hook 在这个析构调用前后保存并恢复 emitter 的注释视图，让外层 annotator 为实际新增的槽记录来源。四种 store（普通 / indexed × 有 / 无 offset）的调用点分别登记。 - **load 合并 hook**：load/store 优化器的 `SimplifyVectorCombineWithSublanesPerStrideInternal` 会把两个各取一半 sublane 的 load 合成一个新 load（注释为 `combined load` 或 `combine strided load`），原来的两个 load 被删除。不处理时，新 load 没有来源，Pallas 的 `get` 在清单中就没有记录。hook 同样先调用注释 setter，再把两个原 load 的来源传播给新值；两个原 load 的指针保存在调用方栈帧中，由 trampoline 前缀代码读入第三、第四个参数。 - **DMA 展开 hook**：编译器把一个 `dma.done` 展开为 wait 和 sync decrement 两条新指令。hook 在新指令生成后、原指令销毁前，把原指令的来源和注释关联到新指令。依据是这次展开本身，不按最终的 DMA flag 配对。 - **MXU prep 改写 hook**：为适配 MXU 尺寸而合并序列时，编译器为 latch 和 matmul 重新构造 prep 指令，原指令的来源不会随之传递。hook 在构造返回后，把正被改写的原指令的来源和注释关联到新 prep；来源取自原指令，不取自 prep 的操作数或相邻的 push。 - **v6e prep 注释 hook**：v6e 发射 matprep 时没有像其他 MXU 指令那样安装 annotator，当前 LLO 的注释因此落不到槽上。hook 只在这一次发射前后补上原生 annotator，为实际新增的槽记录注释；原 LLO 本身没有来源时仍然没有。 - **位置注释 hook**：MLIR `llo` dialect 翻译为 libtpu 内部的 LLO 时，编译器用 op 的 `loc(...)` 整体覆盖新指令的注释，region builder 先前写下的说明（例如取模的实现方式、循环退出判断、vreg 切片范围）随之丢失。hook 把这次覆盖改为拼接，格式与编译器自己追加注释时相同：`loc(...) :: 原注释`。它只改注释文本。 - **MLIR CSE hook**：MLIR CSE 删除等价 op 时丢弃它的位置，同一个表达式写在两处（例如两个辅助函数中的 `index * 8`）时，最终指令只剩第一处来源。两个 op 计算的是同一个值，hook 在删除后把两处位置合并为保留 op 的 `FusedLoc`。它只改位置，不改 IR 结构。 - **合并位置 hook**：LLO SourceMap 从每个 op 的位置只解析出一个 SourceInfo，`FusedLoc` 只有第一部分生效。hook 对每一部分分别解析，使每处位置都带上该 op 的 ordinal。 hook 维护五个计数器：emitted、annotated、failures、rewrites、propagated。退出上下文时，failures 不为零就抛出 `RuntimeError`；annotated 为零则发出警告，这通常意味着编译命中了缓存。 ### 版本相关的部分每个 libtpu 版本的来源后端由两部分组成：`source_backends/` 中的版本文件给出 hook 调用的各函数的 VA，包括发射、注释读写、protobuf 消息序列化（`MessageLite::SerializeToArray`）、HLO module 获取、操作数访问、ordinal 追加、两类替换、合并候选的访问、`ScopedAnnotator` 析构和 flag 读写；`tc_source_backend.py` 中的 `SourceBackend` 记录每个调用点的原始字节、hook 名和可选前缀（`calls`），以及安装前核对的字节（`signatures`，分类见下文）。各版本登记的 hook 不必相同：来源后端只登记在该版本上核对过的调用点，没有登记的路径照常编译，只是相应指令缺少来源。发射、替换、BF16 合并和 store 注释 hook 在所有已支持版本中登记；其余 hook 目前只登记了 0.0.49，它们的实现和用到的对象偏移也写在该版本文件中。所有版本共用的 hook 读取的对象偏移不在版本文件中，而是直接写在公共的 `tc_source_native.cc` 中，因为已支持的各版本布局相同。它们涉及：指令的 region、ordinal 和操作数个数；region 的 module；module 的 SourceMap、root region 和 HLO 名；HLO module 的名称和 ID；SourceMap 与 SourceInfo 中的 repeated 字段；emitter 的当前注释；libc++ 字符串的表示。这些偏移由 `signatures` 中的第三类字节守护。若某个版本的布局不同，应把这些偏移改为按版本定义的常量，而不是修改公共文件，以免破坏已有版本。记录的内容也随版本变化，因为 SourceInfo 由 libtpu 从 MLIR location 转换而来。使用同一 JAX 和同一补丁离线编译 v4 示例时，0.0.48 与 0.0.48 nightly 的 22 份清单相同，但与 0.0.49 的清单有 21 份不同。其中 20 份只有来源注释不同：0.0.48 系列的大部分记录在 scope 栈末尾多出一段方程自身的 primitive 名称，例如 0.0.49 的 `[get; …]`、`[top_k/argmax/top_k; …]` 在 0.0.48 中为 `[get/get; …]`、`[top_k/argmax/argmax/top_k; …]`。另外，`rms_norm` 的机器码也不同；两个版本各自打补丁与不打补丁编译出的机器码相同，所以这是编译器调度的差异，与来源捕获无关。仓库中的清单用 0.0.49 生成。 ### 补丁的安装与恢复 [tc_compiler.py](../../src/tpuasm/tc_compiler.py) 中的 `CompilerSourceMapping` 负责安装和恢复补丁： 1. 先选出编解码后端，再按其标识在 `SOURCE_BACKENDS` 中查找来源后端。来源后端单独登记，有编解码后端不代表支持来源捕获。 2. 核对 `signatures` 中的每一段字节，磁盘上的 ELF 文件（经 `PT_LOAD` 把 VA 换算为文件偏移）和内存中已加载的代码都要核对，且所在映射必须是私有的 r-x 映射。`signatures` 包含三类字节：hook 调用的每个函数的入口、各个调用点本身，以及 libtpu 中访问同一批对象偏移的代码片段。第三类是对象布局的防护：hook 按固定偏移读取 emitter、module、SourceMap 等对象，libtpu 若改变了这些布局，相应代码片段的字节也会改变，安装就会被拒绝。 3. 从 `/proc/self/maps` 取得 libtpu 的加载基址。 4. 在距所有调用点 ±2 GiB 以内（`call rel32` 的可达范围）用 `MAP_FIXED_NOREPLACE` 申请一页，作为 gateway。每个调用点在 gateway 中对应一段 32 字节的 trampoline，内容是可选的前缀代码加上 `jmp [rip+0]` 和 hook 的绝对地址。写入后读回核对，再把该页设为 RX。 5. 安装时把每个调用改为 `call rel32` 指向对应的 trampoline（原来是 `jmp rel32` 尾调用时保持 `jmp`，hook 直接返回原调用方）；原调用较长时用 NOP 填满原跨度。修改期间代码页暂时可写，改完恢复 RX；任何一处失败都回滚已改的调用点。 6. 退出时先确认补丁字节没有被改动，再恢复原调用、flag、环境变量和 lowering 函数，确认所有原字节都已复原，最后释放 gateway。如果恢复失败，gateway 可能仍被执行中的代码引用，因此保留状态对象、不释放 gateway，并拒绝在本进程中再次安装，要求重启。嵌套的上下文复用同一个状态。锁只能协调通过本接口进行的调用，不能保护绕过接口、同时进行编译的其他线程。命中编译缓存时不会产生来源记录，使用时需要清除 JAX 缓存并禁用持久编译缓存。 ## 恢复 [tc_source_mapping.py](../../src/tpuasm/tc_source_mapping.py) 的恢复流程只依赖 executable 中保存的数据和实际的解码结果。 ### 程序身份 `ProgramSourceMap` 分别保存以下身份字段：record、映像索引、segment set / segment 索引、初始化数据中的偏移、segment set 的 hash、core program 的 fingerprint、可选的 compilation ID、metadata 所在 record 和 metadata 的 program ID。这些字段不要求彼此相等。不同程序映像之间不按名称或 ordinal 拼接。 metadata 取紧跟在 core 记录之后、且含有 overlay 信息的那条记录。如果一条 core 记录中有多份代码映像，而 metadata 中没有能区分映像的标识，就无法确定 metadata 属于哪份映像。这时每份映像都得到空映射和一条歧义诊断，不能把同一组来源复制给所有映像。 ### 坐标注释表在两种坐标下给出：field 4 以发射时（链接前）的 PC 为键；field 30 以程序映像 PC 为键，内容是链接器已经重定位的 overlay 前缀和后缀。发射 PC 按所在的 overlay 换算为程序映像 PC： ```text image_start = encoded_word_offset * 编码字节数 * 每块 bundle 数 / 块字节数 body_start = image_start + prefix_size image_pc = body_start + emitted_pc - start_bundle_number ``` `encoded_word_offset` 的单位随目标不同：v4 为 512 字节的块（每块 10 个 bundle，`image_start = encoded_word_offset * 10`），v6e 为 32 字节（每 512 字节 8 个 bundle，`image_start = encoded_word_offset / 2`）。v6e 的单位是从示例的元数据推断的：按 32 字节换算后，示例的注释都落在解码结果中被占用的槽上，没有诊断；按 64 字节或 512 字节换算，第二个 overlay 会落到程序映像之外。`Overlay` 同时保存原始的 `encoded_word_offset` 和换算后的 `image_start`。换算以 map 的键为准：continuation 尾部的 `Annotation.pc` 可能从零重新开始计数，不能代替键。 ### 槽与来源注释中的槽名是 libtpu 的枚举名（如 `SLOT_VECTOR_EXTENDED_0`），先按目标映射为物理槽名，再检查实际解码的 bundle 中该槽是否被占用；不符合时记录诊断，并丢弃这条注释。槽的归属以解码器为准。v6e 的枚举名为 `SLOT_SCALAR_ALU_0/1`、`SLOT_DMA`、`SLOT_VECTOR_ALU_0..3`、`SLOT_VECTOR_STORE`、`SLOT_VECTOR_LOAD_0/1`、`SLOT_VECTOR_MISC`、`SLOT_VECTOR_EXTENDED_0/1` 和 `SLOT_VECTOR_RESULT_0/1`，与物理槽一一对应。注释中的 v1 记录解析为 `InstructionOrigin`，并去重；其余文本保留为 `compiler_annotation`。文本中 `loc(...)` 明确记录的位置单独解析为 `annotation_locations`，但不据此补造 primitive 或 scope。 `SlotSource.source_frames` 是给一般调用者的统一视图：合并捕获记录中的 frame 与 `annotation_locations`，按文件和行列去重。源码浏览器只需读这个字段；原始字段继续保留，供需要 primitive、scope、ordinal 或原始证据的调用者使用。`source_kind` 标明位置来自 tpuasm 的捕获（`captured`）、只来自编译器的 `loc(...)`（`compiler_location`），还是没有位置（`unknown`）。函数区间内被占用但没有注释的槽也会列出，来源为空。恢复端不补造来源：操作数依赖不说明两条指令来自同一源码操作，所以不从寄存器 def-use、DMA flag 或相邻指令推导来源。 ### 函数区间 HLO 为 custom-call 的 symbol 视为函数。它的每段 `child_instructions` 按 overlay 分别裁切、换算，保留区间之间的空洞，不用最小和最大 PC 把区间填满。去重后共享的机器代码可以同时属于多个函数，它们的区间允许重叠；每个槽记录覆盖该 PC 的所有函数。 `FunctionSource.ranges` 原样保留这些编译器区间，listing 的函数标记也用它们。函数边界只由编译器区间决定，不按源码覆盖率裁切：区间内没有来源的指令和空 bundle 同样属于该函数。 ### 状态至少一个槽有 v1 来源时，状态为 `captured`，但这不表示所有指令都有来源。否则状态为 `absent`，附带提示重新编译的诊断，同时仍保留编译器原生注释和函数归属。没有来源的指令保持 `unknown`，不能据此判断它属于 runtime。排查来源缺口时，应沿编译阶段追踪同一条指令，找到位置丢失的那一步，再在那一步的实际改写关系上保存来源，而不是从最终 bundle 反推或借用相邻指令的位置。上文的 DMA 展开和 MXU prep 改写 hook 就是这样找到的。 ### 依赖的元数据字段代码按字段号读取以下消息： | 消息 | 使用的字段 | |---|---| | `TpuCoreProgramProto` | fingerprint=3、compilation_id=4、tensor_core=5、memory_segments=8 | | TensorCore / sequencer | sequencer=1；sequencer_type=3，值 1 表示 TensorCore | | `TpuMemorySegmentSetProto` | segments=2、initialized_data=3、hash=4、compression=5 | | memory segment / range | segment_type=1（CODE=1）、range=3；start_byte_offset=1、size=2 | | `CompilerMetadata` | annotation_metadata=4、symbol_table=8、program_id=9、overlay_metadata=10、module_name=20、overlayer_annotation_metadata=30 | | `AnnotationMetadata` / `Annotation` | annotations=1（map）；pc=1、slot_to_text=2、annotation=3 | | `Overlay` | prefix_size=1、start_bundle_number=2、end_bundle_number=3、suffix_size=4、encoded_word_offset=5 | | `Symbol` / HLO | parents=21、child_instructions=23、hlo_instruction=70；op=3、name=4、deduplicated_name=8 | ## 输出这些 dataclass 定义在同一模块中。JSON 由 `dataclasses.asdict` 生成并加上 `schema_version`，所以修改任何 dataclass 字段都会改变 JSON schema，需要同时提高版本号。当前为版本 3：版本 2 相对版本 1 增加了 `ProgramSourceMap.target` 和 `Overlay.image_start`；版本 3 增加 `source_frames`、`source_kind`，并列出函数范围中没有注释的已占用槽。本模块的 `SourceLocation` 表示带 primitive/scope/ordinal 的 Pallas 捕获位置，`SourceFrame` 表示统一投影中的文件行列 frame；汇编诊断位置是另一个类型 `AssemblyLocation`。 `source_comments` 把映射渲染为 `#` 注释：函数区间、bundle 注释和诊断放在 bundle 之外，逐槽来源放在对应指令行的末尾。每个捕获位置显示为一条从内到外的调用链 `path:内层 <- 调用方 [scope; LLO n]`，调用方可能是外层方程，也可能是上文 jit 缓存留下的第一次调用。同一槽上相同的调用链只显示一次，方括号中合并它的全部 scope 和 LLO ordinal。编译器 `loc(...)` 的坐标已出现在调用链中时不再重复，否则单独列出且不带方括号，因此带方括号的是 tpuasm 捕获的来源，不带的是编译器的位置。其余注释文本原样保留，JSON 保留完整的 `compiler_annotation`。注释不参与编码。 ## 验证 [tests/reproduce_tpu_v4_tc_source_mapping.py](../../tests/reproduce_tpu_v4_tc_source_mapping.py) 需要在 TPU 上运行，它检查以下内容：机器码与数值不受补丁影响；load/get、push、matmul、pop、store/swap 各自的来源关联；normal 与 xpose 两种布局下 DWG 和 push 的操作数；完全静态展开循环中每条指令保留自己的 primitive；安装中途失败时的回滚，以及嵌套上下文与异常退出后，调用点、flag、环境变量和两个 lowering 函数都能恢复原状。 v6e 没有对应的设备测试：[examples/pallas/](../../examples/pallas/) 的示例在 `run_all.sh --aot tpu-v6e-tc` 下离线编译时捕获来源，只能核对机器码，不能核对数值。CI 重新导出全部示例清单，要求与仓库逐字节相同。新增或修改补丁（JAX lowering 或 libtpu hook）时，比较 v4 与 v6e 全部示例修改前后的清单：去掉 `#` 注释后应完全相同，注释中原有的来源和编译器说明不应丢失；再在 v4 上运行复现脚本，确认数值和补丁生命周期。