Infeed / Outfeed 队列

本页所有地址都适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（构建 libtpu_lts_20260413_b_RC00，build-id md5 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节）。该镜像未被 strip；反混淆后的 C++ 符号名按原文逐字引用。VA == 文件分析地址。其他版本会有所不同。

摘要

Infeed 和 outfeed 是 TPU 的流式主机↔设备通道：运行中的程序发出会阻塞在硬件 FIFO 上的 Infeed/Outfeed HLO ops，同时主机并发地将输入 literal 推入 infeed 队列，并从 outfeed 队列中排出结果 literal。这在结构上不同于批量缓冲区传输（PJRT_Buffer 复制，见主机↔设备 DMA）：缓冲区复制的目标是某个特定 HBM 分配，并且只完成一次；infeed/outfeed 传输的目标是一个按核心命名的队列，仅由 TpuCoreLocation + 一个小整数队列索引标识，设备侧会在可执行程序运行时按程序顺序消费/产生条目。XLA 参考框架与上游使用的相同，即 TransferManager::TransferLiteralToInfeed / TransferLiteralFromOutfeed 和 PjRtDevice::TransferToInfeed / TransferFromOutfeed，但 libtpu 提供了它的两套并行实现，重新实现者不得将它们混为一谈。

现代路径是 xla::TpuDevice::TransferToInfeed(const LiteralSlice&) / TransferFromOutfeed(MutableBorrowingLiteral)（learning/45eac/research/pjrt/tpu_pjrt_client.cc），即基于 TFRT 原生 tpu::System runtime 的 PJRT 设备表面。它把主机 literal 线性化为设备布局缓冲区，按 TransferSizeUtil 规定的大小将每个缓冲区切分成硬件 span，并通过 tpu::System::EnqueueInfeed / DequeueOutfeed 入队每个 span。这是一个以 TpuCoreLocation 为键、基于回调、完全阻塞的传输。遗留路径是 tensorflow::tpu::TpuTransferManager::TransferLiteralToInfeed / TransferLiteralFromOutfeed(StreamExecutor*, …)，一个很薄的 C-ABI shim，它通过 ApiConverter::ToC 封送 literal，并调用 TPU 驱动的 TfTpu_ExecutorApiFn 表（EnqueueInfeed 槽 +560，TransferLiteralFromOutfeed 槽 +576），最终落到 deepsea::executor::DeepseaExecutor::EnqueueInfeed。两者最终都到达同一套片上队列硬件；差别只在主机侧抽象（TFRT async runtime 对 StreamExecutor C-shim）。

本页负责说明 infeed/outfeed 的队列机制、两个 transfer-manager 入口点、shape/layout 线性化契约，以及阻塞/异步语义。片上队列驱动（TpuInfeedQueue / TpuOutfeedQueue 及其 Pxc/Jxc 代际实现）只说明到足以理解契约的深度；通用批量缓冲区 DMA 机制见主机↔设备 DMA；与这些队列交错执行的 execute 路径见在 Stream 上异步执行。

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

两个入口点 — 基于 tpu::System 的 PJRT TpuDevice::TransferToInfeed/FromOutfeed，以及基于 ExecutorApiFn C-shim 的遗留 TpuTransferManager；并且它们是相互独立的代码路径，不是一方包装另一方。
队列句柄 — 一次传输用 tpu::TpuCoreLocation + 一个 int 队列索引来命名目标，不是用设备地址。队列对象按调用解析（TpuChipConfig::GetInfeedQueues / tpu::System::EnqueueInfeed 的拓扑遍历），调用者从不持有它。
layout/线性化契约 — 主机 literal 的 shape 会被转换为设备 shape（TransferSizeUtil::HostShapeToDeviceShape）、分块平铺并线性化（LiteralLinearizer::LinearizeToBuffers）；outfeed 执行逆过程（Delinearize）。元素类型由 HardwareLayout::SupportedPrimitiveType 门控。
span 切块 + 阻塞语义 — 线性化缓冲区会被拆成 TensorCoreInfeedSpanSizeBytes / TensorCoreMaxOutfeedSpanSizeBytes 大小的 chunk，每个 chunk 作为由回调完成的任务入队到 AsyncTaskGroup / BlockableAsyncTaskGroup，调用会阻塞（WaitTillDone / 在剩余计数谓词上的 Mutex::LockWhenCommon），直到每个 span 完成。


PJRT infeed 入口	`xla::TpuDevice::TransferToInfeed(const LiteralSlice&)` @ `0xf7ff540` (`tpu_pjrt_client.cc:2153`)
PJRT outfeed 入口	`xla::TpuDevice::TransferFromOutfeed(MutableBorrowingLiteral)` @ `0xf7ffca0` (`:2194`)
PJRT outfeed 辅助函数	`xla::TransferFromOutfeedHelper(TpuCoreLocation, Layout, System, MutableBorrowingLiteral)` @ `0xf8436e0` (`outfeed_utils.cc`)
PJRT infeed span 循环	`tpu::TransferLinearizedBufferToInfeed(Span<uint8>, TpuCoreLocation, AnyInvocable, System*)` @ `0xf8d5cc0`
驱动入队/出队	`tpu::System::EnqueueInfeed` @ `0x1d0b5d00`, `tpu::System::DequeueOutfeed` @ `0x1d0b5f00`
遗留 SE infeed 入口	`tensorflow::tpu::TpuTransferManager::TransferLiteralToInfeed(StreamExecutor*, LiteralSlice&)` @ `0xe9721c0`
遗留 SE outfeed 入口	`tensorflow::tpu::TpuTransferManager::TransferLiteralFromOutfeed(StreamExecutor*, MutableBorrowingLiteral)` @ `0xe972660`
遗留 C-shim 槽位	`ExecutorApiFn()+560`（infeed 入队），`+576`（outfeed 出队）；status `+392/+400/+408/+384`
遗留驱动叶节点	`TpuExecutor_EnqueueInfeed` @ `0xeab9680` → `deepsea::executor::DeepseaExecutor::EnqueueInfeed`; `TpuExecutor_DequeueOutfeed` @ `0xeab96c0`
队列句柄	`tpu::TpuCoreLocation` + `int` 队列索引（不是设备地址）
队列对象	`tpu::TpuInfeedQueue` / `tpu::TpuOutfeedQueue`，Pxc/Jxc 每代实现
证据等级	重新实现级 / 与 IDA 反编译逐字节确认（两条路径均端到端追踪）

1. 两个 Transfer Manager，一套队列硬件

libtpu 携带了完整 XLA transfer-manager 表面的两份实现。逐字节确认的拆分如下：

方面	PJRT 路径	遗留 StreamExecutor 路径
Infeed 入口	`xla::TpuDevice::TransferToInfeed` @ `0xf7ff540`	`tensorflow::tpu::TpuTransferManager::TransferLiteralToInfeed` @ `0xe9721c0`
Outfeed 入口	`xla::TpuDevice::TransferFromOutfeed` @ `0xf7ffca0`	`…::TpuTransferManager::TransferLiteralFromOutfeed` @ `0xe972660`
调用者表面	`PJRT_Device_*` / `PjRtDevice`（JAX, PyTorch-XLA）	`xla::LocalClient` / `xla::Service`，TF-TPU op kernels
设备 runtime	`tpu::System`（TFRT async-value 原生）	`TfTpu_ExecutorApiFn` C-ABI 表（`SE_StreamExecutor*`）
封送方式	进程内 `LiteralLinearizer`	`ApiConverter::ToC` → C-shim → 驱动
源码根目录	`learning/45eac/research/pjrt/`	`third_party/tensorflow/compiler/xla/stream_executor/tpu/`
驱动叶节点	`tpu::System::EnqueueInfeed` / `DequeueOutfeed`	`DeepseaExecutor::EnqueueInfeed` / `DequeueOutfeed`

陷阱 — 这些路径并非分层关系。 PJRT TpuDevice 路径对 ExecutorApiFn / SE_StreamExecutor / TpuExecutor 有零引用（通过扫描 TpuClient text 范围逐字节确认）；它直接与 tpu::System 通信。假设 PJRT infeed 会转发到 TpuTransferManager 的重新实现者会出错：PJRT 自己线性化并入队。两条路径只在 TPU 驱动核心（片上队列）内部汇合，且到达那里的最后一跳不同。置信度：CONFIRMED。

还有第三个按代际区分的变体，即 xla::DeepseaTransferManager::TransferLiteralToInfeed @ 0xeac3cc0 / TransferLiteralFromOutfeedLocked @ 0xeac4b80，它是 jellyfish/Deepsea 的 TransferManager 子类；它与 PJRT 路径共享 linearizer/TransferSizeUtil 机制。CPU backend 提供自己的实现（xla::CpuTransferManager @ 0xf93fda0/0xf93fde0 → TransferLiteralToInfeedOnCpu @ 0xf940080），仅在 HLO 落到主机 CPU 设备时使用。本页记录 TPU infeed/outfeed；CPU infeed 是作用域外的 posix 风格主机内存环。

2. 队列句柄和 Layout 契约

2.1 什么会命名一个队列

一次传输不携带设备地址。infeed/outfeed 目标由以下内容命名：

一个 tpu::TpuCoreLocation — 队列所属的物理 TensorCore；以及
一个 int 队列索引 — 使用该核心的若干 infeed/outfeed 队列中的哪一个。

tpu::System::EnqueueInfeed(const TpuCoreLocation&, int queue_id, Span<const uint8>, AnyInvocable<void(const Status&)>)（0x1d0b5d00）会在每次调用时从这些信息解析出实际队列对象：

// tpu::System::EnqueueInfeed   sub_1D0B5D00
// a2 = TpuCoreLocation, a3 = queue_id, a4/a5 = Span<uint8>, a6 = completion callback
chip   = TpuCoreLocation::Chip(core_loc);                 // @0x1d0b5d2f
topo   = system.topology()->chip_to_node(chip);           // vtable +80
node   = topo->core_for(core_loc.core_index);             // vtable +32  (core_loc+0x30)
queue  = node->infeed_queue(queue_id);                    // vtable +48
status = queue->Validate();                               // vtable +32
if (status != OkStatus) { invoke callback(status); return; }
// resolve the per-core TpuAllocator by its shared-memory location
alloc  = system.allocators().at(core_loc.LocalSharedMemory(0));   // flat_hash_map<TpuSharedMemoryLocation, unique_ptr<TpuAllocator>>
view   = alloc->MapDmaBuffer(span_ptr, span_len);          // vtable +168
if (view.mapped)  queue->Enqueue(view, …, callback);       // vtable +56  (mapped DMA buffer)
else              queue->EnqueueImpl(span, callback);      // vtable +40  (raw span)
```text

因此队列句柄是*临时的*：调用者持有一个 `TpuCoreLocation` 和一个索引，runtime 在每次 enqueue 时遍历 `topology → chip-node → core → infeed_queue[id]`。PJRT `TpuDevice` 中没有持久队列指针。`DequeueOutfeed`（`0x1d0b5f00`）是针对 `outfeed_queue(queue_id)` 的镜像。

> **注意 — 实践中队列索引为 0。** 两个 PJRT span 循环（`TransferLinearizedBufferToInfeed` @ `0xf8d5cc0` 和 `TransferFromOutfeedHelper` 中的 outfeed 循环）对每个 span 都传入 `queue_id = 0`；索引参数为多队列核心而存在，但 literal-transfer 路径使用队列 0。每核心的 infeed 队列*列表*来自 `TpuChipConfig::GetInfeedQueues(TpuCoreType)` @ `0x20afcc80`，并且 `TpuDevice::TransferToInfeed` 在线性化前会 RET_CHECK `!infeed_queues.empty()`（`tpu_pjrt_client.cc:2158`）。置信度：CONFIRMED。

### 2.2 shape → 设备 layout 线性化

主机 literal 不能原样进入队列；设备需要它自己的 tiled layout。两个 PJRT 入口点上的契约相同：

1. **主机 shape → 设备 shape。** `xla::jellyfish::TransferSizeUtil::HostShapeToDeviceShape(out_shape, system, host_shape, 0, 0, 1)` 计算设备上的 shape（tiling、padding）。Infeed：`TpuDevice::TransferToInfeed` 第 `0xf7ff540+…` 附近。Outfeed：`TransferFromOutfeed` @ `0xf7ffca0`。
2. **元素类型门控。** Outfeed 会验证 `xla::jellyfish::HardwareLayout::SupportedPrimitiveType(element_type)`；不支持的 dtype 会返回 `Unimplemented("Attempted to transfer array of shape %s from a TPU device. Transferring data with element type %s has not been implemented on TPUs.")`（`outfeed_utils.cc`）。Infeed 依赖 linearizer 自身的检查。
3. **线性化（infeed）/ 反线性化（outfeed）。** Infeed 调用 `xla::jellyfish::LiteralLinearizer::LinearizeToBuffers(system, literal, device_shape, &buffers, …)` 产生一个设备 layout 字节缓冲区向量；最后一个参数 `v12 = (GetInfeedQueues()[0]+64 == 1)` 会选择一种打包变体。Outfeed 分配一个 `posix_memalign` 对齐的接收缓冲区，dequeue 到其中，然后 `xla::jellyfish::LiteralLinearizer::Delinearize(topology, device_shape, raw_buffer, byte_count/4, layout, system)` 将原始设备字节重塑回调用者的 `MutableBorrowingLiteral`。
4. **Tuple 处理（outfeed）。** `TransferFromOutfeed` 遍历 `ShapeUtil::TupleElementCount`，并对每个 tuple leaf 调用一次 `TransferFromOutfeedHelper`（每个 leaf 都获得索引 `i` 处自己的 `MutableBorrowingLiteral` view）；非 tuple shape 则调用一次。Outfeed 还为 64 位和 128 位元素宽度提供了快速 `TransferX64OrX128FromOutfeedHelper` 分支（`ElementHasBitWidth(64)` / `(128)`）。

> **陷阱 — layout 必须填充 tile。** `TransferFromOutfeedHelper` 先运行 `LayoutUtil::ValidateLayoutForShape`，然后运行 `TransferSizeUtil::GetCompactTiles`。如果设备 layout 没有 tile，它会失败并报 `"Device layout of an array needs to be populated with tiles, got <layout>"`（`outfeed_utils.cc:92`）。重新实现者若把未 tiled 的 layout 喂给 outfeed，会得到 `InvalidArgument`，而不是静默的原始复制。置信度：CONFIRMED。

---

## 3. PJRT 路径 — Span 切块和阻塞

### 3.1 Infeed：`TransferToInfeed` → `TransferLinearizedBufferToInfeed`

将 literal 线性化为缓冲区向量后，`TpuDevice::TransferToInfeed` 会在设备的 infeed mutex（`this+376`）下，把每个缓冲区调度进一个 `tpu::BlockableAsyncTaskGroup`（大小等于缓冲区数量），对每个缓冲区调用 `tpu::TransferLinearizedBufferToInfeed`，然后在 `WaitTillDone` 上阻塞调用线程：

```c
// xla::TpuDevice::TransferToInfeed   sub_F7FF540   (tpu_pjrt_client.cc:2153)
RET_CHECK(!infeed_queues.empty());                       // :2158
device_shape = TransferSizeUtil::HostShapeToDeviceShape(system, literal.shape());
linear_ok    = LiteralLinearizer::LinearizeToBuffers(system, literal, device_shape,
                                                      &buffers, …, pack_variant);
if (linear_ok != Ok) return InitRep(linear_ok);          // :2176

mutex.lock(this+376);                                     // infeed serialization lock
group = BlockableAsyncTaskGroup{ total = buffers.count }; // util.h:124 CHECK total>0
for (buf : buffers) {
    CHECK(--unscheduled_count_ >= 0);                     // util.h:138 capacity guard
    next_done = group.NextDone();                         // per-task completion callback
    queue = walk to first ready infeed queue node;        // skip nodes with (flags&3)!=0
    tpu::TransferLinearizedBufferToInfeed(buf.ptr, buf.len, core_location,
                                          next_done, queue /*+64*/);   // @0xf8d5cc0
}
group.WaitTillDone();                                     // Mutex::LockWhenCommon on count==0
// destroy group, release buffers, unlock

tpu::TransferLinearizedBufferToInfeed（0xf8d5cc0）是逐缓冲区 span 循环。它会把设备缓冲区拆分成 TransferSizeUtil::TensorCoreInfeedSpanSizeBytes 字节大小的硬件 span，并逐个入队：

// tpu::TransferLinearizedBufferToInfeed   sub_F8D5CC0
// a1 = buffer bytes, a2 = byte_count, a3 = TransferSizeUtil*, a4 = callback, a5 = System*
span = TransferSizeUtil::TensorCoreInfeedSpanSizeBytes(util, core, 1, callback);  // @0xf8d5d…
n_full   = byte_count / span;
n_tasks  = n_full + (byte_count % span != 0);            // ceil-div task count
group    = AsyncTaskGroup(operator new(64,16), n_tasks, callback);
ptr = a1;
while (remaining >= span) {                              // full spans: enqueue in place
    CHECK(--unscheduled_count_ >= 0);                    // util.h:52
    tpu::System::EnqueueInfeed(system, core, /*queue=*/0, ptr, span, group.NextDone());
    ptr += span; remaining -= span;
}
if (remaining > 0) {                                     // tail span: pad to span size
    pad_len = TensorCoreInfeedSpanSizeBytes(util, remaining-core, 1, …);
    posix_memalign(&tail, 0x20, pad_len);                // 32-byte aligned
    memcpy(tail, ptr, remaining);
    memset(tail + remaining, 0, pad_len - remaining);    // zero-pad the partial span
    tpu::System::EnqueueInfeed(system, core, 0, tail, pad_len, owning_callback);  // frees tail
}
```text

重新实现者必须复现两个细节：(a) 每个 span enqueue 都是一次独立的 `EnqueueInfeed`，并带有由 `AsyncTaskGroup` 提供的自己的完成回调；(b) 最后一个不完整 span 会被复制进一个新的 `posix_memalign(32)` 缓冲区，并**零填充**到完整 span 宽度后再 enqueue（设备读取完整 span）。拥有所有权的回调变体会在完成时释放这个 padded 缓冲区。

### 3.2 Outfeed：`TransferFromOutfeed` → `TransferFromOutfeedHelper`

outfeed 镜像路径会计算 padded 设备字节数，然后逐 span dequeue 到一个对齐的接收缓冲区，按计数谓词阻塞，并执行反线性化：

```c
// xla::TransferFromOutfeedHelper   sub_F8436E0   (outfeed_utils.cc)
ValidateLayoutForShape(core, layout);                          // :82 on failure
device_shape = HostShapeToDeviceShape(...);  GetCompactTiles(...);  // tile check :92
if (!HardwareLayout::SupportedPrimitiveType(elt)) return Unimplemented(...);
padded = TransferSizeUtil::ShapeSizeCompactForDma(topology, device_shape);
RET_CHECK(padded % sizeof(uint32_t) == 0);                     // :143
if (elt is x64/x128 array && fast-path layout)                 // direct into literal
    dst = MutableLiteralBase::untyped_data(literal);
else { posix_memalign(&dst, 0x20, padded); }                   // staging buffer
max_span = TransferSizeUtil::TensorCoreMaxOutfeedSpanSizeBytes(topology);
remaining = padded; off = 0; outstanding = 0;
while (remaining > 0) {
    chunk = min(remaining, max_span);
    ++outstanding;                                             // count guarded by mutex
    tpu::System::DequeueOutfeed(core, system, /*queue=*/0, dst+off, chunk, on_chunk_done);
    off += chunk; remaining -= chunk;
}
Mutex::LockWhenCommon(predicate: outstanding == 0);            // block until all spans done
// then reshape device bytes back into the literal
LiteralLinearizer::Delinearize(topology, device_shape, dst, padded/4, layout, system);  // :231
free(dst);

TransferFromOutfeed @ 0xf7ffca0 会在整个 tuple-leaf 循环期间持有设备的 outfeed mutex（this+0x180），并为每个 leaf 调用 TransferFromOutfeedHelper。layout 验证失败时，helper 在 outfeed_utils.cc:82 返回；dequeue 错误会通过 CreateStatusAndConditionallyLog(212, …) 暴露；反线性化失败会通过 :231 暴露。

怪癖 — outfeed 会先 dequeue 到 staging buffer，再反线性化；infeed 先线性化，再 enqueue。 两者在缓冲区处理上并不对称。Infeed 预先产生 N 个设备缓冲区，并将每个缓冲区作为 span 流式传输。Outfeed 将原始设备字节接收到一个对齐缓冲区（或在 x64/x128 fast path 中直接接收到 literal 的 untyped_data），然后再运行 Delinearize，把字节布置成主机 literal。试图逐 span 反线性化的重新实现者会破坏多 span shape。置信度：CONFIRMED。

3.3 阻塞 / 异步语义

两个 PJRT 入口点对调用者是同步的：TransferToInfeed 阻塞在 BlockableAsyncTaskGroup::WaitTillDone，TransferFromOutfeed 阻塞在 Mutex::LockWhenCommon，直到每个 span 的完成回调都已触发。设备侧是异步的：每个 EnqueueInfeed/DequeueOutfeed 在把 span 交给队列后立即返回，TPU 驱动完成时会触发 AnyInvocable<void(const Status&)> 回调，从而递减 outstanding-task 计数。主机线程停在 absl::Mutex 条件上，而不是自旋。设备侧错误会作为回调的 Status 参数交付，并作为传输的返回 status 传播出去。每设备序列化由 infeed mutex（TpuDevice+376）和 outfeed mutex（TpuDevice+0x180）强制执行：同一设备上的两个并发 TransferToInfeed 调用不会交错它们的 span。

4. 遗留 StreamExecutor 路径

tensorflow::tpu::TpuTransferManager 是 LocalClient/Service 和 TF-TPU op kernels 使用的 xla::TransferManager 子类。它的 infeed/outfeed 是纯 C-ABI 封送，libtpu 自身代码中没有线性化；驱动在 shim 后面完成这件事。

// tensorflow::tpu::TpuTransferManager::TransferLiteralToInfeed   sub_E9721C0
fn  = ExecutorApiFn();
ctx = fn[360]();                          // make SE_Status / status carrier   (slot +360)
ApiConverter::ToC(literal_slice);         // marshal LiteralSlice → XLA_Literal C struct
fn[560]( this->se_handle,                 // ExecutorApiFn slot +560 = infeed enqueue
         executor->se_executor,           // *(StreamExecutor+7)
         &c_literal, ctx );
ApiConverter::Destroy(&c_literal);
status = fn[408](ctx) ? Ok               // slot +408 = StatusOk?
       : MakeRep(fn[400](ctx)*4|1,       // slot +400 = code, +392 = message ptr
                 fn[392](ctx), strlen,   // status_helper.h:38
                 38, "…/status_helper.h");
fn[384](ctx);                            // slot +384 = free status carrier
return status;
```text

Outfeed（`TransferLiteralFromOutfeed` @ `0xe972660`）除了会封送**两个** C struct（`ApiConverter::ToC(literal.shape())` 到一个 shape carrier，*以及* `ApiConverter::ToC(LiteralSlice(literal))` 到一个 literal carrier），并调用 **slot +576** 而不是 +560 之外，其余相同。叶级导出 `TpuExecutor_EnqueueInfeed` @ `0xeab9680` 和 `TpuExecutor_DequeueOutfeed` @ `0xeab96c0` 会转发到同一二进制中的 `deepsea::executor::DeepseaExecutor::EnqueueInfeed`（@ `0x1d0dbbe0`）/ `DequeueOutfeed`（@ `0x1d0dc160`）层：

```c
// TpuExecutor_EnqueueInfeed   sub_EAB9680   (the C-ABI export the +560 slot targets)
result = deepsea::executor::DeepseaExecutor::EnqueueInfeed(*executor);
// result is an absl::Status; refcount-managed into *out_status

注意 — 槽号就是契约。 遗留路径唯一在 libtpu 可见的旋钮是 TfTpu_ExecutorApiFn 表偏移：+560 = infeed enqueue，+576 = outfeed dequeue，+360 = status-carrier ctor，+392/+400/+408/+384 = message/code/ok/destroy。这些是 TPU 驱动 C-ABI 的一部分（单独任务；不是完整表）。遗留路径的重新实现者要复现封送和这些槽调用，而不是 linearizer，因为 layout 由驱动拥有。槽号和调用形态的置信度：CONFIRMED。DeepseaExecutor::EnqueueInfeed(int, Span<uint8>) 主体确实存在于这个二进制中（0x1d0dbbe0；携带回调的重载为 0x1d0dbe20）：它取得 DeepseaPlatform::ScopedRef，通过 executor 的队列表（vtable +48）解析每索引队列对象，并在已经封送好的 span 上调用该队列的 enqueue（vtable +48）。从 libtpu.so 看仍然不透明的是在此之前驱动执行的 layout/线性化，以及队列对象之下的硅 FIFO 写入，而不是 executor enqueue 本身。

5. 片上队列对象（驱动深度）

在 tpu::System::EnqueueInfeed/DequeueOutfeed 下方是一个按代际区分的队列对象。符号集确认存在 tpu::TpuInfeedQueue / tpu::TpuOutfeedQueue 抽象基类及按硅实现：

类	代际	关键方法（已确认符号）
`tpu::TpuInfeedQueuePxcDriverImpl`	Pxc（Pufferfish/BarnaCore）	`EnqueueImpl(Span<uint8>, AnyInvocable)` @ `0xe802600`
`tpu::TpuOutfeedQueuePxcDriverImpl`	Pxc	`RawDequeue(Span<uint8>/DmaBuffer*, AnyInvocable)` @ `0xe803de0`/`0xe803e80`
`tpu::TpuInfeedQueueJxcDriverImpl`	Jxc（Jellyfish）	`EnqueueImpl(TpuMappedDmaBufferView/Span, …)` @ `0xe7380e0`/`0xe737c60`
`tpu::TpuOutfeedQueueJxcDriverImpl`	Jxc	`RawDequeue(...)` @ `0xe73a100`/`0xe73a320`, `MapDmaBuffer` @ `0xe73a520`
`tpu::TpuPxcDriver`	Pxc 驱动	`EnqueueInfeed(TpuInfeedQueueOnChip, Span, AnyInvocable)` @ `0xe8109c0`, `DequeueOutfeed` @ `0xe811380`
`tpu::JfSoftwareInfeedQueueController`	Jxc 软件队列	`ThrottledSharedMemoryWrite(...)` @ `0xe738800`/`0xe738900`

Pxc infeed enqueue 展示了队列索引分派和 sparsecore offload 分支：

// tpu::TpuInfeedQueuePxcDriverImpl::EnqueueImpl   sub_E802600
queue_index = this[16];                                  // signed int
driver      = this[136];                                 // TpuPxcDriver*
if (queue_index < 0)                                     // negative index = BarnaCore/sparsecore
    TpuPxcDriver::EnqueueBarnaCoreHmf(driver);           // HMF (host-memory-feed) offload
else
    TpuPxcDriver::EnqueueInfeed(driver, on_chip_queue, queue_index, span_ptr, span_len, cb);
```text

重新实现者在这一层只需要知道：(a) 队列由通过拓扑解析的整数索引选择，(b) 负索引会改路由到 BarnaCore/sparsecore host-memory-feed 路径，而不是 TensorCore infeed FIFO，以及 (c) 实际硅 enqueue/dequeue（DMA 缓冲区映射、FIFO 写入）位于 TPU 驱动核心内部，从这个二进制看是不透明的。`TpuChipConfig::InfeedQueue::PerCore` / `Hardware::PerCore` 变体类型（`0xe7274c0`, `0xe7386c0`）确认队列按核心配置，并具有硬件或软件表示。类结构和分派的置信度：HIGH；硅 FIFO 机制的置信度：LOW。

---

## 6. 重新实现说明

- **选对入口。** PJRT consumer（JAX/PyTorch-XLA）使用 `TpuDevice::TransferToInfeed`/`FromOutfeed`，从不接触 `TpuTransferManager`。只有 `LocalClient`/`Service`/TF-TPU kernels 使用遗留 `TpuTransferManager` + C-shim。实现其中之一并不等于实现另一个。
- **句柄是 `{TpuCoreLocation, int}`。** 不要把 infeed/outfeed 队列建模成设备指针。runtime 在每次 enqueue 时重新解析 `topology → chip-node → core → queue[id]`；literal-transfer 路径始终使用 `queue_id = 0`。
- **两个方向都必须有 layout。** Infeed：`HostShapeToDeviceShape` → `LinearizeToBuffers`。Outfeed：dequeue raw → `Delinearize`。Outfeed 还要求一个已填充 tile 的 layout（`GetCompactTiles`）和一个受支持的元素类型（`SupportedPrimitiveType`）；字节数必须满足 `% 4 == 0`。
- **进入队列的是 span，不是 buffer。** 按 `TensorCoreInfeedSpanSizeBytes`（infeed）/ `TensorCoreMaxOutfeedSpanSizeBytes`（outfeed）切分每个设备 buffer。把尾部不完整的 infeed span 在 32 字节对齐缓冲区中零填充到完整 span 宽度。
- **阻塞是 API 语义。** 两个 PJRT 入口都会阻塞调用者，直到所有 span 回调触发（`BlockableAsyncTaskGroup::WaitTillDone` / `Mutex::LockWhenCommon`）。设备侧通过逐 span 的 `AnyInvocable<void(const Status&)>` 回调异步运行；错误以回调的 `Status` 到达。每设备 infeed/outfeed mutex 会序列化并发传输。
- **这是流式通道，不是缓冲区复制。** 对比 [主机↔设备 DMA](../dma/host-device-dma.md)：批量 `PJRT_Buffer` 复制以一个 HBM 分配为目标，并只完成一次；infeed/outfeed 则喂给一个按程序顺序的 FIFO，运行中的可执行程序在遇到 `Infeed`/`Outfeed` HLO ops 时消费/产生它。

---

## 交叉引用

- [主机↔设备 DMA](../dma/host-device-dma.md) — *批量缓冲区*传输路径（`TpuRawBuffer` copies, UHI/HIB DMA spans）；对比：infeed/outfeed 是流式队列路径，不同于批量传输
- [在 Stream 上异步执行](execute-async-on-stream.md) — execute 路径，其运行中的程序驱动这些队列所服务的 `Infeed`/`Outfeed` HLO ops
- [Stream 语义](stream-semantics.md) — enqueue/dequeue 回调完成所依赖的 `tpu::System` 排序 / sequence-point 模型
- [主机回调](host-callbacks.md) — 兄弟主机↔设备控制通道（`DoHostCallbackWithStatus`），即运行中的程序抵达主机的另一种方式
- [UHI 主机接口 DMA](../dma/uhi-host-interface.md) — profiler 中承载这些传输片上侧的 host-interface DMA band
- [Runtime 概览](overview.md) — transfer managers 和 `tpu::System` 在 libtpu runtime stack 中的位置

Infeed / Outfeed 队列 ​

摘要 ​

1. 两个 Transfer Manager，一套队列硬件 ​

2. 队列句柄和 Layout 契约 ​

2.1 什么会命名一个队列 ​

3.3 阻塞 / 异步语义 ​

4. 遗留 StreamExecutor 路径 ​

5. 片上队列对象（驱动深度） ​