Alias di XLA

Dokumen ini menjelaskan API aliasing XLA, yang memungkinkan Anda menentukan aliasing antara buffer input dan output saat mem-build program XLA.

Menentukan aliasing pada waktu kompilasi

Misalnya, pertimbangkan modul HLO sederhana yang hanya menambahkan 1 ke inputnya:

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Modul ini akan mengalokasikan dua buffering 4 byte: satu untuk %p input, dan satu untuk %out output.

Namun, sering kali lebih disarankan untuk melakukan update langsung (misalnya, jika di frontend yang menghasilkan ekspresi, variabel input tidak lagi aktif setelah komputasi, seperti pada p++ inkremental).

Untuk melakukan update tersebut secara efisien, Anda dapat menentukan aliasing input:

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Format tersebut menentukan bahwa seluruh output (ditandai dengan {}) diberi alias ke parameter input 0.

Untuk menentukan aliasing secara terprogram, lihat XlaBuilder::SetUpAlias API.

Menentukan aliasing saat runtime

Aliasing yang ditentukan pada langkah sebelumnya ditentukan selama kompilasi. Selama eksekusi, Anda dapat menggunakan LocalClient::RunAsync API untuk memilih apakah akan mendonasikan buffer atau tidak.

Buffering input ke program digabungkan dalam ExecutionInput, yang kemudian berisi hierarki MaybeOwningDeviceMemory. Jika memori ditentukan sebagai memiliki (kepemilikan buffer diteruskan ke runtime XLA), buffer akan didonasikan, dan update dijalankan di tempat, seperti yang diminta oleh API aliasing waktu kompilasi.

Namun, jika buffer yang di-aliaskan pada waktu kompilasi tidak didonasikan pada runtime, copy-protection dimulai: O buffer output tambahan akan dialokasikan, dan konten buffer input P yang dimaksudkan untuk dijadikan alias akan disalin ke O (sangat efektif program ini dapat dieksekusi seolah-olah buffer O didonasikan pada runtime).