PJRT C++ Device API 概览

背景

PJRT 是我们希望添加到机器学习生态系统中的统一设备 API。长期愿景是：

  1. 框架（JAX、TF 等）将调用 PJRT，后者具有对框架不透明的特定于设备的实现；
  2. 每个设备都专注于实现 PJRT API，并且对框架来说是不透明的。

PJRT 同时提供 C API 和 C++ API。在任一层级插入都可以，C++ API 使用类来抽象化一些概念，但与 XLA 数据类型的关联也更紧密。本页重点介绍 C++ API。

PJRT 组件

PJRT 组件

PjRtClient

如需完整参考信息，请访问 pjrt_client.h > PjRtClient。

客户端管理设备与框架之间的所有通信，并封装通信中使用的所有状态。它们具有一组用于与 PJRT 插件交互的通用 API，并且拥有给定插件的设备和内存空间。

PjRtDevice

完整参考资料：pjrt_client.h > PjRtDevice 和 pjrt_device_description.h

设备类用于描述单个设备。设备具有设备说明，可帮助识别其类型（用于识别 GPU/CPU/xPU 的唯一哈希），以及在本地和全局设备网格中的位置。

设备还知道其关联的内存空间以及其所属的客户端。

设备不一定知道与其关联的实际数据缓冲区，但可以通过查看其关联的内存空间来确定这一点。

PjRtMemorySpace

如需完整参考信息，请访问 pjrt_client.h > PjRtMemorySpace。

内存空间可用于描述内存的位置。这些应用可以不固定，可以位于任何位置，但必须能从设备访问；也可以固定，但必须位于特定设备上。

内存空间知道其关联的数据缓冲区、与内存空间关联的设备（多个）以及内存空间所属的客户端。

PjRtBuffer

如需完整参考信息，请访问 pjrt_client.h > PjRtBuffer。

缓冲区以某种格式在设备上保存数据，以便在插件内轻松处理，例如 MLIR 元素属性或专有张量格式。框架可能会尝试以 xla::Literal 的形式（即作为模块的输入实参）将数据发送到设备，该数据必须克隆（或借用）到设备的内存中。当不再需要缓冲区时，框架会调用 Delete 方法进行清理。

缓冲区知道自己属于哪个内存空间，并且可以传递性地确定哪些设备能够访问它，但缓冲区不一定知道自己的设备。

为了与框架通信，缓冲区知道如何转换为 xla::Literal 类型以及如何从 xla::Literal 类型转换：

// Literal to Buffer
absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(...) {...}

// Buffer to Literal
xla::Future<> ToLiteral(xla::MutableLiteralBase* literal) override {...}

用于创建缓冲区的 API 具有缓冲区语义，可帮助确定是否可以共享、复制或更改主机缓冲区中的字面数据。

最后，如果缓冲区分配给框架层中的变量 x = jit(foo)(10)，则缓冲区可能需要比其执行范围更长的生命周期；在这些情况下，缓冲区允许构建外部引用，这些引用提供指向缓冲区所保存数据的临时所有权指针，以及用于解释底层数据的元数据（dtype / 维度大小）。

PjRtCompiler

如需完整参考信息，请访问 pjrt_compiler.h > PjRtCompiler。

PjRtCompiler 类为 XLA 后端提供了有用的实现细节，但插件不必实现该类。从理论上讲，PjRtCompiler 或 PjRtClient::Compile 方法的职责是获取输入模块并返回 PjRtLoadedExecutable。

PjRtExecutable / PjRtLoadedExecutable

完整参考信息请参阅 pjrt_executable.h > PjRtExecutable 和 pjrt_client.h > PjRtLoadedExecutable。

PjRtExecutable 知道如何获取已编译的制品和执行选项，并对其进行序列化/反序列化，以便根据需要存储和加载可执行文件。

PjRtLoadedExecutable 是已编译到内存中的可执行文件，可随时接收输入实参以供执行，它是 PjRtExecutable 的子类。

可执行文件通过客户端的 Execute 方法之一进行交互：

// Execute on addressable devices
absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
Execute(absl::Span<const std::vector<PjRtBuffer*>> argument_handles, ...) {...}

// Execute assigned replica/partition on the specified device
absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
ExecuteSharded(absl::Span<PjRtBuffer* const> argument_handles,
              PjRtDevice* device, ...) {...}

// Execute on specified device, single replica / partition
absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
ExecutePortable(absl::Span<PjRtBuffer* const> argument_handles,
                PjRtDevice* device, ...) {...}

在调用 Execute 之前，框架会将所有必需数据转移到执行客户端拥有的 PjRtBuffers，但会返回以供框架引用。然后，这些缓冲区会作为实参提供给 Execute 方法。

PJRT 概念

Future 和异步计算

如果插件的任何部分是以异步方式实现的，则必须正确实现 future。

请考虑以下计划：

@jax.jit
def foo(x): return x + 1

x = foo(1)
# [...] other logic not using `x`
print(x + 1)

异步插件能够将计算 x 加入队列，并立即返回一个尚未准备好读取的缓冲区，但执行会填充该缓冲区。在 x 之后，执行可以继续将不需要 x 的必要计算（包括在其他 PJRT 设备上的执行）加入队列。一旦需要 x 的值，执行就会阻塞，直到缓冲区通过 GetReadyFuture 返回的 future 声明自身已准备就绪。

Future 可用于确定对象（包括设备和缓冲区）何时可用。

高级概念

除了实现基本 API 之外，还可以扩展 JAX 的功能，以便供插件使用。从某种意义上来说，这些都是选择性加入的功能，因为典型的 JIT 和执行工作流程在没有这些功能的情况下也能正常运行，但对于生产质量的流水线，可能需要考虑对 PJRT API 支持的任何这些功能的支持程度：

  • 内存空间
  • 自定义布局
  • 通信操作，例如发送/接收
  • 主机分流
  • 分片

典型的 PJRT 框架-设备通信

日志示例

以下是为加载 PJRT 插件和执行 y = jax.jit(lambda x: jnp.power(x, jnp.int32(2)))(1) 而调用的方法的日志。在本例中，我们记录了 JAX 与 StableHLO 参考 PJRT 插件的交互。

日志示例 


//////////////////////////////////
// Load the plugin
//////////////////////////////////

I client_cpp_pjrt.cc:55] StablehloReferencePjrtClient(0x23bac400)
I device.cc:53] StablehloReferenceDeviceDescription(0x23bac4f8)
I device.cc:104] StablehloReferenceDevice(0x23bac4e0)
I device.cc:123] client(0x23bac4e0)
I device.cc:123] client(0x23bac4e0)
I client_cpp_pjrt.cc:71] process_index(0x23bac400)
I client_cpp_pjrt.cc:67] platform_name(0x23bac400)
I device.cc:143] AttachDefaultMemorySpace(0x23bac4e0)
I client_cpp_pjrt.cc:67] platform_name(0x23bac400)
I client_cpp_pjrt.cc:86] devices(0x23bac400)
I client_cpp_pjrt.cc:81] addressable_device_count(0x23bac400)
I device.cc:168] description(0x23bac4e0)
I device.cc:168] description(0x23bac4e0)
I device.cc:86] Attributes(0x23bac4f8)
I device.cc:128] IsAddressable(0x23bac4e0)
I device.cc:168] description(0x23bac4e0)
I device.cc:61] process_index(0x23bac4f8)
I device.cc:123] client(0x23bac4e0)
I client_cpp_pjrt.cc:71] process_index(0x23bac400)
I client_cpp_pjrt.cc:81] addressable_device_count(0x23bac400)
I client_cpp_pjrt.cc:95] memory_spaces(0x23bac400)
I device.cc:128] IsAddressable(0x23bac4e0)
I device.cc:168] description(0x23bac4e0)
I device.cc:61] process_index(0x23bac4f8)
I device.cc:123] client(0x23bac4e0)
I client_cpp_pjrt.cc:71] process_index(0x23bac400)
I device.cc:148] memory_spaces(0x23bac4e0)
Creating PJRT Client from client
I client_cpp_pjrt.cc:108] platform_version(0x23bac400)
I client_cpp_pjrt.cc:67] platform_name(0x23bac400)
I device.cc:57] id(0x23bac4f8)
I device.cc:70] device_kind(0x23bac4f8)
I device.cc:70] device_kind(0x23bac4f8)
I device.cc:80] ToString(0x23bac4f8)
I device.cc:80] ToString(0x23bac4f8)
I device.cc:75] DebugString(0x23bac4f8)
I device.cc:75] DebugString(0x23bac4f8)
I device.cc:61] process_index(0x23bac4f8)
I device.cc:128] IsAddressable(0x23bac4e0)
I device.cc:168] description(0x23bac4e0)
I device.cc:61] process_index(0x23bac4f8)
I device.cc:123] client(0x23bac4e0)
I client_cpp_pjrt.cc:71] process_index(0x23bac400)
I device.cc:153] default_memory_space(0x23bac4e0)
I client_cpp_pjrt.cc:71] process_index(0x23bac400)

//////////////////////////////////
// RUN: `y = jax.jit(lambda x: jnp.power(x, jnp.int32(2)))(1)`
//////////////////////////////////

I executable.cc:309] num_partitions(0x240bab70)
I executable.cc:305] num_replicas(0x240bab70)
I executable.cc:309] num_partitions(0x240bab70)
I client_cpp_pjrt.cc:233] BufferFromHostBuffer(0x23bac400)
I buffer.cc:285] CreateMlirBufferFromLiteral
I buffer.cc:98] CreateFromLiteral
I buffer.cc:99] CreateFromLiteral: s32[] 2
I buffer.cc:64] MlirPjrtBuffer(0x240bb050)
I buffer.cc:102] CreateFromLiteral -> 0x240bb050
I buffer.cc:158] device(0x240bb050)
I buffer.cc:154] memory_space(0x240bb050)
I buffer.cc:154] memory_space(0x240bb050)
I executable.cc:328] GetHloModules(0x240bab70)
I executable.cc:240] Execute(0x240bab70)
I executable.cc:197] ExecuteWithReferenceInterpreter(0x240bab70)
I buffer.cc:303] GetAttributeFromBuffer
I buffer.cc:229] IsDeleted(0x240bb050)
I buffer.cc:311] GetAttributeFromBuffer(0x240bb050) -> dense<2> : tensor<i32>
I executable.cc:205] EvalModule:
module @jit attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<i32> {mhlo.layout_mode = "default"}) -> (tensor<i32> {jax.result_info = "", mhlo.layout_mode = "default"}) {
    // ...
    return %3 : tensor<i32>
  }
}
I executable.cc:206] Inputs: [dense<2> : tensor<i32>]
I executable.cc:213] Results: [dense<2> : tensor<i32>]
I device.cc:153] default_memory_space(0x23bac4e0)
I buffer.cc:291] CreateMlirBufferFromAttribute
I buffer.cc:116] CreateFromAttribute
I buffer.cc:64] MlirPjrtBuffer(0x22cea630)
I buffer.cc:122] CreateFromAttribute(dense<2> : tensor<i32>) -> 0x22cea630

//////////////////////////////////
// RUN: `print(y)`
//////////////////////////////////

I buffer.cc:263] GetReadyFuture(0x22cea630)
I buffer.cc:264] GetReadyFuture(0x22cea630)
I buffer.cc:154] memory_space(0x22cea630)
I buffer.cc:154] memory_space(0x22cea630)
I buffer.cc:158] device(0x22cea630)
I buffer.cc:158] device(0x22cea630)
I buffer.cc:154] memory_space(0x22cea630)
I buffer.cc:154] memory_space(0x22cea630)
I buffer.cc:229] IsDeleted(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:158] device(0x22cea630)
I buffer.cc:154] memory_space(0x22cea630)
I buffer.cc:154] memory_space(0x22cea630)
I client_cpp_pjrt.cc:71] process_index(0x23bac400)
I buffer.cc:229] IsDeleted(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:269] IsOnCpu(0x22cea630) # Returns true, allows external references.
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:129] on_device_shape(0x22cea630)
I buffer.cc:168] AcquireExternalReference(0x22cea630)
I buffer.cc:73] MlirClonedExternalReference(0x2404d560)
I buffer.cc:303] GetAttributeFromBuffer
I buffer.cc:229] IsDeleted(0x22cea630)
I buffer.cc:311] GetAttributeFromBuffer(0x22cea630) -> dense<2> : tensor<i32>
I buffer.cc:291] CreateMlirBufferFromAttribute
I buffer.cc:116] CreateFromAttribute
I buffer.cc:64] MlirPjrtBuffer(0x240bb050)
I buffer.cc:122] CreateFromAttribute(dense<2> : tensor<i32>) -> 0x240bb050
I buffer.cc:168] AcquireExternalReference(0x22cea630)
I buffer.cc:73] MlirClonedExternalReference(0x240b6010)
I buffer.cc:303] GetAttributeFromBuffer
I buffer.cc:229] IsDeleted(0x22cea630)
I buffer.cc:311] GetAttributeFromBuffer(0x22cea630) -> dense<2> : tensor<i32>
I buffer.cc:291] CreateMlirBufferFromAttribute
I buffer.cc:116] CreateFromAttribute
I buffer.cc:64] MlirPjrtBuffer(0x23b2db60)
I buffer.cc:122] CreateFromAttribute(dense<2> : tensor<i32>) -> 0x23b2db60
I buffer.cc:263] GetReadyFuture(0x22cea630)
I buffer.cc:264] GetReadyFuture(0x22cea630)