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çindir. Kılavuz adım adım değildir ve LLVM, Bazel ve XLA hakkında bilgi sahibi olduğunu varsayar.

XLA, yeni bir mimarinin veya hızlandırıcının XLA tarafından üretilen ML programları çalıştırmak üzere bir arka uç oluşturmak için uygulayabileceği soyut bir arayüz sağlar. XLA'nın yeniden hedeflenmesi, yeni donanım için PyTorch veya TensorFlow gibi mevcut her işlemi ön uç çerçevesinde uygulamaktan çok daha basit ve ölçeklenebilir olmalıdır.

Çoğu uygulama, aşağıdaki senaryolardan birine girer:

  1. Mevcut CPU mimarisi, mevcut bir LLVM arka ucuyla veya bu arka ucu olmadan 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.

1. Senaryo: Mevcut CPU mimarisi henüz XLA tarafından resmi olarak desteklenmiyor

Bu senaryoda, mevcut XLA CPU arka ucuna bakarak başlayın. XLA, LLVM kullanarak farklı CPU'ların hedeflenmesini kolaylaştırır. Bunun nedeni, CPU'lar için XLA arka uçları arasındaki temel 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şturulmuş 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çin hedef mimariyi yapılandırmak üzere bir LLVM üçlüsünü sağlayabilir.

Mevcut LLVM arka ucu yoksa ancak başka bir kod oluşturma aracı varsa mevcut CPU arka ucunun büyük bir kısmı yeniden kullanılabilir.

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

Zaten LLVM IR yayan 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 özelliğinin değiştirilmesi gerekebilir ancak mevcut arka uçlarla çok fazla kod paylaşılabilir.

Bunun için iyi bir örnek olarak XLA'nın GPU arka ucu verilebilir. GPU arka ucu, CPU benzeri 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. Altıgen 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 olacaktır.

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

LLVM kullanmak mümkün değilse en iyi seçenek, istenen donanım için XLA'ya yönelik yeni bir arka uç uygulamaktır. Bu seçenek en fazla ç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ıfta, bir HLO hesaplamasının bir xla::Executable kodunda derlenmesi ele alınmaktadır.
  • xla::Executable: Bu sınıf, platformda derlenmiş bir hesaplama başlatmak için kullanılır.
  • xla::TransferManager: Bu sınıf, arka uçların belirli cihaz bellek yerlerinden XLA değişmez verilerini oluşturmak için platforma özel mekanizmalar sağlamasına olanak tanır. Başka bir deyişle, ana makineden cihaza ve geri veri aktarımını kapsamaya yardımcı olur.