Architektura XLA

Zadbaj o dobrą organizację dzięki kolekcji Zapisuj i kategoryzuj treści zgodnie ze swoimi preferencjami.

Dlaczego zbudowaliśmy XLA?

Aby XLA współpracował z TensorFlow, mieliśmy kilka celów:

  • Popraw szybkość realizacji. Kompiluj podgrafy, aby skrócić czas wykonywania krótkotrwałych Ops, aby wyeliminować obciążenie środowiska wykonawczego TensorFlow, łączyć operacje potokowe w celu zmniejszenia obciążenia pamięci i specjalizować się w znanych kształtach tensorów, aby umożliwić bardziej agresywną stałą propagację.

  • Popraw wykorzystanie pamięci. Analizuj i planuj wykorzystanie pamięci, w zasadzie eliminując wiele buforów pamięci pośredniej.

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

  • Zredukuj mobilny ślad. Wyeliminuj środowisko wykonawcze TensorFlow, kompilując z wyprzedzeniem podwykres i emitując parę obiekt/plik nagłówkowy, które można połączyć bezpośrednio z inną aplikacją. Wyniki mogą zmniejszyć ślad dla wnioskowania mobilnego o kilka rzędów wielkości.

  • Popraw przenośność. Spraw, aby stosunkowo łatwo było napisać nowy backend dla nowego sprzętu, w którym to momencie duża część programów TensorFlow będzie działać na tym sprzęcie bez modyfikacji. Stoi to w sprzeczności z podejściem polegającym na specjalizacji pojedynczych monolitycznych Ops pod kątem nowego sprzętu, które wymaga przepisania programów TensorFlow, aby mogły korzystać z tych Ops.

Jak działa XLA?

Język wejściowy do XLA nazywa się „HLO IR” lub po prostu HLO (Operacje wysokiego poziomu). Semantyka HLO opisane są na Działanie Semantyka strony. Jest to najwygodniejszy myśleć HLO jako IR kompilatora .

XLA pobiera grafy („obliczenia”) zdefiniowane w HLO i kompiluje je w instrukcje maszynowe dla różnych architektur. XLA jest modułowy w tym sensie, że jest łatwy do gniazda w alternatywnym backend aby kierować jakąś nową architekturę HW . Zaplecze procesora dla procesorów x64 i ARM64, a także zaplecze GPU NVIDIA znajdują się w drzewie źródłowym TensorFlow.

Poniższy diagram przedstawia proces kompilacji w XLA:

XLA jest kilka optymalizacji i analizy przebiegów, które są niezależne od docelowych, takich jak CSE , docelowym niezależne fuzji pracy, a bufor do analizy alokacji pamięci wykonania dla obliczeń.

Po kroku niezależnym od celu, XLA wysyła obliczenia HLO do zaplecza. Backend może przeprowadzać dalsze optymalizacje na poziomie HLO, tym razem mając na uwadze specyficzne informacje i potrzeby docelowe. Na przykład backend XLA GPU może wykonywać fuzję operacji, która jest korzystna dla modelu programowania GPU i określać sposób podziału obliczeń na strumienie. Na tym etapie backendy mogą również dopasowywać pewne operacje lub ich kombinacje do zoptymalizowanych wywołań bibliotek.

Następnym krokiem jest generowanie kodu specyficznego dla celu. CPU i GPU backends dołączone XLA stosowania LLVM dla IR niskim poziomie, optymalizacji i generowania kodu. Te backendy emitują LLVM IR niezbędne do reprezentowania obliczeń XLA HLO w wydajny sposób, a następnie wywołują LLVM, aby wyemitować kod natywny z tego LLVM IR.

Backend GPU obsługuje obecnie procesory graficzne NVIDIA za pośrednictwem backendu LLVM NVPTX; backend procesora obsługuje wiele procesorów ISA.