Dialekt 'sdy'

Dialekt Shardy (SDY) definiuje reprezentację podziału tensorów na podstawie osi oraz dodatkowe komponenty interfejsu API służące do dołączania podziału do tensorów.

Operacje

sdy.all_gather (sdy::AllGatherOp)

Wykonuje komunikację all-gather wzdłuż osi

Składnia:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Zbiera fragmenty tensora wzdłuż osi określonych w gathering_axes.

gathering_axes to lista list osi. Lista zewnętrzna jest większa niż wymiary tensora. Każda lista wewnętrzna określa osie, wzdłuż których należy wykonać oddzielne zbieranie danych w odpowiednim wymiarze. Zostanie ono zastosowane do podziału operandu (tensor), aby uzyskać podział wyniku (out_sharding).

Pamiętaj, że parametr out_sharding nie jest używany do określania podziału wyników. Zamiast tego podział na fragmenty jest określany przez podział operandu i elementu gathering_axes, a element out_sharding musi być zgodny z tym przypuszczalnym podziałem.

Przykład:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

Ograniczenia:

• Musi spełniać ograniczenia wymienione w dokumentacji Sdy_CollectiveOpInterface.
• Elementy w elementach gathering_axes muszą spełniać ograniczenia wymienione w elementach AxisRefListAttr.
• Zastosowanie gathering_axes do podziału operandu daje out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, Sdy_CollectiveOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.all_reduce (sdy::AllReduceOp)

Wykonywanie komunikacji All-Reduce wzdłuż osi

Składnia:

operation ::= `sdy.all_reduce` $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Zmniejsz elementy tensora wzdłuż osi określonych w reduction_axes. Kolejność elementów reduction_axes nie ma znaczenia dla wyniku, ale może wpływać na kolejność odpowiednich grup replik.

Ograniczenia:

• Musi spełniać ograniczenia wymienione w dokumentacji Sdy_CollectiveOpInterface.
• reduction_axes musi spełniać ograniczenia wymienione w AxisRefListAttr;
• reduction_axes nie może pokrywać się z osią podziału operandu;

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.all_slice (sdy::AllSliceOp)

Wykonuje operację dynamicznego przekroju wzdłuż osi

Składnia:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

dzieli elementy tensora wzdłuż osi określonych w slicing_axes. Istnieje algebroidalna dualność między sdy.all_slice a sdy.all_gather.

slicing_axes to lista list osi. Lista zewnętrzna jest większa niż wymiary tensora. Każda lista wewnętrzna określa osie, wzdłuż których należy wykonać przecięcie w odpowiednim wymiarze. Zostanie ono zastosowane do podziału operandu (tensor), aby uzyskać podział wyniku (out_sharding).

Pamiętaj, że parametr out_sharding nie jest używany do określania podziału wyników. Zamiast tego podział na fragmenty jest określany przez podział operandu i elementu slicing_axes, a element out_sharding musi być zgodny z tym przypuszczalnym podziałem.

Przykład:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

Ograniczenia:

• Elementy w elementach slicing_axes muszą spełniać ograniczenia wymienione w elementach AxisRefListAttr.
• Musi spełniać ograniczenia wymienione w dokumentacji Sdy_CollectiveOpInterface.
• Zastosowanie slicing_axes do podziału operandu daje out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.all_to_all (sdy::AllToAllOp)

Wykonuje komunikację typu „wszystko do wszystkich” wzdłuż osi

Składnia:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

W przypadku każdej tupla (axes, src_dim, tgt_dim) na liście parametrów ta operacja dzieli elementy tensora wzdłuż wymiaru tgt_dim i osi określonych w elementach axes, rozprasza te elementy wzdłuż osi i konkatenuje je wzdłuż wymiaru src_dim.

Ta operacja jest w podstawie kombinacją zbiorczego zbierania wzdłuż osi src_dim i axes, a następnie zbiorczego wycinania wzdłuż osi tgt_dim i axes, czyli sufiks wymiaru podziału osi src_dim w wewnętrznym tensorze jest dołączany do wymiaru podziału osi tgt_dim w wyjściowym tensorze.

Operacja „all-to-all” zostanie zastosowana do podziału operandu (tensor), aby uzyskać podział wyniku (out_sharding).

Pamiętaj, że parametr out_sharding nie jest używany do określania podziału wyników. Zamiast tego podział na fragmenty jest określany przez podział operandu src_dim, tgt_dim i axes, a wartość out_sharding musi być zgodna z tym przypuszczalnym podziałem.

Przykład:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

Ograniczenia:

• Musi spełniać ograniczenia wymienione w dokumentacji Sdy_CollectiveOpInterface.
• Lista parametrów nie może być pusta.
• W przypadku każdego parametru w sekcji params:
  • Elementy w elementach axes muszą spełniać ograniczenia elementu AxisRefAttr.
  • src_dim i tgt_dim muszą być prawidłowymi wymiarami (nieujemnymi i mniejszymi od wymiaru tensora).
  • Wartości src_dim lub tgt_dim muszą być unikalne we wszystkich parametrach.
  • src_dim musi być posortowany w kolejności rosnącej według wszystkich parametrów.
• Przeniesienie axes z src_dim do tgt_dim w ramach podziału operandu powoduje otrzymanie out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, Sdy_CollectiveOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.collective_permute (sdy::CollectivePermuteOp)

Wykonuje komunikację z kolekcjonowaniem permutacji w celu zastąpienia osi

Składnia:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Wysyła fragment tensora wejściowego z każdego urządzenia na inne, aby zmienić kolejność osi lub zastąpić je w tensorze.

Zbiór permutacji może przekształcić podział na części danych wejściowych tak, aby każda wymiar był podzielony na części danych w taki sam sposób jak wcześniej, czyli aby był podzielony na części danych wzdłuż osi, których iloczyn rozmiarów odpowiada iloczynowi osi, które wcześniej podzieliły tensor.

Jest to przydatne, gdy chcesz zmienić kolejność osi w pojedynczym wymiarze lub w różnych wymiarach albo zamienić osi podzielonych na kopie z osiami replikowanymi.

W przykładzie poniżej rozmiar podzielonego tensora to tensor<1x4x2xf32>, a ta wartość jest zachowana przez zbiorczą permutację.

Przykład:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

Ograniczenia:

• Musi spełniać ograniczenia wymienione w dokumentacji Sdy_CollectiveOpInterface.
• Jeśli dzielenie na części wejściowe i wyjściowe ma różne siatki, siatki te muszą mieć dokładnie te same osie i różną kolejność identyfikatorów urządzeń.
• W przypadku każdego wymiaru iloczyn rozmiarów podziału osi w elementach out_sharding musi być zgodny z elementami podziału wymiaru operandu.

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.constant (sdy::ConstantOp)

Operacja na stałej

Tworzy tensor output z elementu stałego value.

Patrz: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Przykład:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

Cechy: AlwaysSpeculatableImplTrait

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Wyniki:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Operacja krawędziowa Przepływ danych

Składnia:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

Krawędzia przepływu danych w operacji X definiuje most między zbiorem źródeł (każde z nich jest albo operandem X, albo operandem terminatora bloku X) i zbiorem docelów (każdy z nich jest albo wynikiem X, albo argumentem bloku X), tak aby wszystkie źródła i docelowe miały taką samą liczbę fragmentów.

Operacja może mieć wiele krawędzi przepływu danych, które są do siebie prostopadłe.

Na przykład:

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

Ta operacja while ma n krawędzi przepływu danych, a i-ta krawędź przepływu danych znajduje się między źródłami x_i i return_value_i oraz celami y_i, pred_arg_i i body_arg_i.

Funkcja sdy.data_flow_edge przyjmuje jako argument właściciela krawędzi (może to być dowolna z docelowych wartości, ale najlepiej wynik op zamiast argumentu block), który nie powinien mieć żadnego innego zastosowania. Ta opcja nie jest czysta, ponieważ może przyjmować dane wejściowe, które pierwotnie nie miały żadnego zastosowania.

sdy.data_flow_edge zawiera też opcjonalne podziały dla wszystkich celów na krawędzi. Podczas propagacji należy zaktualizować te podziały zamiast podziału celów (jeśli można go zastosować). Jest to przydatne, gdy opcja ma wiele krawędzi, ponieważ znacznie wydajniej można:

• rozprzestrzeniać się po każdej krawędzi osobno.
• zaktualizować podział każdego krawędzi oddzielnie zamiast wszystkich docelowych naraz (np. operacja ma jeden niezmienny TensorShardingPerValueAttr dla podziału wyników).
• dodawać każdą krawędź do listy zadań osobno, gdy zmieni się podział źródła.

Propagowanie spowoduje rozpropagowanie podziału na wszystkie źródła i cele funkcji sdy.data_flow_edge tak, jakby była to zwykła operacja ze źródłami jako operandami i celami jako wynikami oraz z tożsamością sdy.op_sharding_rule. Oznacza to, że propagacja do przodu odbywa się ze źródeł do miejsc docelowych, a propagacja wsteczna – z miejsc docelowych do źródeł.

Nie zezwalamy na definiowanie danych wejściowych sdy.data_flow_edge przez opcję SdyDialect, więc możemy założyć, że jest ona zdefiniowana przez opcję z atrybutem sdy.sharding, który nie jest zarejestrowany.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.manual_computation (sdy::ManualComputationOp)

Operacja równoległości na wielu urządzeniach z ręcznymi zbiorami

Składnia:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

Przejdź do regionu napisanego w języku kodu lokalnego na urządzeniu z wyraźnymi zbiorami, w których kształty logiczne pasują do lokalnych kształtów bufora fizycznego na urządzeniu, a zbiory dokładnie odpowiadają fizycznej komunikacji między urządzeniami.

Ciało jest lokalne względem manual_axes. Propagowanie będzie odbywać się przez ciało na dowolnych wolnych osiach – tych, które nie znajdują się na liście manual_axes.

Ograniczenia:

• Elementy w elementach in_shardings i out_shardings muszą spełniać ograniczenia wymienione w elementach TensorShardingAttr.
• Liczba globalnych i lokalnych wejść/wyjść tensora w regionie operacji musi być taka sama.
• Osie ręczne muszą znajdować się przed dowolnymi osiami swobodnymi w każdej wymiarze wymiarowania.
• Osie ręczne nie mogą wprowadzać wypełnień. Mianowicie rozmiar wymiaru musi być podzielny przez odpowiedni rozmiar osi ręcznej.
• Globalny i lokalny kształt argumentów/wyników op regions muszą być takie same.
• Nie ma podziałów osi ręcznych.

Cechy: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfejsy: ShardableDataFlowOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.mesh (sdy::MeshOp)

Siatka nazwana

Składnia:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

Definiuje nową siatkę o nazwie. Wszystkie siatki w module muszą mieć tę samą liczbę urządzeń (z wyjątkiem siatek z jednym identyfikatorem device_id). Sieć jest operacją Symbol, która pojawia się w SymbolTable modułu i do której można się odwoływać za pomocą jej name.

Cechy: HasParent<ModuleOp>

Interfejsy: Symbol

Atrybuty:

sdy.named_computation (sdy::NamedComputationOp)

Nazwa operacji obliczeniowej

Składnia:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

Grupuje obliczenia, czyli blok operacji, i nadaje mu nazwę. Rozchodzenie się zmian będzie przebiegać w regionie tak, jakby wszystko było wbudowane.

Można go użyć do propagowania instrukcji wywołania do innych funkcji. Wszyscy użytkownicy Shardy powinni napisać pass import/export, który konwertuje ich wywołania funkcji na operacje sdy.named_computation, duplikując lub kopiując ciało wywoływanej funkcji do ciała named_computation.

Typ argumentów każdego bloku i zwracanych wartości w regionie musi być taki sam jak typ operandów i typ wyników operatora.

Przykład:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

Cechy: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Operacja bariery propagacji

Składnia:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

Ta operacja działa jak operacja tożsamościowa, zwracając tę samą wartość, którą przyjęła jako wejście. Jednak w przypadku propagacji pozwoli to na przepływ tylko w określonym kierunku.

Zapobiega to propagowaniu podziału między użyciami wyniku operacji zapory i jej operandu.

• FORWARD oznacza, że elementy mogą być dzielone tylko od operandu do wyniku.
• BACKWARD oznacza, że elementy mogą przepływać tylko z wyniku do operandu.
• NONE oznacza, że żaden podział nie może być propagowany przez tę opcję.
• Nie można określić BOTH, ponieważ ta opcja byłaby zbędna.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

sdy.reshard (sdy::ReshardOp)

Zmienia podział tensora na inny

Składnia:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

Przeszarpuje wejściowy tensor za pomocą określonego podziału, który różni się od dotychczasowego podziału wejściowego tensora.

Zarówno ShardingConstraintOp, jak i ReshardOp dołączają podział do tensora. Ich żywotność:

1. Przed propagowaniem podziału na fragmenty użytkownicy dodają operację ShardingConstraintOp.
2. Rozprzestrzenianie partycjonowania wykorzystuje operację ShardingConstraintOp. W wynikach propagacji podziału nie ma operacji ShardingConstraintOp. Zamiast tego w razie potrzeby można dodać ReshardOp.
3. Moduły partycjonowania przekształcają operację ReshardOp w operację zbiorczą (lub operację tożsamości). W wynikach partycjonowania nie powinno być operacji ReshardOp.

// TODO(b/331680067). Dodaj wzór kanoniczny, aby usunąć zbędne operacje reshard.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

sdy.return (sdy::ReturnOp)

Operacja sdy.return kończy regiony przypisane do operacji opartych na regionie sdy oraz wszystkie inne operacje oparte na regionie Shardy. Jest zmienna: jako argumenty przyjmuje listę wartości, których typy mogą być dowolne (ale tego samego rodzaju, np. AnyTensor), i dlatego może być używana wielokrotnie na różnych poziomach zbioru reguł IR Shardy.

Składnia:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

Cechy: AlwaysSpeculatableImplTrait, Terminator

Interfejsy: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Operandy:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Ogranicza tensor do określonego podziału

Składnia:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

Dodaje podział do pośredniego tensora (np. wyniku matmul), aby wskazać, że tensor lub jego podzbiór powinien być podzielony w ten sposób.

Jeśli dzielenie na części ma otwarte wymiary i nieograniczone osie, oznacza to, że tensor można dalej dzielić na części według otwartych wymiarów.

Ta opcja może:

• nie mają zastosowania (wiszące) – oznacza to, że dołączone shardingiowanie określa sposób, w jaki należy podzielić tensor wejściowy;
• Have uses – oznacza to, że dołączone podziały to sposób, w jaki należy dzielić elementy ograniczeń podziału, podczas gdy inne elementy tensora wejściowego mogą mieć inny podział (jeśli tensor wejściowy nie ma innych zastosowań, zachowanie jest takie samo jak w przypadku braku zastosowań).

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.sharding_group (sdy::ShardingGroupOp)

Ogranicza tensory w grupie, aby miały ten sam podział.

Składnia:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

Ta opcja udostępnia interfejs do przypisywania tensorów do grup podziału (grup tensorów, które będą miały identyczne podziały). Podczas propagacji, gdy tylko jeden element grupy zostanie podzielony na fragmenty, wszystkie pozostałe elementy zostaną podzielone w identyczny sposób. Ta operacja przyjmuje identyfikator grupy argumentów i nie zwraca żadnego wyniku, ale zamiast tego modyfikuje wewnętrzną reprezentację grupy dzielenia na części, aby dodać wejściowy tensor do grupy o danym identyfikatorze.

Interfejsy: InferTypeOpInterface

Atrybuty:

Operandy:

Atrybuty

AllToAllParamAttr

Parametr „wszystko do wszystkich”

Składnia:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

Tupla zawierająca osie i wymiary źródła/docelowego, na których ma być wykonywane działanie typu all-to-all.

Parametry:

AlltoAllParamListAttr

Lista parametrów typu „wszystko do wszystkich”

Składnia:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

Parametry:

AxisRefAttr

Odniesienie do pełnej osi lub podzielonej podosi

Składnia:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

Ograniczenia:

• W ograniczonym zakresie MeshAttr musi być obecna wartość name.
• Jeśli występuje parametr sub_axis_info, musi on spełniać ograniczenia parametru SubAxisInfoAttr.

Parametry:

AxisRefListAttr

Lista odwołań osi

Składnia:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

Ograniczenia:

• Elementy w elementach value muszą spełniać ograniczenia elementu AxisRefAttr.
• Nie ma żadnych duplikatów odwołań do osi ani podosi, które się na siebie nakładają.
• Żadne 2 sąsiednie odwołania do osi nie są kolejnymi podosiami tej samej pełnej osi, tzn. nie można ich scalić w jedną podoś lub pełną oś.

Parametry:

DimMappingAttr

Lista indeksów czynników wymiaru

Pusta lista wskazuje, że jest to mapowanie puste (jest ono analizowane i drukowane z użyciem *), co oznacza, że wymiar nie jest mapowany na żadne czynniki.

Ograniczenia:

• występuje co najmniej 1 indeks czynnika;
• Indeksy czynników muszą mieścić się w zakresie [0, $factor_sizes].
• Jeśli jest ich więcej niż 1, żaden z nich nie może mieć rozmiaru 1.
• Brak zduplikowanych indeksów czynników.

Parametry:

DimensionShardingAttr

Podzielanie wymiarów

Lista nazw osi, według których należy podzielić wymiar tensora od głównej do podrzędnej, wartość logiczna wskazująca, czy wymiar można podzielić na mniejsze części, oraz opcjonalna liczba całkowita określająca priorytet podziału wymiaru, który będzie uwzględniany podczas propagacji podziału. Priorytety pochodzą z adnotacji podziału użytkowników, a niższa wartość oznacza wyższy priorytet. Jeśli w adnotacji brakuje priorytetu, przyjmuje się najwyższy priorytet.

Ograniczenia:

• Elementy w elementach axes muszą spełniać ograniczenia wymienione w elementach AxisRefListAttr.
• Jeśli podział wymiaru ma priorytet:
  • Priorytet jest większy lub równy 0.
  • Wymiar ma co najmniej 1 oś, jeśli jest zamknięty.

Parametry:

ListOfAxisRefListsAttr

Lista list odwołań osi

Składnia:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

Parametry:

ManualAxesAttr

Lista osi, na których operacja ManualComputationOp jest ręczna

Składnia:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

Parametry:

MeshAttr

Siatka osi i lista urządzeń

Składnia:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

Siatka to lista osi i opcjonalna lista identyfikatorów urządzeń określająca kolejność urządzeń.

Jeśli lista osi jest pusta, siatka ma nienazwaną oś o wielokrotności 1. W tym przypadku, jeśli nie podano listy identyfikatorów urządzeń, domyślna lista identyfikatorów urządzeń to [0]; jeśli podano listę identyfikatorów urządzeń, musi ona zawierać jedną liczbę całkowitą o dowolnej wartości nieujemnej. Nazywamy to przypadkiem maksymalnego podziału.

W przypadku wszystkich przypadków niemaksymalizowanego dzielenia, jeśli podana jest lista identyfikatorów urządzeń, iloczyn rozmiarów osi powinien odpowiadać liczbie urządzeń. Jeśli lista identyfikatorów urządzeń nie jest określona, domyślna lista identyfikatorów urządzeń to iota(product(axes)). Ze względu na prostotę nie zezwalamy też na podanie listy identyfikatorów urządzeń, która jest taka sama jak iota(product(axes)); w tym przypadku nie należy podawać listy identyfikatorów urządzeń.

Oto kilka przykładów siatek:

• Pusty układ reprezentuje układ zastępczy, który można zastąpić podczas propagacji: <[]>
• Siatka z nienazwaną osią i wyraźnym identyfikatorem urządzenia, który zwykle służy do reprezentowania maksymalnego podziału: <[], device_ids=[3]>
• Siatka z 2 osiami i ukrytymi identyfikatorami urządzeń iota(6): <["a"=2, "b"=3]>
• Siatka z 2 osiami i wyraźnymi identyfikatorami urządzeń określającymi kolejność urządzeń: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Ograniczenia:

• Elementy w elementach axes nie mogą mieć zduplikowanych nazw.
• Jeśli podano wartość device_ids:
  • Iloczyn rozmiarów osi musi być zgodny z liczbą urządzeń.
  • Wszystkie jego elementy muszą być nieujemne.
  • Wartość device_ids nie powinna być równa wartości iota(product(axis_sizes)).
  • Posortowane device_ids musi być iota(product(axis_sizes)).

Parametry:

MeshAxisAttr

Nazwana oś w siatce

Składnia:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

Parametry:

OpShardingRuleAttr

Określa sposób dzielenia operacji na partycje.

Składnia:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  bool   # is_custom_rule
>

Reguła podziału określa, jak operacja może być dzielona na podstawie różnych właściwości operacji, np. atrybutów, kształtu operandów, kształtu wyników itp. Przykład:

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>

%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

Pamiętaj, że zezwalamy na czynniki o rozmiarze 1, mimo że nie można ich dzielić. Robimy to głównie ze względu na kompletność, ponieważ wiele operacji, np. operacje punktowe, ma wymiary o rozmiarze 1, które odpowiadają sobie w przypadku operandów i wyników.

Typy czynników:

• reduction_factors zawiera indeksy czynników wymagających zmniejszenia, takich jak wymiary kontraktowe w operacji kropki.
• need_replication_factors zawiera indeksy czynników wymagających pełnej replikacji, takich jak posortowany wymiar w operacji sortowania.
• permutation_factors zawiera indeksy czynników wymagających zbiorczej permutacji, jeśli są one dzielone, np. wymiary wypełnienia w operacji wypełnienia.
• Wszystkie pozostałe czynniki są traktowane jako czynniki przekazywane, czyli czynniki, które nie wymagają żadnej komunikacji, jeśli są dzielone w ten sam sposób na wszystkich tensorach, do których są mapowane.

blocked_propagation_factors zawiera czynniki, według których nie można propagować podziału. Jest ona ortogonalna do typów czynników. Mianowicie, czynnik blokujący może należeć do dowolnego typu czynników.

is_custom_rule określa, czy jest to reguła zdefiniowana przez użytkownika. Użytkownicy mogą definiować reguły dzielenia na części w przypadku wywołań niestandardowych lub zastąpić zdefiniowane wstępnie reguły dzielenia na części w przypadku operacji standardowych. Niestandardowa reguła jest zawsze zachowywana i nigdy nie jest usuwana.

Ograniczenia:

• Liczba mapowań operandów/wyników musi być zgodna z liczbą operandów/wyników operacji.
• Istnieje co najmniej 1 mapowanie (nie można mieć reguły dla operacji bez operandów/wyników).
• Pozycja każdego elementu TensorMappingAttr odpowiada pozycji odpowiadającego mu typu tensora.
• W przypadku każdej grupy czynników (reduction_factors, need_replication_factors, permutation_factors):
  • Elementy muszą mieścić się w zakresie [0, $factor_sizes].
  • Brak zduplikowanych indeksów czynników w poszczególnych grupach i między grupami.

Parametry:

SubAxisInfoAttr

Informacje o tym, jak ta oś podrzędna jest tworzona na podstawie osi pełnej

Składnia:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

Podczas dzielenia osi pełnej na n osi podrzędnych zmienia się jej kształt na [k_1,...,k_n], a i oś podrzędna i rozmiar k_i można wyrazić jako iloczyn wszystkich rozmiarów osi po lewej stronie m=prod(k_1,...,k_(i-1)) (czyli przed rozmiarem). Dlatego atrybut sub-axis-info zawiera te 2 liczby i jest oznaczony w ten sposób: (m)k dla rozmiaru wstępnego m i rozmiaru k.

Ograniczenia:

• pre-size ma wartość co najmniej 1.
• Wartość w polu size jest większa niż 1.
• pre-size musi dzielić rozmiar pełnej osi, czyli zarówno pre-size, jak i size muszą dzielić rozmiar pełnej osi, a pod-oś nie może wykraczać poza pełną oś.
• Rozmiar podrzędnej osi nie jest równy rozmiarowi odpowiadającej pełnej osi. W takim przypadku należy użyć pełnej osi.

Parametry:

TensorMappingAttr

Mapowania czynników dla każdego wymiaru tensora.

Składnia:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

Ograniczenia:

• Elementy w elementach dim_mappings muszą spełniać ograniczenia określone w elementach DimMappingAttr.
• Brak zduplikowanych indeksów czynników w różnych wymiarach.

Parametry:

TensorShardingAttr

Dzielenie tensorów

Składnia:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

Podział tensora jest powiązany z konkretną siatką i może odwoływać się tylko do nazw osi z tej siatki. Podziały wymiarów podają, na których osiach (lub podrzędnych osiach) tensor jest dzielony od głównej do podrzędnej. Wszystkie inne osie, które nie dzielą wymiaru, są powielane (lub powielane pośrednio, jeśli występują na liście powielonych osi).

Sieć, do której jest przypisany podział, może być określona za pomocą nazwy symbolu, która odwołuje się do odpowiadającego mu symbolu MeshOp, lub za pomocą wbudowanego elementu MeshAttr.

Ograniczenia:

• Elementy w elementach dim_shardings muszą spełniać ograniczenia wymienione w elementach DimensionShardingAttr.
• Elementy w elementach replicated_axes muszą spełniać ograniczenia wymienione w elementach AxisRefListAttr.
• Jeśli odpowiadający typ tensora nie jest typu ShapedType, podział musi mieć rangę 0 i nie może zawierać powielonych osi.
• Tensor powinien mieć rangę.
• Liczba podziałów wymiaru jest równa rangowi tensora.
• Wymiary o rozmiarze 0 nie są dzielone.
• Elementy w elementach replicated_axes są uporządkowane zgodnie z elementem mesh_or_ref (patrz AxisRefAttr::getMeshComparator).

Parametry:

TensorShardingPerValueAttr

Dzielenie tensora według operandu lub wyniku operacji

Składnia:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

Lista wartości TensorShardingAttr, po jednej dla każdego operandu lub wyniku operacji.

Ograniczenia:

• Elementy w elementach shardings muszą spełniać ograniczenia elementu TensorShardingAttr.

Parametry:

Wartości w polu enum

PropagationDirection

Typ propagacji

Przypadki: