XLA (Accelerated Linear Algebra) to kompilator systemów uczących się (ML), który optymalizuje algebrę liniową, co poprawia 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 platformy 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ę opracowywanych wspólnie przez czołowe organizacje zajmujące się sprzętem i oprogramowaniem ML. Przed utworzeniem projektu OpenXLA XLA powstała w ramach tego projektu, ale podstawowe cele pozostają takie same:
Zwiększ szybkość wykonywania. Kompiluj podgrafy, aby skrócić czas wykonywania krótkotrwałych operacji i wyeliminować obciążenie środowiska wykonawczego, scalać operacje potoku, aby zmniejszyć obciążenie pamięci, i specjalizować znane kształty tensorów, aby umożliwić bardziej agresywną stałą propagację.
Popraw wykorzystanie pamięci. Analizuj i planuj wykorzystanie pamięci, eliminując wiele pośrednich buforów pamięci.
Zmniejsz zależność od operacji niestandardowych. Wyeliminuj potrzebę wielu niestandardowych operacji, poprawiając wydajność automatycznie połączonych operacji niskiego poziomu do wydajności niestandardowych operacji, które był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ę działała bez zmian. Jest to przeciwieństwo podejścia do specjalizacji w przypadku pojedynczych operacji monolitycznych na nowym sprzęcie, które wymaga przeredagowania modeli, aby można było z nich korzystać.
Jak to działa
Kompilator XLA wykorzystuje wykresy modeli ze platform ML zdefiniowanych w StableHLO i kompiluje je w instrukcje maszynowe dla różnych architektur. StableHLO określa zestaw operacji obsługujących wersje (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:
XLA przeprowadza na wykresie StableHLO kilka wbudowanych procesów optymalizacji i analiz, które są niezależne od celu, np. CSE, fuzję operacji niezależnych od celu i analizę bufora, która służy do przydzielania pamięci środowiska wykonawczego do obliczeń. Na tym etapie optymalizacji XLA również przekształca dialekt StableHLO w wewnętrzny dialekt HLO.
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 konkretnym celem. Na przykład backend GPU może przeprowadzać fuzje operacji, które są przydatne tylko 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 zoptymalizowanych wywołań biblioteki.
Backend generuje następnie kod dostosowany do miejsca docelowego. Backendy procesora i GPU zawarte w środowisku XLA używają LLVM do niskopoziomowego podczerwieni, optymalizacji i generowania kodu. Te backendy emitują sygnał IR LLVM niezbędny do efektywnego reprezentowania obliczeń HLO, a następnie wywołują LLVM, aby wyemitować kod natywny z tej instancji IR LLVM.
W ramach tego procesu kompilator XLA jest modułowy w tym sensie, że można go łatwo umieszczać w alternatywnym backendzie, aby kierować reklamy na nową architekturę HW. Backend GPU obsługuje obecnie procesory graficzne NVIDIA przez backend LLVM NVPTX. Backend procesora obsługuje wiele ISA dla wielu procesorów.