Architektura XLA

XLA (Accelerated Linear Algebra) to kompilator systemów uczących się, który optymalizuje algebrę liniową, zapewniając lepszą szybkość wykonywania i wykorzystanie pamięci. Ta strona zawiera krótkie omówienie celów i architektury kompilatora XLA.

Cele

Obecnie XLA obsługuje kilka frontendów ML (w tym PyTorch, TensorFlow i JAX) i jest częścią projektu OpenXLA – ekosystemu technologii kompilatorów open source dla systemów uczących się opracowywanego wspólnie przez czołowe organizacje zajmujące się sprzętem i oprogramowaniem ML. Przed utworzeniem projektu OpenXLA program XLA był opracowywany w ramach projektu TensorFlow, ale jego podstawowe cele się nie zmieniły:

  • Popraw szybkość wykonywania. Kompiluj podgrafy, aby skrócić czas wykonywania krótkotrwałych operacji i wyeliminować obciążenia ze środowiska wykonawczego, scalać operacje potoku, aby zmniejszyć obciążenie pamięci, i specjalizować znane kształty tensorów, aby umożliwić agresywną stałą propagację.

  • Lepsze wykorzystanie pamięci. Analizuj i planuj wykorzystanie pamięci, eliminując wiele pośrednich buforów.

  • Zmniejsz zależność od operacji niestandardowych. Wyeliminuj potrzebę wielu niestandardowych operacji, poprawiając wydajność automatycznie połączonych operacji niskiego poziomu, aby dopasować je do skuteczności operacji niestandardowych, które zostały pierwotnie połączone ręcznie.

  • Większa przenośność. Ułatw napisanie nowego backendu na nowy sprzęt, aby duża część modeli systemów uczących się mogła działać na tym sprzęcie bez modyfikacji. Stanowi to przeciwieństwo specjalizacji w poszczególnych operacjach monolitycznych związanych z nowym sprzętem, co wymaga przepisania modeli, aby można było z nich korzystać.

Jak to działa

Kompilator XLA pobiera wykresy modeli ze platform ML zdefiniowanych w StableHLO i kompilowa je w instrukcje maszynowe dla różnych architektur. StableHLO określa zestaw operacji w wersji (HLO = operacje wysokiego poziomu), który zapewnia warstwę przenośności między platformami ML a kompilatorem.

Ogólnie proces kompilacji, który przekształca wykres modelu w plik wykonywalny zoptymalizowany pod kątem celów, obejmuje te kroki:

  1. XLA wykonuje kilka wbudowanych testów do optymalizacji i analizy na wykresie StableHLO, które są niezależne od celu, na przykład CSE, łączenie operacji niezależnych od celu i analizę bufora do przydzielania pamięci środowiska wykonawczego do obliczeń. Na tym etapie optymalizacji XLA również konwertuje dialekt StableHLO na wewnętrzny dialekt HLO.

  2. XLA wysyła obliczenia HLO do backendu w celu dalszej optymalizacji na poziomie HLO, tym razem z uwzględnieniem informacji i potrzeb związanych z miejscem docelowym. Backend GPU może na przykład przeprowadzać fuzje operacji, które są szczególnie przydatne w przypadku modelu programowania GPU, i określać sposób partycjonowania obliczeń na strumienie. Na tym etapie backendy mogą również dopasowywać określone operacje lub ich kombinacje do wzorców zoptymalizowanych wywołań biblioteki.

  3. Backend wykonuje następnie generowanie kodu dla poszczególnych celów. Backendy procesora i GPU zawarte w pakiecie XLA używają LLVM do obsługi niskiego poziomu podczerwieni, optymalizacji i generowania kodu. Te backendy emitują IR LLVM, niezbędną do reprezentowania obliczeń HLO w efektywny sposób, a następnie wywołują LLVM, aby wyemitować kod natywny z tej bazy IR LLVM.

W ramach tego procesu kompilator XLA jest modułowy, co ułatwia tworzenie boksów w alternatywnym backendzie w celu kierowania na nowatorską architekturę sprzętu HW. Backend GPU obsługuje obecnie procesory graficzne NVIDIA przez backend LLVM NVPTX. Backend procesora obsługuje wiele identyfikatorów ISA dla wielu procesorów.