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

Bu ön kılavuz, TensorFlow'u donanımlarına verimli bir şekilde kolayca yeniden hedeflemek isteyen ilk kullanıcılar içindir. Kılavuz adım adım değildir ve LLVM , Bazel ve TensorFlow hakkında bilgi sahibi olunduğunu varsayar.

XLA, TensorFlow grafiklerini çalıştırmak için bir arka uç oluşturmak için yeni bir mimarinin veya hızlandırıcının uygulayabileceği soyut bir arabirim sağlar. XLA'yı yeniden hedefleme, yeni donanım için mevcut tüm TensorFlow Op'u uygulamaktan çok daha basit ve ölçeklenebilir olmalıdır.

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

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

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

Bu senaryoda, mevcut XLA CPU arka ucuna bakarak başlayın. XLA, CPU'lar için XLA arka uçları arasındaki temel fark, LLVM tarafından oluşturulan kod olduğundan, LLVM kullanarak TensorFlow'u farklı CPU'lara yeniden hedeflemeyi kolaylaştırır. Google, x64 ve ARM64 mimarileri için XLA'yı test eder.

Donanım satıcısının donanımı için bir LLVM arka ucu varsa, arka ucu XLA ile oluşturulmuş LLVM ile bağlamak kolaydır. JIT modunda, XLA CPU arka ucu, ana bilgisayar CPU'su için kod yayar. Önceden derleme için xla::AotCompilationOptions , hedef mimariyi yapılandırmak için bir LLVM üçlüsü sağlayabilir.

Mevcut bir LLVM arka ucu yoksa ancak başka bir tür kod oluşturucu varsa, mevcut CPU arka ucunun çoğunu yeniden kullanmak mümkün olmalıdır.

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

Mevcut xla::CPUCompiler ve xla::GPUCompiler sınıflarında yeni bir xla::Compiler uygulamasını modellemek mümkündür, çünkü bunlar zaten LLVM IR yaymaktadır. Donanımın doğasına bağlı olarak, LLVM IR oluşturma özelliklerinin birçoğunun değiştirilmesi gerekebilir, ancak mevcut arka uçlarla çok sayıda kod paylaşılabilir.

İzlenecek iyi bir örnek, XLA'nın GPU arka ucudur. GPU arka ucu, CPU benzeri olmayan bir ISA'yı hedefler ve bu nedenle kod üretiminin bazı yönleri GPU alanına özgüdür. Diğer donanım türleri, örneğin Hexagon (yukarı akış LLVM arka ucuna sahip olan) gibi DSP'ler, LLVM IR emisyon mantığının parçalarını yeniden kullanabilir, ancak diğer parçalar benzersiz olacaktır.

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

LLVM'yi 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 fazla çabayı gerektirir. Gerçekleştirilmesi gereken sınıflar aşağıdaki gibidir:

  • StreamExecutor : Birçok cihaz için StreamExecutor tüm yöntemlerine ihtiyaç duyulmaz. Ayrıntılar için mevcut StreamExecutor uygulamalarına bakın.
  • xla::Compiler : Bu sınıf, bir HLO hesaplamasının derlenmesini bir xla::Executable içine alı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 aygıt bellek tanıtıcılarından XLA değişmez verileri oluşturmak için platforma özgü mekanizmalar sağlamasına olanak tanır. Başka bir deyişle, ana bilgisayardan cihaza ve geri veri aktarımının kapsüllenmesine yardımcı olur.