Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

MLIR CodeGen dla XLA

XLA działa na HloInstruction i wykonuje wiele optymalizacji na tej reprezentacji, udostępniając wiele z nich między docelowymi urządzeniami. W pewnym momencie obliczany jest harmonogram liniowy, a bufor pamięci jest przypisywany do każdej wartości statycznie. Kodegen specyficzny dla urządzenia działa poprzez przechodzenie przez tę sekwencję i wywoływanie „emiterów” w celu wygenerowania reprezentacji odpowiedniej dla urządzenia (na przykład pojedyncza funkcja LLVM na obliczenia XLA na CPU lub sekwencja „uderzeń” hermetyzujących operacje GPU i prawdopodobnie generowana PTX, gdy ukierunkowane na GPU).

W ramach etapu przejściowego jesteśmy obecnie w trakcie przechwytywania procesu zaraz po tym, jak XLA zakończy fazę przypisywania bufora i zamiast tego emituje moduł lhlo dialekcie lhlo . Stamtąd wykonujemy codegen przy użyciu komponentów MLIR (głównie Linalg, affine i dialekt GPU) w zależności od urządzenia.

Poniżej znajduje się plan przyrostowej migracji XLA / GPU przy użyciu lhlo jako wejścia codegen.

Zadania

Gospodarz Urządzenie
Format wejściowy HloInstruction * (Zadanie 1) HloInstruction * (Zadanie 1)
Format wyjściowy xla :: Thunk (Zadanie 2) LLVM IR (Zadanie 3)
  • Zadanie 1 zmienia format wejściowy hosta i urządzenia z HloInstruction * na LHLO.
  • Zadanie 2 zmienia format wyjściowy hosta z thunks na „jakieś lądowisko dla hosta” (patrz poniżej).
  • Zadanie 3 przenosi dane wyjściowe urządzenia z LLVM IR do jakiejś formy MLIR. W przypadku tego projektu jest to opcjonalne. Szczegółowe informacje można znaleźć w sekcji „Migracja urządzenia LLVM IR”.

Ten projekt kładzie nacisk na posiadanie kompleksowych modeli, które można uruchomić, z włączonymi emiterami LHLO w jak największym stopniu. Oznacza to, że następująca lista celów według priorytetów:

  • Spraw, aby XLA / GPU działały z emiterami LHLO, z istniejącymi Thunkami i emiterami niezmodyfikowanymi.
  • Usuń odniesienia do HloInstruction * w LHLO, przypadek po przypadku:
    • Zmień starszy emiter na emiter oparty na MLIR (np. Linalg) lub
    • Przetłumacz mechanicznie istniejący emiter na reprezentację MLIR (migracja do standardu z dialektem GPU).

Migrating Thunks (zadanie 2)

xla :: gpu :: Thunk to struktura danych, która:

  • Może zostać wywołany z hosta (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Zawiera różne dane w swoich podklasach.
  • Współdziała z BufferAllocation :: Slice i StreamExecutor.
  • Uruchamia jądra
  • Wywołuje wszystkie biblioteki uruchomieniowe.

Koszt obejmuje:

  • Reprezentowanie danych konfiguracyjnych specyficznych dla operacji (np. Konfiguracje splotu).
  • Migracja kształtów op i operandów.
  • Reprezentujące drzewo thunks (podczas, stan itp.).

Migracja jest niezależna od migracji LHLO / emiter. Przy ograniczonych zasobach ma priorytet za migracją LHLO / emiter.

Mamy kilka możliwości obniżenia części po stronie hosta z LHLO:

  • TFRT
    • (Pro) świetne opakowania CUDA i HIP do użytku.
    • (Pro) łatwe do zaimplementowania wywołania bibliotek (cuDNN, cuBLAS, cuFFT, itp.), Ponieważ operacje TFRT są interpretowane przez kod C ++.
    • (Con) strona hosta jest w fazie rozwoju i nie jest testowana.
  • Kod Jitted CPU
    • (Pro) wielka niższa zdolność. Utwórz kilka pętli i warunków i gotowe.
    • (Con) GPUDialect nie modeluje jeszcze łańcuchów / strumieni / asynchroniczności / alokacji urządzeń.
    • (Wady) Obsługa środowiska uruchomieniowego CUDA / HIP jest minimalna (ścieżka do zestawu narzędzi, wersja, ładowanie dynamiczne itp.).
  • Istniejące (interpretujące) środowisko wykonawcze XLA

Decyzja: zastosuj TFRT, ale obsługuj jitting kod procesora w TFRT.

Migracja urządzenia LLVM IR (Zadanie 3)

Emiter elementarny generuje docelową operację, wypełniając ją element po elemencie. Każdy element wyjściowy zależy od zestawu elementów z operandów. Wszystkie elementy są opisane poprzez połączenie bufora z dynamicznymi indeksami. Wystarczy opisać prawie wszystkie operacje matematyczne, ale ze względu na wydajność tylko duży podzbiór operacji matematycznych jest zaimplementowany bezpośrednio w (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter jest wyjątkowy pod tym względem:

  • Duża część kodu jest współdzielona między XLA / GPU i CPU.
  • Reprezentuje dużą część operacji widocznych w modelach, w tym wszystkie operacje oparte na elementach.
  • Większość fuzji zależy wyłącznie od ElementalIrEmitter.
  • Jest strukturalnie prosty, ponieważ opisuje zależność danych DAG między elementami op i elementami operandu.
  • Jest w większości przenośny i wysokopoziomowy (np. W przeciwieństwie do GPU kReduce i GPU kCopy).
  • Dynamiczne wsparcie kształtów jest łatwe w przypadku operacji przynajmniej elementarnych.

Teraz, dla wszystkich operacji, emitowanych elementarnie lub nie, istnieje kilka odmian stanu końcowego każdej operacji XLA:

  1. Kod urządzenia pozostaje jako LLVM IR.
  2. Refaktoryzuj stary emiter tak, aby wyglądał jak LHLO -> MLIR LLVM Dialect:
    • (Koszt) Będzie pracą wyrzucaną, jeśli ostatecznie chcemy przejść na Standard.
    • (Korzyści) To jest łatwe i mechaniczne. Można to zrobić w krótkim czasie.
    • (Korzyści) Nie przynosi więcej korzyści w porównaniu z (1).
  3. Zmień wygląd starych emiterów na LHLO -> MLIR GPU + Standard + Loops:
    • (Koszt) Podniesienie istniejących emiterów do standardu stwarza pewne wyzwania. Wskaźniki i GEP należy przekonwertować na MemRefs i SubViews. Zapewnienie kompletności amdgpu to inna sprawa.
    • (Koszt) XLA / GPU w dużym stopniu opiera się na metadanych LLVM:
      • range dla indeksów bloków / wątków.
      • align , dereferenceable , invariant.load , alias.scope , noalias dla ładowania / przechowywania.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable dla pętli sekwencyjnych.
    • (Korzyści) Może być długoterminowe. Bardziej przenośny.
  4. Przebuduj stare emitery na LHLO -> Linalg i napisz nowe emitery Linalg
    • (Koszt) Jest to przypadek indywidualny. W porównaniu z poprzednimi opcjami nowa implementacja, która odpowiada wydajności XLA, musi przejść przez test porównawczy <-> zoptymalizować przepływ pracy, co może być znacznym kosztem w przypadku niektórych operacji.
    • (Korzyści) zunifikowany stos; Społeczność; ruchliwość; więcej potencjałów optymalizacji.

Wnioski:

  • Nie idź po (2). (1) lub (3) są po prostu lepsze niż (2). (2) kosztuje więcej niż (1), ponieważ wymaga dużo mechanicznej refaktoryzacji. Dzięki (1) nadal możemy osiągnąć cel, jakim jest umożliwienie XLA wychwytywania emiterów MLIR. Odbywa się to poprzez wykonanie LHLO -> LLVM IR -> uruchomienie starszych emiterów urządzeń.
  • Operacje ElementalIrEmitter idą na (4), ale nie przyrostowo. Nie ma sposobu, aby to zrobić op po op, ponieważ wszystkie emitowane elementarnie operacje są połączone w ten sam graf. Ta praca może również służyć jako punkt unifikacji kilku bieżących sił (xla / service / mlir_gpu, generator jądra, Linalg).
  • Wszystkie inne operacje są dostępne dla (1). W celu rozciągnięcia mogą zostać przeniesione do (3) lub (4).

Priorytetyzacja

Chociaż wszystkie trzy zadania wymienione powyżej można zrównoleglenie, przy ograniczonych zasobach muszą być serializowane. Priorytetyzacja koncentruje się na widocznych rezultatach ukończenia każdego zadania.

Priorytetyzacja jest następująca: Zadanie 1 (LHLO dla starszych emiterów)> Zadanie 2 (Thunks)> Zadanie 3 (emiter MLIR).

Pod koniec Zadania 1 użytkownicy XLA mogą wygenerować LHLO (np. Generator jądra) i je wykonać. Format kompilacji nie będzie możliwy do serializacji MLIR.

Pod koniec Zadania 2 LHLO obniża się do właściwego, serializowalnego MLIR. Umożliwia to kompilację w trybie offline.

Pod koniec Zadania 3, wszystkie nadajniki XLA są w swojej implementacji oparte na MLIR.

Szczegółowy projekt

Krok 1: (Zadanie 1) Ukończ LHLO i spraw, aby starsze emitery przejmowały LHLO

Ten krok sprawia, że ​​wszystkie istniejące nadajniki XLA / GPU współdziałają z operacjami MLIR. Ten krok to czysta refaktoryzacja i NFC.

Ten krok jest głównie mechaniczny, ale warto zauważyć następujące rozbieżności między niezagospodarowanym HloComputation a LHLO:

  • Każda HloInstrukcja ma bezpośredni dostęp do swoich argumentów (DAG przepływu danych). Wręcz przeciwnie, każda operacja LHLO ma dostęp tylko do swoich buforów operandów (dwudzielność między operacjami i buforami). Operacje LHLO muszą przejść przez łańcuchy use-def, aby uzyskać dostęp do swoich operacji operandów.
  • Niezaufane tradycyjne emitery empirycznie prawie nigdy nie uzyskują dostępu do swoich operandów. Jedynym wyjątkiem jest kReduce.
  • Niezaufane starsze emitery uzyskują dostęp do BufferAssignment tylko w celu pobierania wycinków, a nie w celu uzyskania dostępu do struktur danych pomocniczych, takich jak dataflow_analysis () lub alias_analysis (). llvm_ir buduje swój własny alias_analysis () na podstawie informacji o wycinku.

Wniosek jest taki, że LHLO powinno pasować prosto bez większych kłopotów.

Krok 2: (opcjonalnie) obsługa profilowania

Ten krok jest potrzebny tylko wtedy, gdy zaczniemy odrzucać część logiki XLA Thunk (zobacz następny krok).

Przed faktycznym włączeniem emiterów opartych na MLIR potrzebujemy profilowania dla emiterów opartych na MLIR.

Obecnie XLA wykonuje własne profilowanie, wywołując licznik czasu StreamExecutor. Licznik czasu pod maską wstawia dwa zdarzenia przed i po uruchomieniu jądra i mierzy czas synchronizacji między tymi dwoma zdarzeniami.

Istnieją trzy podejścia do obsługi profilowania w MLIR:

  • Uruchom profilera od końca do końca
  • Dodaj opcję profilu dla każdej operacji w LHLO, używając wstrzykiwanego profilera.

Podejście „od końca do końca” jest przezroczyste dla MLIR, ale występuje ten sam problem, który sprawia, że ​​XLA nie używa go w pierwszej kolejności: wywołania biblioteki zebrane przez profiler (nvprof / ...) nie mogą łatwo odnosić się do HLO ops. Na przykład cuDNN uruchamia wiele jąder dla każdego HLO i trudno powiedzieć, które jądra odpowiadają któremu HLO.

Podejście „wstrzykiwany profiler” wymaga:

  • LHLO, aby wziąć profiler jako parametr.
  • wstawianie profile.start / profile.end przed i po każdej operacji.
  • przejście z tego obniża profil. {początek, koniec} do implementacji w C ++.

Dokładnego profilowania nie można łatwo wykonać dla operacji generowanych przez MLIR, ponieważ:

  • MLIR nie ma licznika czasu ani nie zależy od TFRT / StreamExecutor.
  • MLIR nie jest łatwo wywoływać funkcje C o skomplikowanych parametrach.

Krok 3: (Zadanie 2) Migracja Thunks

Uwaga: istnieją mniej więcej trzy rodzaje bzdur:

  • KernelThunk, który uruchamia jądro.
  • Sterowanie przepływem thunks, które ma logikę sterowania hostem (sekwencję warunkową, while, for) i uruchamia jądra treści.
  • Biblioteki: cuDNN, cuBLAS, cuFFT, NCCL itp.

Plan jest następujący:

  • Umożliwienie serializacji Thunks (de).
  • Pomóż ulepszyć TFRT do stanu, w którym może obsługiwać tę semantykę.
  • Wraz z poprawą stanu stopniowo migruj poszczególne elementy.

Te działania są tylko częściowo uporządkowane. Rzeczywiste zlecenie wykonania / równoległość inżynierska ma być oceniane na bieżąco.

Krok 4: (Zadanie 3) Migrated ElementalIrEmitter

Gdy profilowanie jest gotowe, możemy uzupełnić i dostroić wszystkie emitery oparte na ElementalIrEmitter w MLIR. Następnie włączamy je domyślnie, zakładając, że wszystkie te emitery oparte na MLIR używają jednego strumienia.

Zauważ, że korzystna jest również migracja ElementalIrEmitter XLA / CPU, ponieważ współdzielą one dużą część kodu.

Po wykonaniu wszystkich testów porównawczych i poszukiwań wydajności (TODO: zdefiniuj parzystość wydajności), włączamy nowy emiter elementarny oparty na MLIR i usuwamy starszą wersję ElementalIrEmitter.

Ten krok zapewnia również łatwe przejścia fusion (operacje zagnieżdżone) na potrzeby późniejszej migracji.

Krok 5: Obsługa wielu strumieni lub upuszczanie

Nie możemy usunąć niektórych emiterów, dopóki nie obsługujemy ich w MLIR lub nie usuniemy tej funkcji. To stosunkowo duża ilość pracy w MLIR i niewielki zysk dla XLA. Powinniśmy zbadać obecnych użytkowników wielostrumieniowych użytkowników XLA / GPU i spróbować usunąć tę funkcję, jeśli jest to uzasadnione.

Krok 6: (Zadanie 3) Migrated Device Ops

Ten krok powoduje migrację wszystkich niezagnieżdżonych operacji, a następnie możemy usunąć wszystkie niezaufane emitery.

To wymaga przepisania / refaktoryzacji dla kCopy i kReduce. Na kReduce już dużo pracowano, więc faktyczna ilość pracy, jaką trzeba wykonać, dopiero się okaże.