Ten przewodnik jest przeznaczony dla inżynierów systemów, którzy chcą, aby XLA umożliwiało sprawne 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, aby utworzyć backend do uruchamiania programów ML wyjściowych, które są generowane przez XLA. Ponowne kierowanie XLA powinno być znacznie prostsze i bardziej skalowalne niż wdrażanie wszystkich istniejących operacji z platformy frontendowej, takiej jak PyTorch czy TensorFlow w przypadku nowych narzędzi.
Większość implementacji pasuje do jednego z tych scenariuszy:
- Istniejąca architektura procesora nie jest jeszcze oficjalnie obsługiwana przez XLA, niezależnie od obecnego backendu LLVM.
- Sprzęt inny niż procesor z istniejącym backendem LLVM.
- Sprzęt inny niż procesor bez backendu LLVM.
Scenariusz 1: istniejąca 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 za pomocą LLVM, ponieważ główną różnicą między backendami XLA dla procesorów jest kod wygenerowany przez LLVM.
Jeśli dostawca sprzętu ma backend LLVM do swojego sprzętu, można go łatwo połączyć z tą platformą. W trybie JIT backend procesora XLA wysyła 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 istniejącego backendu procesora.
Scenariusz 2. Sprzęt inny niż procesor z istniejącym backendem LLVM
Można modelować nową implementację xla::Compiler na podstawie istniejących klas xla::CPUCompiler i xla::GPUCompiler, ponieważ te już emitują IR LLVM. W zależności od charakteru sprzętu może się zdarzyć, że wiele aspektów generowania IR LLVM będzie musiało zostać zmienione, ale znaczna część kodu zostanie udostępniona istniejącym backendom.
Dobrym przykładem jest backend GPU XLA. Backend GPU jest kierowany na ISA inne niż procesor, dlatego niektóre aspekty jego generowania kodu są unikalne dla domeny GPU. Inne rodzaje sprzętu, np. platformy DSP, takie jak Hexagon (który ma wychodzący backend LLVM), mogą wykorzystywać fragmenty logiki emisji podczerwień podczerwieni LLVM, ale pozostałe elementy są niepowtarzalne.
Scenariusz 3. Sprzęt niebędący procesorem bez istniejącego backendu LLVM
Jeśli nie można użyć LLVM, najlepszym rozwiązaniem jest wdrożenie nowego backendu dla XLA dla wybranego sprzętu. Ta opcja wymaga najwięcej wysiłku. Klasy, które należy zaimplementować, to:
StreamExecutor: W przypadku wielu urządzeń nie wszystkie metodyStreamExecutorsą potrzebne. Szczegółowe informacje znajdziesz w dotychczasowych implementacjachStreamExecutor.xla::Compiler: ta klasa zawiera kompilację obliczeń HLO wxla::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 konstruowania danych literału XLA z odpowiednich uchwytów pamięci urządzenia. Innymi słowy, pomaga energotować przesyłanie danych z hosta do urządzenia i z powrotem.