本文件說明 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
一樣執行)。