XLA'da Takma Ad Oluşturma

Bu dokümanda, bir XLA programı oluştururken giriş ve çıkış arabellekleri arasındaki takma adı belirtmenize olanak tanıyan XLA takma ad oluşturma API'si açıklanmaktadır.

Derleme zamanında takma ad tanımlama

Örneğin, girişine sadece 1 ekleyen sıradan bir HLO modülünü düşünün:

HloModule increment

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

Bu modül, biri %p girişi ve diğeri %out çıkışı için olmak üzere iki tane 4 baytlık arabellek ayıracaktır.

Bununla birlikte, genellikle güncellemenin yerinde yapılması tercih edilir (örneğin, ifadeyi oluşturan ön uçta, giriş değişkeni hesaplamadan sonra artık aktif değilse (p++ artışında olduğu gibi).

Bu tür bir güncellemeyi verimli bir şekilde gerçekleştirmek için giriş takma adını belirtebilirsiniz:

HloModule increment, input_output_alias={ {}: 0 }

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

Bu biçim, çıkışın tamamının ({} ile işaretlenir) diğer giriş parametresinin 0 olarak atandığını belirtir.

Takma adı programatik olarak belirtmek için XlaBuilder::SetUpAlias API'ye bakın.

Çalışma zamanında takma ad kullanımını tanımlama

Önceki adımda tanımlanan takma ad, derleme sırasında belirtilir. Yürütme sırasında, arabelleğin bağışlanıp bağışlanmayacağını seçmek için LocalClient::RunAsync API'sini kullanabilirsiniz.

Programa verilen giriş arabellekleri ExecutionInput içine alınır. Bu da MaybeOwningDeviceMemory ağacı içerir. Bellek sahiplik olarak belirtilirse (arabelleğin sahipliği XLA çalışma zamanına aktarılır) tampon aslında bağışlanır ve güncelleme, derleme zamanı takma adı oluşturma API'sinin istediği şekilde yerinde yürütülür.

Bununla birlikte, derleme sırasında takma adlı arabellek çalışma zamanında bağışlanmazsa kopya koruması devreye girer: Ekstra bir çıkış arabelleği O ayrılır ve takma ad olarak verilmesi amaçlanan P giriş arabelleğinin içeriği O öğesine kopyalanır (böylece program, O arabelleği çalışma zamanında bağışlanmış gibi çalışabilir).