Obecny stan dynamizmu jest bardziej formalnie opisany w RFC dotyczącym dynamizmu. Na tej stronie znajdziesz ogólne omówienie tego dokumentu oraz ważne interfejsy API i narzędzia do interakcji z programami dynamicznymi.
Terminologia dotycząca dynamizmu i omówienie pomocy
Najpierw omówimy kilka terminów, które pojawią się w tym dokumencie, a także krótko przedstawimy ich obsługę w StableHLO:
Wymiary dynamiczne
Wymiary dynamiczne to wymiary, których rozmiar jest nieznany.
W StableHLO wymiary dynamiczne są reprezentowane za pomocą symbolu ?, czyli tensor<16x?xf32>.
Ograniczona dynamika
Ograniczona dynamika odnosi się do wymiaru dynamicznego, którego wartość ma znaną górną granicę. Zwykle jest to przydatne do wypełniania tensora podczas wykonywania.
W StableHLO ograniczoną dynamikę reprezentujemy za pomocą #stablehlo.bounds jako kodowania tensora, tzn. tensor rzędu 2 z 1 wymiarem dynamicznym ograniczonym do 16 i drugim bez ograniczenia można przedstawić jako tensor<?x?xf32, #stablehlo.bounds<16, ?>>.
StableHLO może reprezentować ograniczoną dynamikę, ale ma ograniczoną obsługę w platformach, która pochodzi z TensorFlow i jest częściowo obsługiwana w PyTorch/XLA.
Nieograniczona dynamika
Nieograniczona dynamika, jak sama nazwa wskazuje, odnosi się do wymiaru dynamicznego, którego rozmiar nie jest ograniczony. Ten rodzaj dynamizmu jest bardzo powszechny w StableHLO, z obsługą JAX, PyTorch/XLA i TF, często używany do eksportowania modeli z dynamicznym rozmiarem pakietu lub długością sekwencji.
W StableHLO po prostu pomijamy kodowanie granic dla tej formy dynamizmu, tzn.tensor<?x?xf32>.
Polimorfizm kształtów
Polimorfizm kształtu to termin, który przejęliśmy z JAX.
Polimorfizm kształtu ma 2 kluczowe implikacje:
- Cała dynamika programu wynika z argumentów wejściowych.
- Cała dynamika dotyczy tylko kształtów tensorów, czyli nie zależy od danych.
Dzięki tym 2 regułom, gdy znane są statyczne kształty programu, możemy wziąć program dynamiczny i w pełni przekształcić go w program statyczny na potrzeby kompilacji (patrz „Etapy kompilacji służące do przekształcania programów dynamicznych”).
Polimorfizm kształtu zwykle wykorzystuje nieograniczoną dynamikę. Jeśli znane kształty argumentów mogą prowadzić do w pełni statycznego programu, nie ma potrzeby zgadywania, jak ograniczyć wartości.
Dynamiczność zależna od danych
Dynamizm zależny od danych odnosi się do dynamicznych rozmiarów wymiarów, które dotyczą danych w tensorze. Typowym przykładem jest funkcja nonzeros, która zwraca indeksy wszystkich elementów, które mają wartość 0 w wartości tensora. Kształtu nie można określić bez oceny danych, ale często można go skompilować za pomocą ograniczonej dynamiki, poświęcając dodatkową pamięć na potencjalny rozmiar tensora wyjściowego.
Wiele dynamicznych operacji zależnych od danych można modelować za pomocą ograniczonej dynamiki, w której określa się górną granicę rozmiaru tensora, a sprzęt zwykle implementuje to za pomocą dopełniania tensora. Obecnie PyTorch/XLA i TensorFlow w pewnym stopniu obsługują dynamizm zależny od danych, ale JAX nie śledzi obecnie operacji, które prowadzą do dynamizmu zależnego od danych.
Eksportowanie programów z dynamicznymi wymiarami
Informacje o eksportowaniu programów z dynamicznymi rozmiarami partii lub długościami sekwencji znajdziesz w naszych samouczkach dotyczących StableHLO:
- Samouczek JAX > Eksportowanie z dynamicznym rozmiarem partii
- Samouczek PyTorch/XLA > Eksportowanie z dynamicznym rozmiarem partii
Fazy kompilatora służące do ulepszania programów dynamicznych
Usuwanie potoku przekazywania danych o dynamicznym dostosowywaniu
Istnieje kilka przydatnych przebiegów do dopracowywania kształtów, które są wygodnie zgrupowane w potoku przebiegów createStablehloRemoveDynamismPipeline:
void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
TypeRange refinedTypes);
Pojedyncze przejścia do dopracowywania dynamiki
Przepustki, które zwykle są przydatne do dopracowywania kształtu, to:
stablehlo-refine-arguments– zastępuje argumenty wejściowe konkretnymi typami tensorów.stablehlo-refine-shapesaby rozpowszechnić nowe informacje o kształcie argumentu wejściowego w całym programie.stablehlo-canonicalize-dynamism– zastępuje dynamiczne operacje ich statycznymi odpowiednikami.stablehlo-check-shape-assertions– sprawdza i usuwa niestandardowe wywołania asercji kształtu.
Aktualne informacje i przykłady znajdziesz w dokumentacji, do której prowadzą linki.
Przykład: jak przydatna jest dynamika i jak mogę z niej korzystać?
Dynamiczność ma wiele zastosowań. Skupimy się tu głównie na typowym przypadku użycia polimorfizmu kształtu – tworzeniu elastycznej reprezentacji wyeksportowanego modelu, która jest zwykle używana do reprezentowania dynamicznego rozmiaru pakietu lub długości sekwencji.
Statyczny model add_one
Aby to zademonstrować, użyjemy tego prostego modelu add_one:
def add_one(x):
return x + 1
Po śledzeniu za pomocą tensor<4xf32> otrzymamy ten program StableHLO:
// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
%0 = stablehlo.add %arg0, %cst : tensor<4xf32>
return %0 : tensor<4xf32>
}
Ten model będzie działać tylko w przypadku argumentów wejściowych o kształcie tensor<4xf32>. Jeśli kiedykolwiek zmienimy rozmiar partii lub długość sekwencji, będziemy musieli ponownie prześledzić kod źródłowy i przekształcić go do StableHLO. Nie ma też gwarancji, że nadal będziemy mieć dostęp do kodu źródłowego.
Dynamiczny model add_one
W takiej sytuacji przydaje się dynamiczna zmiana kształtu. Zamiast tego JAX i PyTorch/XLA mogą emitować model add_one z dynamicznie prawidłowym IR, który będzie rozgłaszać stałą, aby dopasować ją do dynamicznego kształtu danych wejściowych w ten sposób:
// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
%cst = stablehlo.constant dense<1.0> : tensor<f32>
%0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
%1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
%2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
%3 = stablehlo.add %arg0, %2 : tensor<?xf32>
return %3 : tensor<?xf32>
}
Taka reprezentacja modelu jest znacznie bardziej elastyczna i umożliwia odroczone określanie wartości, takich jak rozmiar pakietu czy długość sekwencji. Ten model można wdrożyć na platformach obsługujących dynamiczne kształty (takich jak AI Edge) lub udoskonalić za pomocą przejść dynamicznych wymienionych w tej dokumentacji.
Ulepszanie modelu dynamicznego
Na przykład następująca kolejność przekazywania może w pełni udoskonalić ten program:
stablehlo-opt add_one_dynamic.mlir \
--stablehlo-refine-arguments='types=tensor<16xf32>' \
--stablehlo-refine-shapes \
--stablehlo-canonicalize-dynamism
Program będzie stopniowo przekształcany w ten sposób:
// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
%c = stablehlo.constant dense<16> : tensor<1xi64>
%0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
...
%3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
%4 = stablehlo.add %0, %3 : tensor<?xf32>
return %4 : tensor<?xf32>
}
// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
%c = stablehlo.constant dense<16> : tensor<1xi32>
%0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
%1 = stablehlo.add %arg0, %0 : tensor<16xf32>
return %1 : tensor<16xf32>
}
// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
%0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
%1 = stablehlo.add %arg0, %0 : tensor<16xf32>
return %1 : tensor<16xf32>
}
// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
%cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
%0 = stablehlo.add %arg0, %cst : tensor<16xf32>
return %0 : tensor<16xf32>
}