XLA 中的別名

本文件說明 XLA 別名 API,可讓您在建構 XLA 程式時指定輸入和輸出緩衝區之間的別名。

在編譯期間定義別名

例如,請思考一個僅會在輸入內容中加入 1 的簡單 HLO 模組:

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 一樣執行)。