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敬上 也能使用 Google Cloud CLI 或 Compute Engine API

在執行階段定義別名

上一步驟中定義的別名會在編譯期間指定。 在執行期間，您可以使用 LocalClient::RunAsync敬上 決定是否提供緩衝區的 API。

程式的輸入緩衝區已納入 ExecutionInput、 而該字串中又包含 MaybeOwningDeviceMemory 的樹狀結構。如果記憶體 指定為 owning (緩衝區的擁有權會傳送至 XLA 執行階段)， 緩衝區，更新即會在適當的位置執行， 可能會被編譯時間別名 API 的要求

不過，如果是在編譯期間套用別名的緩衝區，則「不會」在 當執行階段時，copy-protection 會啟動：系統分配額外的輸出緩衝區 O。 以及強制套用別名的輸入緩衝區 P 內容 轉換為 O (因此，程式可以像緩衝區 O 一樣有效地執行) 要在執行階段捐款)。