W tym dokumencie opisujemy proces przekształcania modułu XLA High Level Optimizer (HLO) od stanu początkowego do końcowego pliku wykonywalnego. Czasami pomijamy słowo „moduł” i określamy go po prostu jako „HLO”.
Wstępna optymalizacja HLO
Zaczynamy od modułu HLO przed optymalizacją. Wstępnie zoptymalizowany HLO nie zawiera operacji (ops), które są uważane za wewnętrzne w XLA, takich jak fusion czy bitcast. Na tym etapie operacje nie mają układu lub, jeśli go mają, zostanie on zignorowany. HLO przed optymalizacją jest zwykle generowany przez platformy wyższego poziomu, takie jak TensorFlow i JAX. Gdy używasz flagi XLA -xla_dump_to, wstępnie zoptymalizowany kod HLO jest zapisywany w pliku z sufiksem „before_optimizations.txt”.
Optymalizacja modułu HLO
Potok XLA:GPU przekształca HLO przed optymalizacją w zoptymalizowany HLO, wykonując sekwencję przebiegów. Przejścia można grupować semantycznie i wykonywać w tej kolejności:
Dzielenie powiązanych dokumentów
Obejmuje to przekazywanie, takie jak ShardyPartitioner, lub przekazywanie w przypadku dzielenia SPMD.
Etapy optymalizacji
Może to obejmować zarówno legalizację, jak i uproszczenie.
Przejścia optymalizacji zbiorczej
Podobne do Optymalizacji, ale skupia się na operacjach zbiorczych.
Przekazywanie przypisania układu
Każda operacja HLO ma przypisany układ, który jest częścią kształtu instrukcji. Układ określa, jak tensor jest fizycznie rozmieszczony w pamięci.
Przykład kształtu z układem:
f32[10,20,30]{2,0,1}
Po typie elementu podane są logiczne wymiary kształtu, a następnie permutacja układu w kolejności od najmniejszej do największej. W tym przykładzie najmniejszy wymiar to 30, drugi najmniejszy to 10, a największy to 20.
Celem przypisywania układu jest zminimalizowanie liczby wymaganych transpozycji fizycznych za pomocą strategii zachłannej. Zaczyna się od pewnych ograniczeń układu (np. biblioteki cuDNN/cuBLAS oczekują kolejnych wymiarów) i propaguje układy „w dół”, a potem „w górę” wykresu HLO. Po zakończeniu propagacji układu niektóre instrukcje mogą mieć sprzeczne układy: jeden propagowany z operandu, a drugi – z użytkownika. Aby rozwiązać ten konflikt, wstawiana jest instrukcja copy HLO, która zmienia układ z układu operandu na układ instrukcji.
Etapy normalizacji układu
Ponieważ ustalenie kształtu fizycznego jest dość trudne, normalizacja układu próbuje przekształcić kształt tak, aby używał domyślnego układu {rank-1, rank-2, …, 0}. W przykładzie powyżej znormalizowany kształt tof32[20,10,30]{2,1,0}. Operacje kopiowania, które zmieniają układy, są przepisywane jako połączenie transpose i bitcast. Obecnie nie możemy normalizować wszystkich operacji, więc niektóre z nich mogą mieć inne układy niż domyślne, zwłaszcza gather i dot. Na granicach między znormalizowanymi i nieznormalizowanymi operacjami będą występować bitcast operacje, które reprezentują transpozycję, tzn. transpozycję z przypisanym układem, który sprawia, że fizycznie nie wykonuje ona żadnej operacji.
Normalizacja układu sprawia też, że niektóre niejawne transpozycje stają się jawne, co jest ważne, ponieważ generator kodu może obsługiwać jawne transpozycje za pomocą dedykowanego emitera. Na przykład operacja zmiany kształtu może mieć inny układ fizyczny między operandem a wynikiem (np. z powodu innego rzędu). ReshapeDecomposer etap, który jest wykonywany w ramach etapów normalizacji układu,ReshapeDecomposer przekształca zmianę kształtu w sekwencję transpose, zmianę kształtu bitcast i transpose.
Optymalizacja przypisywania układu po publikacji
Najważniejsze przekształcenia to fuzje Triton (fuzje GEMM + fuzje Softmax/Layernorm) lub przekształcenia w wywołania biblioteczne. Na tym etapie działa też automatyczne dostrajanie, w ramach którego XLA wybiera różne emitery, najlepszy algorytm dla splotów lub iloczynów skalarnych, najlepsze kafelkowanie dla fuzji obsługiwanych przez emiter Triton itp.
Karty Fusion
Główne metody to łączenie PriorityFusion i Multi-Output.
W PriorityFusion tworzymy fuzje na podstawie modelu kosztów. Podczas łączenia zezwalamy na duplikowanie operacji dla kilku użytkowników, jeśli można je połączyć ze wszystkimi użytkownikami. Umożliwimy też rozszerzanie istniejących fuzji Triton Softmax, jeśli będzie to możliwe.
Multi-Output fusion to osobny karnet, który umożliwia łączenie operacji/fuzji, które mają wspólny operand. Może też łączyć operandy lub ich fuzje z użytkownikami bez duplikowania, dodając dodatkowe dane wyjściowe, dzięki czemu inni użytkownicy operacji, która ma zostać połączona, mogą być przekierowywani do tych danych wyjściowych. Ten etap musi uważać, aby nie wprowadzać cykli do wykresu HLO.
Po połączeniu wielu danych wyjściowych następuje eliminacja wspólnych podwyrażeń (HloCSE), która może ponownie połączyć wcześniej zduplikowane operacje, jeśli znalazły się w tym samym połączeniu.
Kilka przepustek po połączeniu
Kilka przebiegów związanych z operacjami zbiorowymi (np. przekształcanie ich w operacje asynchroniczne lub wymuszanie określonej kolejności względnej operacji zbiorowych).
Na koniec wykonujemy operację CopyInsertion, w której dodawane są kopie, aby operacje wykonywane w miejscu nie zastępowały danych, które są jeszcze potrzebne w innych miejscach.
Po zakończeniu optymalizacji zoptymalizowany HLO jest zapisywany w pliku, którego nazwa ma sufiks „after_optimizations.txt”, jeśli używasz flagi -xla_dump_to. Jeśli chcesz zrzucać HLO po przejściach pośrednich, które faktycznie zmieniają HloModule, możesz użyć flagi -xla_dump_hlo_pass_re=.* (lub konkretnego wyrażenia regularnego, aby ograniczyć ją do określonych przejść).
Harmonogram
Moduł HLO bez harmonogramu nadal ma pewną swobodę w kolejności przetwarzania operacji. Każde sortowanie topologiczne uwzględniające relacje między operandami i wynikami oraz zależności sterujące jest prawidłowe. Planowanie określa, którego konkretnego zamówienia użyć. Głównym problemem na tym etapie jest maksymalne zużycie pamięci, które zależy od czasu życia tensorów. W pierwszym kroku wypróbowujemy różne algorytmy planowania i wybieramy ten, który powinien zminimalizować szczytowe zużycie pamięci. Pamiętaj, że na tym etapie nie pracujemy jeszcze z fizycznymi buforami (zajmiemy się tym w sekcji „Przypisywanie buforów”) i symulujemy wykorzystanie pamięci.
Następnie LatencyHidingScheduler wykonuje przebiegi i stara się zmaksymalizować nakładanie się obliczeń i komunikacji. Może to jednak ponownie zwiększyć wykorzystanie pamięci.
Jeśli szczytowe zużycie pamięci jest wyższe niż ilość dostępnej pamięci, uruchamiamy HloRematerialization. Ta przepustka ma na celu zmniejszenie zużycia pamięci kosztem wydajności, ponieważ np. niektóre fuzje mogą zostać podzielone, a niektóre operacje mogą zostać zduplikowane, aby skrócić czas życia bufora. Jeśli nastąpi ponowna materializacja, warto poszukać sposobów na zmniejszenie wymagań dotyczących pamięci po stronie modelu (np. przez używanie mniejszych rozmiarów partii).
Przypisanie bufora
Bezpośrednio przed obniżeniem do LLVM IR uruchamiamy przebiegi przypisywania buforów, które przypisują wycinki buforów do każdej instrukcji na wykresie HLO. Przypisanie bufora odbywa się w kilku krokach:
HloDataflowAnalysisprzypisujeHloValues(w zasadzie bufory logiczne) do instrukcji. W przypadku operacji wykonywanych w miejscuHloValueoperandu może zostać ponownie użyty. Operator może definiować więcej niż 1HloValue(np. z kształtem wyniku w postaci krotki).HloAliasAnalysispróbuje połączyć bufory na potrzeby operacji wygładzania i oblicza mapowanie zHloValuenaHloBuffer.BufferAssignmentoblicza mapowanieHloBuffersna wycinki bufora w dużym buforze w taki sposób, aby ten sam wycinek bufora nie był używany w przypadku różnychHloBuffersz nakładającymi się okresami istnienia. W przypadku operacji, które mogą mieć aliasy, dopuszczalne jest niewielkie nakładanie się (czas zakończenia jednej operacjiHloBuffermoże zbiegać się z czasem rozpoczęcia drugiej operacjiHloBuffer). Jeśli używasz flagi-xla_dump_to, niektóre informacje o przypisaniu bufora są zapisywane w pliku z sufiksem „after_optimizations-buffer-assignment.txt”.
Thunki
Po zoptymalizowaniu i zaplanowaniu wykresu HLO jest on przekształcany w liniową sekwencję funkcji dla konkretnego backendu (procesora lub GPU).
W XLA thunk to abstrakcja samodzielnej jednostki pracy, którą wykonuje środowisko wykonawcze. Może to być uruchomienie skompilowanego jądra, konkretna operacja, wywołanie biblioteki, konstrukcja przepływu sterowania, komunikacja zbiorowa itp. A Sekwencja Thunk reprezentuje cały plik wykonywalny dla konkretnego backendu.
Emisja Thunk
Proces przekształcania zaplanowanego obliczenia HLO w sekwencję thunków nazywa się „emisją thunków”. Zajmuje się tym dedykowana klasa emitera w każdym backendzie.
W przypadku backendu GPU jest to obsługiwane przez funkcję IrEmitterUnnested.
EmitHloComputation iteruje po zaplanowanej liście instrukcji HLO w obliczeniach i wysyła je do wyspecjalizowanej metody Emit... (np.
EmitFusion, EmitConvolutionThunk, EmitWhile). Każda z tych metod tworzy odpowiednie obiekty Thunk i dołącza je do sekwencji Thunk.
W przypadku procesora ThunkEmitter pełni tę rolę i jest zorganizowany w podobny sposób. Ostateczna wartość ThunkSequence jest osadzona w CpuExecutable.
Pamiętaj, że każda instrukcja w obliczeniach wejściowych modułu HLO może odpowiadać 0 (kTuple, kConstant, ..), 1 lub wielu (np. instrukcja sortowania) thunkom w końcowej sekwencji thunków.
Bufory poleceń: optymalizacja wykonywania na GPU
Nowoczesne procesory graficzne umożliwiają jednokrotne zarejestrowanie sekwencji operacji GPU (uruchomienia jądra, kopiowania pamięci itp.), a następnie wielokrotne odtwarzanie tej sekwencji przy minimalnym obciążeniu procesora. Jest to kluczowa optymalizacja wydajności, szczególnie w przypadku zadań z wieloma małymi, szybko uruchamianymi jądrami. XLA używa bufora poleceń jako abstrakcji wykresów CUDA lub wykresów HIP. Główny interfejs jest zdefiniowany w GpuCommandBuffer.
Bufor poleceń jest reprezentowany w sekwencji thunk przez CommandBufferThunk.
Generator nie tworzy tego thunka bezpośrednio na podstawie instrukcji HLO. Zamiast tego robi to CommandBufferConversionPass, który działa na samym ThunkSequence.
Etap ten identyfikuje ciągłe podciągi zgodnych funkcji (np. serię funkcji KernelThunk i GemmThunk). Następnie zastępuje znaleziony podciąg pojedynczym znakiem CommandBufferThunk. Nowy thunk zawiera logikę pierwotnych thunków w postaci listy lekkich obiektów CommandBufferCmd.
Gdy CommandBufferThunk jest wykonywana po raz pierwszy w danym strumieniu GPU, „zapisuje” sekwencję poleceń w buforze poleceń sprzętowych. Przy każdym kolejnym wykonaniu po prostu wysyła do procesora graficznego jedno polecenie „odtwarzania” zarejestrowanej sekwencji. Pozwala to uniknąć obciążenia procesora związanego z uruchamianiem każdego pojedynczego jądra.
Wykonywalne
Produktem końcowym potoku kompilacji XLA jest samodzielny,wykonywalny plik przeznaczony na konkretną platformę. Ten obiekt zawiera wszystkie informacje potrzebne do uruchomienia skompilowanego programu na urządzeniu docelowym, takim jak procesor lub GPU. Jest to pomost między kompilatorem a środowiskiem wykonawczym. Nowoczesne środowiska wykonawcze, takie jak PJRT, używają nieco bardziej zaawansowanych abstrakcji (patrz PjRtExecutable), ale ostatecznie opakowują one plik wykonywalny specyficzny dla backendu.
Executable zawiera kilka kluczowych informacji wygenerowanych podczas kompilacji. Dokładna zawartość różni się w zależności od backendu, ale zwykle obejmuje:
Skompilowany kod: jest to kod maszynowy niskiego poziomu, który będzie działać na urządzeniu. W przypadku procesorów są to zwykle co najmniej 2 pliki obiektowe. W przypadku procesorów graficznych jest to skompilowany kod urządzenia w formacie PTX lub HSACO, który jest wczytywany do procesora graficznego w czasie działania.
Plan wykonania (ThunkSequence): podstawowy element logiki środowiska wykonawczego. Jest to liniowa sekwencja obiektów Thunk. Każdy thunk reprezentuje pojedynczą jednostkę pracy, np.uruchomienie jądra, wywołanie funkcji biblioteki (np. cuBLAS) lub obsługę przepływu sterowania. Środowisko wykonawcze uruchamia program, iterując po tej sekwencji.
Układ pamięci (BufferAssignment): ten kluczowy element metadanych, generowany przez BufferAssigner, opisuje pełny układ pamięci dla obliczeń. Określa rozmiar każdego bufora oraz sposób przydzielania i ponownego wykorzystywania pamięci na potrzeby parametrów, danych wyjściowych i wartości tymczasowych. Środowisko wykonawcze używa tego do przydzielania pamięci urządzenia i przekazywania prawidłowych wskaźników do każdego bloku kodu.
(opcjonalnie) Moduł HLO: na potrzeby debugowania i profilowania plik wykonywalny często zachowuje odniesienie do końcowego, zoptymalizowanego modułu HloModule, z którego został skompilowany.
Proces tworzenia końcowego pliku wykonywalnego jest koordynowany przez kompilator dla każdego konkretnego backendu. Metoda RunBackend w implementacji kompilatora jest ostatnim krokiem procesu kompilacji, który pakuje wszystkie skompilowane artefakty w obiekt wykonywalny.
GpuCompiler
i
CpuCompiler
są odpowiednio przeznaczone dla procesora graficznego i procesora.
Gdy użytkownik wywoła Execute... w pliku wykonywalnym, środowisko wykonawcze użyje BufferAssignment do przydzielenia pamięci, a następnie wywoła ThunkSequence, aby uruchomić operacje na urządzeniu za pomocą skompilowanego kodu.