TPU 缓冲区布局

本页中的所有地址、结构体偏移和字段标签均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d）。该镜像未被 strip；反混淆后的 C++ 符号名按原文引用。其他版本会有所不同。

摘要

TPU 上的 PJRT 设备缓冲区是一段扁平的 HBM 字节，其内部排列不是行主序。用户交给运行时的逻辑 xla::Shape（一个 (element_type, dimensions[], layout) 三元组）和真正落到 HBM 中的字节之间有一次转换，即 xla::jellyfish::TransferSizeUtil::HostShapeToDeviceShape，它会 (a) 选择一个 minor-to-major 维度顺序，(b) 将两个最 minor 的维度填充到芯片的 (sublane, lane) tile，(c) 打包亚字节元素并拆分 >32-bit 元素类型，以及 (d) 在 layout 上盖上描述物理分块的 xla::Tile 列表。填充后的设备形状就是分配器按其确定大小、DMA 描述符寻址、以及 ShapeSizeBytesRaw 度量的对象。本页记录这次转换：shape↔layout 数据模型、tile 填充算术、元素打包规则，以及命名哪些 HBM 地址保存缓冲区叶子的设备端缓冲区驻留记录（xla::ShapedBuffer → XLA_ShapedBuffer）。

熟悉 GPU 上 XLA 的读者应先记住一个类比，然后立刻把它复杂化。在 GPU 上，xla::Layout 是一个 minor_to_major 排列，除此之外几乎没有别的内容；字节是经过排列后的紧密行主序。在 TPU 上，layout 额外携带一个 tile（xla::Layout::tiles()，一个 xla::Tile 向量），物理缓冲区则是将逻辑数组重塑为 (sublane, lane) tile 后按 tile-major 布局。lane 维度在所有代际上都是 128；在这个 v5+/v6e/TPU7x 构建中，sublane 维度是 8（更早的一代使用 16）-- 这正是布局分配盖到每个叶子形状上的同一个 (SublaneCount, LaneCount) tile。硬件将一个 tile 作为一个连续块读取，因此 [3, 5] f32 数组并不占用 60 字节：它会填充到 [8, 128]（一个 tile）并占用一个粒度，尾部 lane/sublane 被填充（在填充路径上为 0xFF）。浪费的填充是真实的 HBM，而最小化它正是布局分配选择器为紧凑字节大小而优化的原因。

本页负责设备端分块缓冲区：shape→device-shape 映射（HostShapeToDeviceShape / SetPaddedShape）、tile 填充规则（Pad2ndMinorCompact、GetCompactTiles）、元素打包因子、字节大小函数（ShapeSizeBytesRaw / ShapeSizeCompact / ShapeWithMetadataSizeBytes）以及驻留记录。它不复述编译期对 minor_to_major 的选择（见 layout-assignment.md）、交出 offset 的 HBM 空闲列表分配器（见 hbm-allocator.md）、DMA 线缆对齐（见 hbm-dma-alignment.md），或 PJRT PJRT_Buffer ABI 与外部引用计数（见 pjrt/buffer-and-memory.md）。

对于重新实现，契约是：

shape↔layout 数据模型 -- 设备 layout 读取的 xla::Shape / xla::Layout 字段，以及带有精确偏移的 XLA_Shape / XLA_Layout C-ABI 镜像。
shape→device-shape 映射 -- HostShapeToDeviceShape → SetPaddedShape：解包亚字节类型、递归、填充 minor 维度、盖上 tile。
tile 填充规则 -- minor dim → LaneCount 的倍数（否则为下一个 pow2）；第 2 minor 维经由 Pad2ndMinorCompact；granule_bytes % sublane_bytes == 0 不变量。
元素打包 -- ElementPackingFactor（亚字节通过 tc_max_packing_factor 打包），以及 64/128-bit 拆分（ComponentShape<64>/<128>）。
驻留记录 -- xla::ShapedBuffer（设备端形状 + 每叶子的 DeviceAddressBase），镜像到 XLA_ShapedBuffer。


Host→device shape	`xla::jellyfish::TransferSizeUtil::HostShapeToDeviceShape`（经由 `HardwareLayout::HostShapeToDeviceShape` @ `0xeab0e20`）
填充形状核心	`TransferSizeUtil::SetPaddedShape` @ `0x1d6ae0e0`
Tile 构造器	`TransferSizeUtil::GetCompactTiles` @ `0x1d6b11c0`（返回 `inlined_vector<xla::Tile,3>`）
第 2 minor 填充	`TransferSizeUtil::Pad2ndMinorCompact` @ `0x1d6af5c0`
元素打包	`TransferSizeUtil::ElementPackingFactor` @ `0x1d6b03e0`（表驱动，`kPackingFactors[35]`）
填充后字节大小	`TransferSizeUtil::ShapeSizeBytesRaw` @ `0x1d6add40`；带元数据为 `ShapeWithMetadataSizeBytes`（经由 `HardwareLayout::ShapeSize` @ `0xeab0ec0`）
紧凑字节大小	`TransferSizeUtil::ShapeSizeCompact` @ `0x1d6ae8a0`（经由 `HardwareLayout::ShapeSizeCompact` @ `0xeab0f20`）
Tile 盖章	`xla::jellyfish::HardwareLayout::PopulateShape`；每叶子 `TransferSizeUtil::UpdateLeafLayout` @ `0x1d6b08c0`
线性化器	`TransferSizeUtil::LinearIndex` @ `0x1d6b0600` → 分块时为 `xla::LayoutUtil::LinearIndex`
HBM 重打包器	`RepackToHardwareLayout<sublane,128>` @ `0x1d5c3b40`（`<16,128>`）、`0x1d5c2880`（`<2,128>`）、`0x1d5c3f80`（`<32,128>`）
驻留记录	`xla::ShapedBuffer` → `XLA_ShapedBuffer`（`ApiConverter::ToC` @ `0xfcb8580`）
默认 tile	`(SublaneCount, LaneCount)` = 本构建中为 `(8, 128)`；更早一代为 `(16, 128)`
置信度	CONFIRMED（字节锚定），除非某行或标注另有说明

1. Shape / Layout 数据模型

目的

本页上的所有内容都操作两个 C++ 对象：xla::Shape 以及其中包含的 xla::Layout。重新实现者必须精确复现这些字段，因为设备 layout 代码会通过原始偏移读取它们，运行时也会跨 C ABI 将它们镜像为 XLA_Shape / XLA_Layout，用于 PJRT/stream_executor 边界。编译期负责填入 minor_to_major 和 memory_space 的 pass 是布局分配；本页是消费者。

设备路径读取的 `xla::Layout` 字段

xla::Layout 是一个富含 absl::InlinedVector 的结构体；设备代码只接触少数字段。C-ABI 镜像 XLA_Layout（由 ApiConverter::ToC(const xla::Layout&, XLA_Layout*) @ 0xfcb7ca0 构建）固定了偏移，设备端函数则通过 xla::Shape::layout()（返回 Layout*）到达同一份数据。

字段	`xla::Layout` 访问	`XLA_Layout` 偏移	含义
`minor_to_major[]`	`layout()[2..3]` inlined vector（计数在 `[2]>>1`）	`+0`（ptr），`+48`（count）	物理维度顺序，minor 在前
`dim_level_types` / flags	`layout()` 低字节	`+400`，`+404`	sparse/dense 层级类型
`index_primitive_type` / `pointer_primitive_type`	`((qword)layout+1)`	`+408`	sparse index/pointer 类型
`element_size_in_bits`	`layout()[1]`	`+416`	对 packed/奇异宽度元素非零；`0` ⇒ 使用 dtype 字节大小
`memory_space`	`((char)layout+2)`	`+424`	层级颜色（HBM/VMEM/...）；见 overview.md
`tiles[]`	`layout()[9]` 是 `tiles().size()`	`+56`（ptr），`+392`（count）	物理 tile 列表（每个 `xla::Tile` 是一个维度向量）
`tail_padding_alignment_in_elements`	`((qword)layout+24)`	`+432`	尾部元素填充量子

最常被查询的字段是 layout()[9] -- tile 计数。LinearIndex（0x1d6b0600）、ShapeSizeBytesRaw（0x1d6add40）和 HasLinearLayout（0x1d6b0160）都会根据 tiles().size() >= 2（或 > 1）分支：已填充的 tile 列表表示缓冲区是分块的，并通过 ShapeUtil::ArraySize / LayoutUtil::LinearIndex 寻址；空 tile 列表表示线性（未分块）缓冲区。

`xla::Shape` 字段

xla::Shape（以及它的 XLA_Shape 镜像，ApiConverter::ToC @ 0xfcb7940）携带维度和 layout。XLA_Shape 偏移直接从转换器解码：

text

XLA_Shape（tuple array 中每个元素 536 字节）:
  +0    element_type      (PrimitiveType, int32)
  +8    dimensions[]      (int64 array)          count @ +56
  +64   dynamic_dims[]    (bool array)           count @ +72
  +80   tuple_shapes[]    (XLA_Shape array, 536 B stride)  count @ +88
  +92   has_layout        (bool)
  +96   layout            (XLA_Layout, present iff has_layout)
```text

> **NOTE --** 设备端函数一开始会特判三个 `element_type` 家族，它们来自用 `_bittest64` 测试的有效性 bitmask `0x2FFF91FFE`（可搜索数组类型，≤ 0x21）、mask-`0x20`-family token-and-opaque 类型，以及用于 64/128-bit 家族的 bitmask `0x400048000`。同一个 `0x2FFF91FFE` mask 也守卫[布局分配选择器](../compiler/layout-assignment.md#per-tensor-chooser--findmemoryminimizinglayout)；它是规范的“这是否为分块数组元素类型”谓词。

### 函数图

| 函数 | 地址 | 角色 |
|---|---|---|
| `ApiConverter::ToC(const xla::Layout&, XLA_Layout*)` | `0xfcb7ca0` | `xla::Layout` → C ABI；固定 layout 偏移 |
| `ApiConverter::ToC(const xla::Shape&, XLA_Shape*)` | `0xfcb7940` | `xla::Shape` → C ABI；递归 tuple shapes |
| `ApiConverter::FromC(XLA_Shape*)` | `0xfcb7400` | C ABI → `xla::Shape`，用于运行时调用 |
| `xla::jellyfish::TransferSizeUtil::HasLinearLayout(const Layout&)` | `0x1d6b0160` | tile 列表为空 ⇒ 线性 |

---

## 2. Host Shape → Device Shape

### 目的

`HostShapeToDeviceShape` 是将逻辑 host `xla::Shape` 转换为 TPU 存储的物理*设备形状*的入口点。它是 transfer manager、分配器定大小路径以及 C shim `Tpu_OnDeviceShape` 调用的内容。工作委托给 `TransferSizeUtil::HostShapeToDeviceShape`，后者分派到每个叶子的 `SetPaddedShape` 并重新盖上 layout。填充后的设备形状在三个方面不同于 host shape：填充后的 minor 维度、打包/拆分后的元素类型，以及已填充的 `tiles()` 列表。

### 入口点

```text
HardwareLayout::HostShapeToDeviceShape (0xeab0e20)        ── C-ABI wrapper
  ├─ ApiConverter::FromC(host XLA_Shape -> xla::Shape)
  ├─ GetRegisteredDeepseaPlatform / DeepseaPlatform::GetTopology   ── the tpu::TpuTopology*
  ├─ GlobalTpuCompEnv()                                    ── reads comp-env+4124, +4373 (pack flags)
  ├─ TransferSizeUtil::HostShapeToDeviceShape(topology, shape, ...)
  │     └─ per leaf: SetPaddedShape (0x1d6ae0e0) -> HardwareLayout::PopulateShape
  └─ ApiConverter::ToC(device xla::Shape -> XLA_Shape)

算法 -- SetPaddedShape

SetPaddedShape（0x1d6ae0e0）是映射的核心。它对打包元素类型递归，并按 minor 维度计算填充后的范围，然后将维度交给 HardwareLayout::PopulateShape 来盖上 tile。

function SetPaddedShape(topology, shape /*buffer leaf*/, out_shape):   // 0x1d6ae0e0
    REQUIRE(LayoutUtil::HasLayout(shape))            // FATAL "Can't pad a shape without knowing its layout", line 1019

    // ---- Step 1: sub-byte / packed element types -> unpack, recurse, re-stamp ----
    pack = ShouldPackPREDAsSingleBit(topology, shape)         // PRED -> 1-bit packing
    if ElementPackingFactor(shape.element_type(), pack) >= 2:
        unpacked = GetUnpackedShape(topology, shape)          // widen the packed dtype
        SetPaddedShape(topology /*unpacked*/, shape, out_shape)   // recurse on unpacked
        out_shape.set_element_type(shape.element_type())      // restore original packed dtype
        ForEachMutableSubshape(out_shape, UpdateLayout)       // re-stamp tile per leaf
        return OK

    pbytes = ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type())

    // ---- Step 2: rank<=1 scalar/vector fast path ----
    if shape.dimensions().size() <= 1:
        CHECK(scalar_layout.minor_to_major().size() == 0)     // line 1040
        chunk_bytes = 4 * topology.chunk_size_elems            // this[53]
        CHECK(chunk_bytes % pbytes == 0)                       // line 1043
        out_shape = PopulateShape(et, {chunk_bytes/pbytes}, tile=1, layout)
        return OK

    // ---- Step 3: only 4-byte (or PRED) element widths are implemented for tiling ----
    if shape.element_type() != PRED and pbytes != 4:
        return Unimplemented("Attempted to map shape ... to on-device TPU padded shape   // line 1050
                              but this is not implemented")

    // ---- Step 4: pad each minor dim up to its chunk bound ----
    padded_dims = copy(shape.dimensions())
    for d in [0 .. rank):                                       // walk in minor_to_major order
        dim_idx  = minor_to_major[rank-1 - d]
        bound    = ChunkBound(topology, d, rank, dim_idx)        // 0x1d6b22e0
        extent   = shape.dimensions(dim_idx)
        padded_dims[dim_idx] = round_up(extent, bound)           // ceil(extent/bound) * bound

    // ---- Step 5: stamp the tile + memory_space ----
    out_shape = HardwareLayout::PopulateShape(et, padded_dims, ntiles, layout)
    return OK
```text

> **GOTCHA --** Step 3 是一堵硬墙：`SetPaddedShape` 只为 **4-byte 元素宽度**（以及 PRED）实现分块。任何更宽的类型都会在到达 4-byte `SetPaddedShape` 之前走 `ShapeSizeBytesRaw`（§4）中的 64/128-bit *拆分*路径；任何更窄的类型都会在 Step 1 中被*打包*到等价 4-byte。试图直接对原始 bf16 或 s8 缓冲区分块的重新实现会命中 line 1050 的 `Unimplemented` -- 设计强制每种元素宽度都先通过打包/拆分进入 4-byte 分块内核。
>
> **QUIRK --** rank≤1 快路径（Step 2）**不会**填充到 2D `(sublane, lane)` tile。标量或 1-D 缓冲区会被重塑为单个包含 `chunk_bytes/pbytes` 个元素的 chunk（`chunk_bytes = 4 * topology.chunk_size_elems`）。只有 rank≥2 的缓冲区才获得 2-minor tile。这就是为什么包含 N 个 f32 元素的一维向量会向上取整到 chunk 倍数，而不是到 `[1, 128]` lane tile。

### 函数图

| 函数 | 地址 | 角色 |
|---|---|---|
| `HardwareLayout::HostShapeToDeviceShape` | `0xeab0e20` | C-ABI host→device shape |
| `TransferSizeUtil::SetPaddedShape` | `0x1d6ae0e0` | 每叶子填充 + tile 盖章 |
| `TransferSizeUtil::GetUnpackedShape` | `0x1d6b2100` | 为递归拓宽 packed dtype |
| `TransferSizeUtil::ShouldPackPREDAsSingleBit` | `0x1d6b0080` | PRED 1-bit 打包谓词 |
| `TransferSizeUtil::ChunkBound` | `0x1d6b22e0` | 填充向上取整使用的每维 chunk bound |
| `HardwareLayout::PopulateShape` | `0x1d6da360` | 盖上 dims + tile + `memory_space` |
| `XlaShapeToTpuPaddedShape` (C shim) | `0xeabf0e0` | `tensorflow::XlaTpuPaddedShapeFn` 边界 |
| `XlaShapeToTpuShapeRepresentation` (C shim) | `0xeabefa0` | Shape-representation 边界 |

---

## 3. Tile 与填充规则

### 目的

设备缓冲区的物理字节是将逻辑数组重塑为 `(sublane, lane)` tile 后按 tile-major 布局。本节给出 tile 几何、精确的第 2 minor 填充算术，以及重新实现者必须强制满足的 `granule % sublane` 不变量。tile 值在启动时从芯片的 `tpu::TpuTopology` 读取，绝不硬编码。

### tile 几何

默认叶子 tile 是 `(SublaneCount, LaneCount)`，与[布局分配盖到每个叶子上](../compiler/layout-assignment.md#depth-aware-layout-cost-and-the-tile)的 `xla::Tile` 相同：

- **`LaneCount` = 128**，跨所有代际相同（topology 字段 `*((qword*)topology+52)`，用作 minor-dim bound）。
- **`SublaneCount`** = topology 字段 `*((qword*)topology+51)` -- 在此 v5+/v6e/TPU7x 构建上为 **8**，在更早一代为 **16**（第 2 minor bound）。
- **`granule_bytes`** = `*(qword*)(topology->[1] + 200)` -- 硬件 DMA 粒度。`GetCompactTiles` 和 `Pad2ndMinorCompact` 都会在 line 77 `CHECK(granule_bytes % sublane_bytes == 0)`，其中 `sublane_bytes = 4 * topology->[51]`（4 字节 × sublane count）。

因此，叶子的物理连续段是：tile 覆盖两个 minor 维度，每个 tile 含 `SublaneCount × LaneCount` 个 4-byte 槽，并被填充满。对于亚字节 dtype，还会附加一个额外的 *subtile*（§3.3）。

### 第 2 minor 填充 -- Pad2ndMinorCompact

`Pad2ndMinorCompact`（`0x1d6af5c0`）计算填充后的第 2 minor 范围和每 tile 元素数。这是重新实现者必须复制的精确规则。

```c
function Pad2ndMinorCompact(topology, extent /*2nd-minor logical*/, element_type):  // 0x1d6af5c0
    lane = topology[52]                                  // LaneCount, 128
    if extent >= lane:
        padded = round_up(extent, lane)                  // ceil(extent/lane) * lane
    else:
        padded = next_pow2(extent)                       // 1 << ceil_log2(extent)

    granule_bytes  = *(topology->[1] + 200)
    sublane_bytes  = 4 * topology[51]                    // 4 * SublaneCount
    CHECK(granule_bytes % sublane_bytes == 0)            // line 77
    per_tile_rows  = (granule_bytes / sublane_bytes) * ElementPackingFactor(element_type, /*pack=*/0)

    return max(padded, per_tile_rows)

QUIRK -- 小于 LaneCount 的分支会向上取整到下一个 2 的幂，而不是 lane count。第 2 minor 范围为 5 会填充到 8，而不是 128；范围为 100 则填充到 128（lane 倍数）。然后填充范围会被 per_tile_rows 设下限，因此一个 tile 永远不会小于一个粒度所需的 sublane。总是将第 2 minor 向上取整到 LaneCount 的重新实现会过度填充每个小缓冲区，并与 libtpu 的 HBM footprint 不一致。

元素打包 -- ElementPackingFactor

ElementPackingFactor（0x1d6b03e0）返回一个 4-byte 物理槽中共享的逻辑元素数量。它完全由表驱动：

function ElementPackingFactor(topology, element_type, pack_pred_as_bit):   // 0x1d6b03e0
    tc_max = topology->[432]                              // tc_max_packing_factor, a power of two
    REQUIRE(popcount(tc_max) == 1)                        // else FATAL "Unsupported tc_max_packing_factor", line 1489
    k = tzcnt(tc_max)                                     // 0..5  (1,2,4,8,16,32)
    REQUIRE(element_type < 0x23)                          // 35 primitive types
    if k == 0: return 1                                   // no packing
    return kPackingFactors[2^k][pack_pred_as_bit][element_type]
```text

`kPackingFactors[35]` 表按 `tc_max`（`<2>`、`<4>`、`<8>`、`<16>`、`<32>`）以及 pack-mode（`pack_pred_as_bit` 标志为 `tc_max >= 8` 选择 `true` 表，将 PRED 打包为单 bit）区分。给定 dtype 的 factor 是该 `tc_max` 下适合放入 4 字节的逻辑元素数量。这与[布局分配 `ShapeSizeCompact`](../compiler/layout-assignment.md) 折叠进其字节计数的打包相同。

> **NOTE --** 亚字节打包会向 tile 列表添加 *subtile*。在 `GetCompactTiles` 中，主 `(sublane, lane)` tile 被 push 之后，如果 `ElementPackingFactor(...) >= 2`，它会附加一个 `GetSubtileForPacking` tile（并且当 minor dim 的范围为 1 时，附加一个 `GetSubtileForBreakingMinorDimensionForPacking` tile）。因此 packed 缓冲区的 `tiles()` 列表长度为 2-3，重新实现者必须发出 subtile，否则字节不会匹配。

### 函数图

| 函数 | 地址 | 角色 |
|---|---|---|
| `TransferSizeUtil::GetCompactTiles` | `0x1d6b11c0` | 为叶子构建 `inlined_vector<xla::Tile,3>` |
| `TransferSizeUtil::Pad2ndMinorCompact` | `0x1d6af5c0` | 第 2 minor 填充 + 每 tile 行数 |
| `TransferSizeUtil::ElementPackingFactor` | `0x1d6b03e0` | 每个 4-byte 槽中的逻辑元素数（表驱动） |
| `TransferSizeUtil::GetSubtileForPacking` | `0x1d6b1f20` | 亚字节 subtile |
| `TransferSizeUtil::GetSubtileForBreakingMinorDimensionForPacking` | `0x1d6b1fe0`（在 `0x1d6b11c0` 中调用） | minor 范围 == 1 时的 subtile |
| `TransferSizeUtil::DoesShapeRequireMultiChunkPacking` | `0x1d6b0720` | Multi-chunk 打包谓词 |
| `Target::LaneCount` / `SublaneCount` | `0x1d60f400` / `0x1d60f300` | 来自芯片描述符的 tile dims |

---

## 4. 字节大小 -- ShapeSizeBytesRaw 及相关函数

### 目的

分配器必须知道设备缓冲区占用多少 HBM 字节。存在三个大小函数，它们都由*填充后的设备形状*驱动，重新实现者必须选对：`ShapeSizeBytesRaw`（填充后的数据字节）、`ShapeSizeCompact`（同上，但使用紧凑分块）以及 `ShapeWithMetadataSizeBytes`（数据 + 每缓冲区元数据，由 transfer manager 使用）。它们分别作为 `HardwareLayout::ShapeSize`（`0xeab0ec0`）、`HardwareLayout::ShapeSizeCompact`（`0xeab0f20`）和 `HardwareLayout::ShapeSizeCompactRaw`（`0xeab0f80`）暴露给 C ABI。

### 算法 -- ShapeSizeBytesRaw

`ShapeSizeBytesRaw`（`0x1d6add40`）是元素类型家族上的分派器。分支顺序就是重新实现契约。

```c
function ShapeSizeBytesRaw(topology, shape):                      // 0x1d6add40
    et = shape.element_type()

    if et == TUPLE (13):                                          // tuple = array of pointers
        n = ShapeUtil::TupleElementCount(shape)
        return round_up(4 * n, ptr_granule)                       // 4 B / element, topology->[1]+200

    if et == TOKEN (17) or IsZeroElementArray(shape) or MaxElementsInPerSplit == 0:
        return 0

    if ElementHasBitWidth(shape, 64):                             // s64/u64/f64/c64-half
        return 2 * ShapeSizeBytesRaw(topology, ComponentShape<64>(shape))   // hi/lo split
    if ElementHasBitWidth(shape, 128):                            // c128
        return 2 * ShapeSizeBytesRaw(topology, ComponentShape<128>(shape))

    if ColorToMemorySpace(layout.memory_space) == kSparseCoreSequencerSflag (12):
        return ExtentProduct(shape) * ByteSizeOfPrimitiveType(et)  // SparseCore: untiled, dense

    if shape.layout().tiles().size() >= 2:                        // already tiled
        return ShapeUtil::ArraySize(shape)

    // untiled 4-byte / packed leaf: pad it, then measure
    tmp = thread_local Shape
    CHECK(SetPaddedShape(topology, shape, &tmp) is OK)            // line 264
    max_elems = LayoutUtil::MaxElementsInPerSplit(tmp, shape)
    if layout.element_size_in_bits == 0:
        return max_elems * ByteSizeOfPrimitiveType(et)
    else:
        bits = element_size_in_bits * max_elems                   // packed: bit-accurate
        return bits/8 + (rounding correction)                     // ceil to byte

Note -- ShapeSizeBytesRaw（0x1d6add40）在 untiled-dense 分支上测试的字面量 12 是规范 LLO MemorySpace enum 中的 sparse_core_sequencer_sflag -- 不是 sparse_core_sequencer_smem，后者是 14（由 MakeSparseCoreSequencerSmemConstant @ 0x1d60bc60 = mov $0xe,%esi 字节确认）。ColorToMemorySpace（0x1d6ffb00）是一个受 color < 0xA 守卫的 byte_B5435CA[color] remap，因此它的输出已经是规范 MemorySpace enum 值，而不是原始 layout color。17 值 owner 表见 memory-space-table.md。
GOTCHA -- 64- 和 128-bit 元素拆分不是“每元素 8 字节”。64-bit 缓冲区由 ComponentShape<64> 分解为两个 32-bit component 缓冲区（高字和低字），每个都作为 4-byte 缓冲区独立分块，并将大小相加（2 * ...）。因此 [N] f64 缓冲区在物理上是两个 [N] f32 分块缓冲区，而不是一个 [N] 8-byte 分块缓冲区。将 f64 视作 8-byte 元素并直接分块的重新实现会错误计算每个 double-precision 缓冲区的大小和地址。同样的拆分也支配 DMA 如何读取该缓冲区。
NOTE -- element_size_in_bits != 0 分支是 packed-buffer 大小：它乘以 bit width 并除以 8，带有向上取整修正。这就是 bit-packed PRED 缓冲区（1 bit/element）报告其字节大小的方式 -- 0 值表示“使用 dtype 的自然字节大小”，非零值表示“每元素这么多 bit，并将总数向上取整到字节”。

函数图

函数	地址	角色
`TransferSizeUtil::ShapeSizeBytesRaw`	`0x1d6add40`	填充后的数据字节（家族分派）
`TransferSizeUtil::ShapeSizeCompact`	`0x1d6ae8a0`	紧凑分块字节大小
`TransferSizeUtil::ShapeSizeCompactRaw`	`0x1d6aea60`	不含元数据的紧凑字节
`TransferSizeUtil::ShapeWithMetadataSizeBytes`	`0x1d6aea00`（经由 `0xeab0ec0`）	数据 + 每缓冲区元数据
`HardwareLayout::ShapeSize` / `ShapeSizeCompact`	`0xeab0ec0` / `0xeab0f20`	C-ABI 字节大小 wrapper
`HardwareLayout::ComponentShape<64>` / `<128>`	`0x1d6d9cc0` / `0x1d6d9e40`	64/128-bit 元素拆分
`xla::ShapeUtil::ArraySize`	(OSS)	分块数组字节大小

5. 线性化与 HBM 重打包器

目的

两个操作会将分块设备形状转换为实际字节移动：线性化（LinearIndex，分块 multi-index → 扁平 offset）和重打包内核（RepackToHardwareLayout，将 host 字节物理重排为分块硬件顺序的 SIMD 例程）。重新实现者需要两者：索引数学用于寻址单个元素，重打包用于为 DMA 暂存整个缓冲区。

算法 -- LinearIndex

function LinearIndex(topology, shape, multi_index):              // 0x1d6b0600
    CHECK(!ElementHasBitWidth(shape, 64))                        // 64-bit handled by component split, line 646
    if shape.layout().tiles().size() > 1:                        // tiled
        return LayoutUtil::LinearIndex(shape, multi_index)       // OSS tiled index math
    // untiled: stamp a default device layout first, then linearize
    tmp = shape
    ForEachMutableSubshape(tmp, UpdateLayout)                    // 0x1d6b05a0 - fill linear layout
    return LayoutUtil::LinearIndex(tmp, multi_index)
```text

> **QUIRK --** `LinearIndex` 直接拒绝 64-bit 元素宽度（`CHECK(!ElementHasBitWidth(shape, 64))`，line 646）。对 f64/s64 缓冲区的寻址是在其 32-bit *component* shapes 上完成的，而从不在 64-bit 逻辑 shape 上完成 -- 这与 `ShapeSizeBytesRaw` 的拆分一致。对 `tiles().size() > 1` 的分支与其他地方的 tiled/linear 测试相同：已填充 tile 列表会路由到 OSS 分块 `LinearIndex`；空列表会先盖上默认 layout。

### 重打包内核

`RepackToHardwareLayout<SUBLANE, 128>`（实例化 `<2,128>` @ `0x1d5c2880`、`<16,128>` @ `0x1d5c3b40`、`<32,128>` @ `0x1d5c3f80`）是手写 AVX 例程，在 DMA 前将连续 host 缓冲区重排为硬件 tile 顺序。它按 **2048-byte（`0x800`）tile** 处理缓冲区（`size >> 11` 个完整 tile + `size & 0x7FF` 剩余）：

```c
function RepackToHardwareLayout<S,128>(dst, size_bytes, src, pad_flag):  // e.g. 0x1d5c3b40
    whole_tiles = size_bytes >> 11          // 2048-byte tiles
    rem         = size_bytes & 0x7FF
    for t in [0 .. whole_tiles):
        // AVX interleave: vpunpcklbw/vpunpckhbw/vpunpcklwd transpose
        // 4-bit/2-bit field packs via vpsllw + vpand with masks
        //   dword_84A2F00 / dword_84A2580 / byte_84A2D4C / dword_84A2A98
        transpose_tile(dst + t*512, src + t*2048)
    if rem:
        copy rem bytes into a scratch tile
        if pad_flag: memset(scratch + rem, 0xFF, 2048 - rem)   // FILL TRAILING PAD WITH 0xFF
        transpose_tile(dst_tail, scratch)

NOTE -- pad_flag 参数就是设备端 tile 填充的物理化：设置时，尾部不完整 tile 会在 SIMD transpose 之前用 0xFF 填充（memset(..., 255, 2048 - rem)）。0xFF 填充值是逻辑范围未填满最后一个 tile 的缓冲区中落入 padding lane/sublane 的字节模式。改用零填充的重新实现会产生字节不同的 HBM 内容（尽管对大多数 reductions 在数值上等价），任何对原始 HBM 取 hash 的 checksum/SDC 路径都会不一致。固定的 2048 tile（= 512 B packed output × 4-way unpack）在 <2/16/32,128> 实例化之间与代际无关。

函数图

函数	地址	角色
`TransferSizeUtil::LinearIndex`	`0x1d6b0600`	tiled/linear multi-index → 扁平 offset
`TransferSizeUtil::UpdateLayout`	`0x1d6b05a0`	盖上默认设备 layout（每 subshape）
`TransferSizeUtil::UpdateLeafLayout`	`0x1d6b08c0`	在单个叶子上盖 tile
`RepackToHardwareLayout<2,128>`	`0x1d5c2880`	SIMD host→tile 重打包，sublane 2
`RepackToHardwareLayout<16,128>`	`0x1d5c3b40`	SIMD host→tile 重打包，sublane 16
`RepackToHardwareLayout<32,128>`	`0x1d5c3f80`	SIMD host→tile 重打包，sublane 32

6. 设备缓冲区驻留记录

目的

一个逻辑缓冲区可以是包含许多叶子的 tuple，每个叶子都位于自己的 HBM 地址。运行时在 xla::ShapedBuffer 中跟踪“哪些设备地址保存此缓冲区的叶子，以及设备端形状是什么”，并跨 C ABI 镜像为 XLA_ShapedBuffer。这是 transfer manager、executable 和 PJRT 传递的驻留记录；分配器交出地址（hbm-allocator.md），此记录则记住它们。

XLA_ShapedBuffer 记录

从 ApiConverter::ToC(const xla::ShapedBuffer&, XLA_ShapedBuffer*)（0xfcb8580）恢复：

text

xla::ShapedBuffer  (the on-device residency record)
  +336  on_device_shape   (xla::Shape; the PADDED device shape from §2)
  +656  device_ordinal    (int)
  +664  buffers           (ShapeTree<DeviceAddressBase>; one leaf per tuple element)

XLA_ShapedBuffer  (C-ABI mirror)
  +0..   on_device_shape  (XLA_Shape, via ApiConverter::ToC)
  +536   device_ordinal   (int)
  +544   addresses[]      (SE_DeviceAddressBase array, 24 B stride)   count @ +552
```text

每个叶子都是一个 `SE_DeviceAddressBase`（`stream_executor::DeviceAddressBase` 镜像），源向量中为 48-byte 记录，复制到一个 24-byte C-ABI 条目 `{opaque_ptr, size, payload}`。`(opaque_ptr, size)` 对就是分配器返回的 HBM `(offset, byte-size)`；`size` 正是该叶子填充后设备形状的 `ShapeSizeBytesRaw`。

> **NOTE --** 该记录存储的是*设备端*（填充后、分块后）形状，而不是 host shape。两个叶子即使有相同的逻辑 host shape，但 `memory_space` 颜色不同，也会是不同的设备形状并占用不同层级。单个 `device_ordinal` 将整个缓冲区固定到一个芯片；此层上的缓冲区永远不会跨设备（跨设备是更上层的 sharding）。`MaybeOwningDeviceMemory` / `ScopedShapedBuffer` 变体（见 `XlaComputationLaunchContext::PopulateOutputs` @ `0xeadbb80`）用一个 ownership 标志包装同一记录，决定丢弃该记录是否释放 HBM -- 见 [buffer-donation-aliasing.md](buffer-donation-aliasing.md)。

### 驱动侧缓冲区句柄

在 XLA `ShapedBuffer` 之下，驱动通过 `MaybeOwningDmaBuffer`（见 `JfHbmWriteQueue` 和 `TpuPxcDriver::WriteToMemoryHelper`）移动字节，这是一个 `(ptr, size, optional<SyncFlag>)` tuple。`optional<SyncFlag>` 是完成握手；对缓冲区的 DMA 在完成时 signal 它。这是驻留记录的线缆端：`ShapedBuffer` 的每叶子地址在 DMA 发起时变成一个 `MaybeOwningDmaBuffer`。SyncFlag 协议见 [sflag-protocol.md](sflag-protocol.md)；约束叶子 offset 的 DMA 对齐下限见 [hbm-dma-alignment.md](hbm-dma-alignment.md)。

### 函数图

| 函数 | 地址 | 角色 |
|---|---|---|
| `ApiConverter::ToC(const xla::ShapedBuffer&, XLA_ShapedBuffer*)` | `0xfcb8580` | 驻留记录 → C ABI |
| `ApiConverter::FromC(XLA_ShapedBuffer*)` | `0xfcb7000` | C ABI → `xla::ShapedBuffer` |
| `ApiConverter::ToC(const DeviceAddressBase&)` | `0xfcb78c0` | 每叶子 `(ptr,size)` → `SE_DeviceAddressBase` |
| `ApiConverter::ToC(const xla::ShapeIndex&)` | `0xfcb82e0` | tuple tree 内的叶子索引 |
| `XlaComputationLaunchContext::PopulateOutputs` | `0xeadbb80` | 执行后构建输出 `ScopedShapedBuffer` |
| `TpuTransferManager::CanShapedBufferBeAccessedNow` | `0xeaba6e0` | 驻留就绪检查 |
| `TpuTransferManager::ReadDynamicShapes` | `0xe9735a0` | 将 dynamic dims 读回 shape |

---

## 相关组件

| 组件 | 关系 |
|---|---|
| `xla::jellyfish::TransferSizeUtil` | 拥有本页中的整个 shape→device-shape 映射、tile 填充和字节定大小 |
| `xla::jellyfish::HardwareLayout`（`PopulateShape` / `ComponentShape`） | 盖上 tile + `memory_space` 并拆分 64/128-bit 元素 |
| `TpuTransferManager`（`0xeab0...` family） | 调用 `HostShapeToDeviceShape` 和大小函数的 C-ABI 表面 |
| `ApiConverter`（`0xfcb7...` family） | 跨 PJRT C ABI 编组 `Shape`/`Layout`/`ShapedBuffer` |
| `RepackToHardwareLayout<S,128>` | 为 DMA 物理实现分块的 SIMD 内核 |

## 交叉引用

- [layout-assignment.md](../compiler/layout-assignment.md) -- 编译期*选择* `minor_to_major` 并盖上本页消费的 `(SublaneCount, LaneCount)` tile 的 pass
- [overview.md](overview.md) -- 六区域内存分类以及为每个叶子层级着色的 `memory_space` enum
- [hbm-allocator.md](hbm-allocator.md) -- 交出驻留记录 `DeviceAddressBase` 中保存的 HBM offset 的 `BestFitAllocator`
- [hbm-dma-alignment.md](hbm-dma-alignment.md) -- 约束每个叶子 offset 的 1024-B DMA 下限与 16-KiB 编译期对齐
- [vmem-allocator.md](vmem-allocator.md) -- 片上 `kAlternate` 层级，其字宽输入每层级 tile/granule
- [buffer-donation-aliasing.md](buffer-donation-aliasing.md) -- 驻留记录上的 `MaybeOwningDeviceMemory` ownership 与 input/output aliasing
- [sflag-protocol.md](sflag-protocol.md) -- 驱动侧 `MaybeOwningDmaBuffer` 上的 `optional<SyncFlag>` 完成握手
- [../compiler/tpu-program-serialization.md](../compiler/tpu-program-serialization.md) -- `ComputationLayout`（input/output device shapes）如何序列化进 executable
- [../pjrt/buffer-and-memory.md](../pjrt/buffer-and-memory.md) -- 此层之上的 `PJRT_Buffer` ABI 与外部引用计数
- [返回索引](../index.md) -- Part X -- On-Chip Memory & DMA

TPU 缓冲区布局 ​

摘要 ​

1. Shape / Layout 数据模型 ​

目的 ​

设备路径读取的 xla::Layout 字段 ​

xla::Shape 字段 ​

算法 -- SetPaddedShape ​

元素打包 -- ElementPackingFactor ​

函数图 ​

5. 线性化与 HBM 重打包器 ​

目的 ​

算法 -- LinearIndex ​

函数图 ​

6. 设备缓冲区驻留记录 ​

目的 ​

XLA_ShapedBuffer 记录 ​

TPU 缓冲区布局

摘要

1. Shape / Layout 数据模型

目的

设备路径读取的 `xla::Layout` 字段

`xla::Shape` 字段

算法 -- SetPaddedShape

元素打包 -- ElementPackingFactor

函数图

5. 线性化与 HBM 重打包器

目的

算法 -- LinearIndex

函数图

6. 设备缓冲区驻留记录

目的

XLA_ShapedBuffer 记录