Aliasing di XLA

Dokumen ini menjelaskan XLA {i>aliasing API<i}, yang memungkinkan Anda menentukan {i>aliasing<i} di antara {i>buffer <i}input dan {i> output<i} saat membangun program XLA.

Mendefinisikan 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 buffer 4 byte: satu untuk %p input, dan satu untuk %out output.

Namun, pembaruan sering kali lebih baik dilakukan di tempat (misalnya, di frontend yang menghasilkan ekspresi, variabel input tidak lagi aktif setelah komputasi, seperti pada penambahan p++).

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 ini menentukan bahwa seluruh output (ditandai dengan {}) diberi alias ke parameter input 0.

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

Menentukan aliasing pada runtime

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

{i>Buffer<i} input ke program dibungkus dalam ExecutionInput, yang pada gilirannya berisi pohon MaybeOwningDeviceMemory. Jika memori ditentukan sebagai memiliki (kepemilikan buffer diteruskan ke runtime XLA), {i>buffer<i} sebenarnya didonasikan, dan pembaruan dieksekusi di tempat, ketika yang diminta oleh API aliasing waktu kompilasi.

Namun, jika buffer yang diberi alias pada waktu kompilasi tidak didonasikan di runtime, copy-protection dimulai: buffer output tambahan O dialokasikan, dan konten buffer input P yang dimaksudkan untuk di-aliaskan telah disalin menjadi O (sehingga program dapat berjalan secara efektif seolah-olah buffer O didonasikan pada runtime).