此页面由 Cloud Translation API 翻译。
Switch to English

混叠XLA

该文件描述了XLA混叠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在运行时捐赠)。

前端互操作

TF / XLA

在与XLA编译TensorFlow程序的群集,所有资源变量更新在编译时别名(在运行时走样取决于是否任何其他人持有的资源变量张量参考)。