Opracowywanie nowego backendu dla XLA

Ten przewodnik jest przeznaczony dla inżynierów systemów, którzy chcą, aby XLA umożliwiało wydajne generowanie programów ukierunkowanych na ich sprzęt. Przewodnik nie jest szczegółowy i wymaga znajomości LLVM, Bazel i XLA.

XLA udostępnia abstrakcyjny interfejs, który może wdrożyć nowa architektura lub akcelerator w celu utworzenia backendu do uruchamiania programów ML wyjściowych. Ponowne kierowanie na XLA powinno być znacznie prostsze i skalowalne niż wdrażanie wszystkich istniejących operacji z platformy frontendowej, takiej jak PyTorch czy TensorFlow w przypadku nowego sprzętu.

Większość implementacji pasuje do jednego z tych scenariuszy:

  1. Istniejąca architektura procesora nie jest jeszcze oficjalnie obsługiwana przez XLA, z istniejącym backendem LLVM lub bez niego.
  2. Sprzęt inny niż procesor z istniejącym backendem LLVM.
  3. Sprzęt inny niż CPU bez backendu LLVM.

Scenariusz 1. Obecna architektura procesora nie jest jeszcze oficjalnie obsługiwana przez XLA

W tym scenariuszu zacznij od sprawdzenia istniejącego backendu procesora XLA. XLA ułatwia kierowanie reklam na różne procesory przy użyciu LLVM, ponieważ główną różnicą między backendami XLA dla procesorów jest kod wygenerowany przez LLVM.

Jeśli dostawca sprzętu ma dla swojego sprzętu backend LLVM, można łatwo połączyć ten backend z tą maszyną wirtualną utworzoną w języku XLA. W trybie JIT backend procesora XLA emituje kod dla procesora hosta. W przypadku kompilacji z wyprzedzeniem xla::AotCompilationOptions może udostępnić potrójną maszynę wirtualną LLVM, aby skonfigurować architekturę docelową.

Jeśli nie ma istniejącego backendu LLVM, ale istnieje inny rodzaj generatora kodu, powinno być możliwe ponowne wykorzystanie większości dotychczasowego backendu procesora.

Scenariusz 2. Sprzęt niepodobny do procesora z istniejącym backendem LLVM

Można wymodelować nową implementację xla::Compiler w dotychczasowych klasach xla::CPUCompiler i xla::GPUCompiler, ponieważ emitują już one LLVM.

Dobrym przykładem jest backend GPU XLA. Backend GPU jest kierowane na ISA inne niż procesor, dlatego niektóre aspekty generowania kodu są unikalne dla domeny GPU. Inne rodzaje sprzętu, np. platformy DSP, takie jak Hexagon (z backendem LLVM), mogą wykorzystywać fragmenty logiki emisji podczerwieni LLVM, ale pozostałe są unikalne.

Scenariusz 3. Sprzęt niepodobny do procesora bez istniejącego backendu LLVM

Jeśli nie można użyć LLVM, najlepszym rozwiązaniem jest wdrożenie nowego backendu dla XLA dla żądanego sprzętu. Ta opcja wymaga najwięcej wysiłku. Klasy, które należy zaimplementować, są:

  • StreamExecutor: W przypadku wielu urządzeń nie wszystkie metody StreamExecutor są potrzebne. Aby dowiedzieć się więcej, zapoznaj się z dotychczasowymi implementacjami StreamExecutor.
  • xla::Compiler: ta klasa zawiera kompilację obliczeń HLO w elemencie xla::Executable.
  • xla::Executable: ta klasa służy do uruchamiania skompilowanych obliczeń na platformie.
  • xla::TransferManager: ta klasa umożliwia backendom udostępnianie specyficznych dla platformy mechanizmów do tworzenia danych literału XLA z określonych uchwytów pamięci urządzenia. Innymi słowy, pomaga energetycznie przesyłać dane z hosta do urządzenia i z powrotem.