Aliasing di XLA

Dokumen ini menjelaskan API aliasing XLA, yang memungkinkan Anda menentukan aliasing antara buffer input dan output saat mem-build 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, sebaiknya lakukan update langsung (misalnya, jika di frontend yang menghasilkan ekspresi, variabel input sudah tidak aktif setelah komputasi, seperti pada p++ pertambahan).

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 oleh {}) 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.

Buffer 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-alias pada waktu kompilasi tidak didonasikan saat runtime, copy-protection dimulai: O buffer output tambahan dialokasikan, dan konten dari P buffer input yang dimaksudkan untuk dialias akan disalin ke O (sehingga program dapat dieksekusi seolah-olah O buffer didonasikan pada runtime).