XLA için yeni bir arka uç geliştirme

Bu kılavuz, XLA'nın, donanımlarını verimli bir şekilde hedefleyen programlar oluşturmasını isteyen sistem mühendisleri için hazırlanmıştır. Kılavuz adım adım açıklamalı değildir ve LLVM, Bazel ve XLA hakkında bilgi sahibi olduğunuz varsayılır.

XLA, XLA tarafından üretilen makine öğrenimi programlarını çalıştırmak için bir arka uç oluşturmak üzere yeni bir mimari veya hızlandırıcının uygulayabileceği soyut bir arayüz sağlar. XLA'yı yeniden hedefleme, mevcut her işlemi yeni donanımlar için PyTorch veya TensorFlow gibi bir ön uç çerçevesinden uygulamaya kıyasla önemli ölçüde daha basit ve ölçeklenebilir olmalıdır.

Çoğu uygulama, aşağıdaki senaryolardan birine karşılık gelir:

  1. Mevcut bir LLVM arka ucu olan veya olmayan mevcut CPU mimarisi henüz XLA tarafından resmi olarak desteklenmemektedir.
  2. Mevcut LLVM arka ucuna sahip CPU benzeri olmayan donanım.
  3. Mevcut LLVM arka ucu olmayan CPU benzeri olmayan donanım.

Senaryo 1: XLA tarafından henüz resmi olarak desteklenmeyen mevcut CPU mimarisi

Bu senaryoda, mevcut XLA CPU arka ucuna bakarak başlayın. XLA, LLVM kullanarak farklı CPU'ları hedeflemeyi kolaylaştırır. Bunun nedeni, CPU'lar için XLA arka uçları arasındaki ana farkın, LLVM tarafından oluşturulan kod olmasıdır.

Donanım tedarikçisinin donanımı için bir LLVM arka ucu varsa arka ucu, XLA ile oluşturulan LLVM'ye bağlamak kolaydır. JIT modunda, XLA CPU arka ucu ana makine CPU'su için kod yayar. xla::AotCompilationOptions, önceden derleme işlemi için hedef mimariyi yapılandırmak üzere bir LLVM (LLVM) üçlü sağlayabilir.

Mevcut bir LLVM arka ucu yoksa ancak başka bir tür kod oluşturucu varsa mevcut CPU arka ucunun çoğunun yeniden kullanılması mümkün olur.

Senaryo 2: Mevcut LLVM arka ucuna sahip CPU benzeri olmayan donanım

LLVM IR yayıldıkları için mevcut xla::CPUCompiler ve xla::GPUCompiler sınıflarında yeni bir xla::Compiler uygulaması modellemek mümkündür. Donanımın yapısına bağlı olarak, LLVM IR neslinin birçok yönünün değiştirilmesi gerekebilir ancak mevcut arka uçlarla çok fazla kod paylaşılabilir.

XLA'nın GPU arka ucu, buna örnek olarak verilebilir. GPU arka ucu, CPU olmayan bir ISA'yı hedeflediğinden kod oluşturma işleminin bazı yönleri GPU alanına özgüdür. Diğer donanım türleri, ör. Hexagon gibi DSP'ler (yukarı akış LLVM arka ucuna sahiptir) LLVM IR emisyon mantığının bazı bölümlerini yeniden kullanabilir ancak diğer parçalar benzersiz olur.

Senaryo 3: Mevcut LLVM arka ucu olmayan CPU olmayan donanım

LLVM kullanmak mümkün değilse en iyi seçenek istenen donanım için XLA için yeni bir arka uç uygulamaktır. Bu seçenek en çok çaba gerektirir. Uygulanması gereken sınıflar şunlardır:

  • StreamExecutor: Birçok cihaz için tüm StreamExecutor yöntemleri gerekli değildir. Ayrıntılar için mevcut StreamExecutor uygulamalarına bakın.
  • xla::Compiler: Bu sınıf, bir HLO hesaplamasının bir xla::Executable içinde derlenmesini içerir.
  • xla::Executable: Bu sınıf, platformda derlenmiş bir hesaplama işlemi başlatmak için kullanılır.
  • xla::TransferManager: Bu sınıf, arka uçların belirli cihaz bellek işleyicilerinden XLA değişmez verileri oluşturmak için platforma özel mekanizmalar sağlamasına olanak tanır. Diğer bir ifadeyle, ana makineden cihaza ve geriye doğru veri aktarımını kapsamaya yardımcı olur.