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 树。如果将内存指定为“拥有”（缓冲区的所有权将传递给 XLA 运行时），则系统会实际捐赠缓冲区，并会根据编译时别名 API 的请求就地执行更新。

不过，如果在运行时未捐献在编译时别名的缓冲区，就会启动复制保护：系统会分配额外的输出缓冲区 O，并将原本要设为别名的输入缓冲区 P 的内容复制到 O 中（以便程序实际上可以像在运行时捐赠缓冲区 O 一样执行）。