XLA 中的走样

本文档介绍了 XLA 别名 API，您可以通过该 API 指定 构建 XLA 程序时，输入和输出缓冲区之间的混叠。

在编译时定义别名

例如，假设有一个简单的 HLO 模块，它只将 1 添加到其输入中：

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

此模块将分配两个 4 字节缓冲区：一个用于输入 %p，另一个用于输入 用于输出 %out。

不过，通常希望就地执行更新（例如，如果 在生成表达式的前端中，输入变量不再存在 如增量 p++）。

为了高效执行此类更新，您可以指定输入别名：

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

该格式指定将整个输出（用 {} 标记）别名为 输入参数 0。

要以编程方式指定别名，请参阅 XlaBuilder::SetUpAlias API。

在运行时定义别名

上一步中定义的混叠是在编译期间指定的。 在执行过程中，您可以使用 LocalClient::RunAsync 用于选择是否捐赠缓冲区的 API。

程序的输入缓冲区被封装 ExecutionInput、 而后者又包含一个 MaybeOwningDeviceMemory 树。如果内存 指定为 owning（缓冲区的所有权会传递给 XLA 运行时）， 缓冲区实际上是捐赠的，然后就地执行更新， 。

不过，如果在编译时设置了别名的缓冲区没有在 运行时，copy-protection 将发挥作用：分配一个额外的输出缓冲区 O， 而原本打算设置别名的输入缓冲区 P 的内容会被复制 到 O 中（因此，程序实际上可以执行，就好像缓冲区 O 是 在运行时捐赠）。