Dialekt 'sdy'

Dialekt Shardy (SDY)

Dialekt Shardy (SDY) definiuje reprezentację dzielenia tensora na osiach i dodatkowe komponenty interfejsu API do dołączania podziałów do tensorów.

Dziennik wersji: 0.0.1: dodaj do TensorShardingAttr osie bez redukcji.

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 przekracza wymiary tensora. Każda wewnętrzna lista określa osie, wzdłuż których należy wykonać osobne zbieranie w odpowiednim wymiarze. Zostanie on zastosowany do podziału operandu (tensor), aby uzyskać podział wyniku (out_sharding).

Pamiętaj, że pole out_sharding nie jest używane do określania podziału wyniku na fragmenty. Zamiast tego podział wyniku na fragmenty jest określany przez podział operandu na fragmenty, a wartości gathering_axes i out_sharding muszą być zgodne z tym wywnioskowanym podziałem na fragmenty.

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 Sdy_CollectiveOpInterface.
• Elementy w gathering_axes muszą spełniać ograniczenia wymienione w AxisRefListAttr.
• Zastosowanie gathering_axes do dzielenia operandu daje out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.all_reduce (sdy::AllReduceOp)

Wykonywanie komunikacji all-reduce wzdłuż osi

Składnia:

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

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

Ograniczenia:

• Musi spełniać ograniczenia wymienione w Sdy_CollectiveOpInterface.
• Parametr reduction_axes musi spełniać ograniczenia wymienione w AxisRefListAttr.
• reduction_axes musi być posortowane względem siatki.
• Podział operandu i out_sharding musi mieć równoważne podziały wymiarów.
• reduction_axes nie może pokrywać się z podziałem wymiaru operandu ani z osią replikowaną (może pokrywać się z osią niezredukowaną).
• reduction_axes nie może pokrywać się z niezredukowanymi osiami out_sharding. Innymi słowy, out_sharding musi być powielony wzdłuż reduction_axes (niejawnie lub jawnie).

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.all_slice (sdy::AllSliceOp)

Wykonuje dynamiczne wycinanie wzdłuż osi

Składnia:

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

Wyodrębnia fragmenty tensora wzdłuż osi określonych w parametrze slicing_axes. Istnieje algebraiczna dualność między sdy.all_slice a sdy.all_gather.

slicing_axes to lista list osi. Lista zewnętrzna przekracza wymiary tensora. Każda wewnętrzna lista określa osie, wzdłuż których należy wykonać wycinek w odpowiednim wymiarze. Zostanie on zastosowany do podziału operandu (tensor), aby uzyskać podział wyniku (out_sharding).

Pamiętaj, że pole out_sharding nie jest używane do określania podziału wyniku na fragmenty. Zamiast tego podział wyniku na fragmenty jest określany przez podział operandu na fragmenty, a wartości slicing_axes i out_sharding muszą być zgodne z tym wywnioskowanym podziałem na fragmenty.

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:

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

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.all_to_all (sdy::AllToAllOp)

Przeprowadza komunikację typu „każdy z każdym” wzdłuż osi

Składnia:

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

W przypadku każdej krotki (axes, src_dim, tgt_dim) na liście parametrów ta operacja wycina fragmenty tensora wzdłuż wymiaru tgt_dim i osi określonych w axes, rozprasza te fragmenty wzdłuż osi i łączy je wzdłuż wymiaru src_dim.

Ta operacja jest w zasadzie połączeniem operacji all-gather wzdłuż osi src_dim i axes, a następnie operacji all-slice wzdłuż osi tgt_dim i axes, tzn. do osi tgt_dim tensora wyjściowego dołączany jest sufiks osi src_dim wymiaru dzielenia tensora wejściowego.

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

Pamiętaj, że pole out_sharding nie jest używane do określania podziału wyniku na fragmenty. Podział wyniku na fragmenty jest określany przez podział operandów src_dim, tgt_dim i axes na fragmenty, a out_sharding musi być zgodny z tym wywnioskowanym podziałem na fragmenty.

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 Sdy_CollectiveOpInterface.
• Lista parametrów nie może być pusta.
• Dla każdego parametru w params:
  • Elementy w axes muszą spełniać ograniczenia AxisRefAttr.
  • src_dim i tgt_dim muszą być prawidłowymi wymiarami (nieujemnymi i mniejszymi niż ranga tensora).
  • Każdy znak src_dim lub tgt_dim musi być unikalny we wszystkich parametrach.
  • src_dim musi być posortowany w kolejności rosnącej we wszystkich parametrach.
• Przeniesienie axes z src_dim do tgt_dim w dzieleniu operandów daje out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.collective_permute (sdy::CollectivePermuteOp)

Wykonuje komunikację collective-permute, aby zastąpić osie

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 do innego, aby zmienić kolejność lub zastąpić osie, które dzielą tensor.

Operacja permutacji zbiorczej może przekształcić podział danych wejściowych w taki sposób, że każdy wymiar musi być podzielony w taki sam sposób jak wcześniej, tzn. musi być podzielony wzdłuż osi, których iloczyn rozmiarów jest zgodny z rozmiarem osi, które wcześniej dzieliły tensor.

Jest to przydatne do zmiany kolejności osi w ramach jednego wymiaru lub w różnych wymiarach oraz do zamiany osi podzielonych na replikowane.

W przykładzie poniżej rozmiar tensora podzielonego na fragmenty wynosi tensor<1x4x2xf32> i jest zachowywany przez operację zbiorowego przestawienia.

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 Sdy_CollectiveOpInterface.
• Jeśli dzielenie danych wejściowych i wyjściowych ma różne siatki, muszą one mieć dokładnie te same osie i różną kolejność identyfikatorów urządzeń.
• W przypadku każdego wymiaru iloczyn rozmiarów osi podziału w out_sharding musi być zgodny z podziałem odpowiedniego wymiaru operandu.

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.constant (sdy::ConstantOp)

Operacja stała

Tworzy tensor output ze stałej value.

Zobacz: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 na krawędzi przepływu danych

Składnia:

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

Krawędź przepływu danych pewnej operacji X definiuje połączenie między zbiorem źródeł (każde z nich jest operandem X lub operandem terminatora bloku X) a zbiorem miejsc docelowych (każde z nich jest wynikiem X lub argumentem bloku X), tak aby wszystkie źródła i miejsca docelowe były dzielone w ten sam sposób.

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

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 ma n krawędzi przepływu danych, a i-ta krawędź przepływu danych znajduje się między źródłami x_i, return_value_i i elementami docelowymi y_i, pred_arg_i, body_arg_i.

sdy.data_flow_edge przyjmuje jako dane wejściowe właściciela krawędzi (może to być dowolny z elementów docelowych, ale najlepiej wynik operacji, a nie argument bloku), który nie powinien mieć żadnych innych zastosowań. Ta operacja nie jest czysta, ponieważ może przyjmować dane wejściowe, które pierwotnie nie miały żadnych zastosowań.

sdy.data_flow_edge zawiera też opcjonalne dzielenie na fragmenty dla wszystkich elementów docelowych krawędzi. Podczas propagacji należy zaktualizować to dzielenie zamiast dzielenia elementów docelowych (jeśli można je dołączyć). Jest to przydatne, gdy operacja ma wiele krawędzi, ponieważ znacznie wydajniej jest:

• propagować przez każdą krawędź osobno.
• aktualizować podział każdego węzła oddzielnie, zamiast wszystkich celów naraz (np. operacja ma pojedynczy, niezmienny atrybut TensorShardingPerValueAttr dla podziału wyników);
• dodawać każdy węzeł do listy zadań osobno, gdy zmieni się podział źródła na fragmenty.

Propagacja będzie propagować podziały między wszystkimi źródłami i celami sdy.data_flow_edge, tak jakby była to zwykła operacja ze źródłami jako operandami i celami jako wynikami oraz tożsamością sdy.op_sharding_rule. Oznacza to, że propagacja w przód odbywa się od źródeł do miejsc docelowych, a propagacja wsteczna – od miejsc docelowych do źródeł.

Nie zezwalamy na określanie danych wejściowych sdy.data_flow_edge przez operację SdyDialect, więc możemy założyć, że są one określane przez operację, która ma niezarejestrowany atrybut sdy.sharding.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

Operacja krawędzi przepływu danych wejściowych/wyjściowych funkcji.

Składnia:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

Operacja krawędzi przepływu danych, ale w przypadku argumentów funkcji lub wyników wywołania. Gdy jego operandem jest BlockArgument, jest to pomost między argumentem wywołania callOp wywołującego a użytkownikami argumentu func. Każdy argument funkcji ma 1 krawędź przepływu danych funkcji. Gdy jego operandem jest OpResult, jest to pomost między wartością zwracaną wywołanej funkcji funcOp a użytkownikami wyniku wywołania. Dla każdego wyniku wywołania istnieje 1 krawędź przepływu danych funkcji.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, SymbolUserOpInterface

Operandy:

Wyniki:

sdy.manual_computation (sdy::ManualComputationOp)

Operacja równoległa na wielu urządzeniach z ręcznie tworzonymi kolekcjami

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 zapisanego w postaci kodu lokalnego na urządzenie z wyraźnymi zbiorami, w którym kształty logiczne odpowiadają kształtom lokalnego bufora fizycznego na urządzenie, a zbiory dokładnie odpowiadają fizycznej komunikacji na innym urządzeniu.

Ciało jest lokalne w odniesieniu do osi manual_axes. Rozprzestrzenianie nastąpi w całym ciele wzdłuż dowolnych osi swobodnych, czyli tych, które nie znajdują się na liście manual_axes.

Pamiętaj, że wszystkie tensory bez rangi powinny mieć podział o randze 0, czyli być w pełni replikowane.

Ograniczenia:

• Elementy w in_shardings i out_shardings muszą spełniać ograniczenia wymienione w TensorShardingAttr.
• Liczba globalnych i lokalnych tensorów wejściowych/wyjściowych regionu operacji musi być zgodna.
• W każdym podziale wymiaru osie ręczne muszą występować przed osiami swobodnymi.
• Osie ręczne nie mogą wprowadzać dopełnienia. Wymiar musi być podzielny przez odpowiedni rozmiar osi ręcznej.
• Globalne i lokalne kształty argumentów/wyników regionów operacji muszą być zgodne.

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

Interfejsy: ShardableDataFlowOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.mesh (sdy::MeshOp)

Nazwana sieć typu mesh

Składnia:

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

Definiuje nową sieć typu mesh. Wszystkie siatki w module muszą mieć tę samą liczbę urządzeń (z wyjątkiem siatek z jednym identyfikatorem urządzenia). Siatka to operacja Symbol, która pojawia się w SymbolTable modułu i do której można się odwoływać za pomocą name.

Cechy: HasParent<ModuleOp>, SymbolName

Interfejsy: Symbol

Atrybuty:

sdy.named_computation (sdy::NamedComputationOp)

Nazwana operacja obliczeniowa

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ę. Propagacja będzie przebiegać w regionie i poza nim tak, jakby wszystko było wbudowane.

Można go używać do przekazywania instrukcji wywoływania do innych funkcji. Użytkownicy Shardy powinni napisać przepustkę importu/eksportu, która przekształca operacje wywołań w operacje sdy.named_computation, duplikując lub kopiując treść wywoływanej funkcji do treści named_computation.

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

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, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Działanie bariery rozprzestrzeniania

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 dane wejściowe. Jeśli chodzi o propagację, będzie ona możliwa tylko w określonym kierunku.

Zapobiega to propagowaniu podziałów między użyciami wyniku operacji bariery a jej operandem.

• Symbol FORWARD oznacza, że podział może przepływać tylko od operandu do wyniku.
• BACKWARD oznacza, że podział może przepływać tylko z wyniku do operandu.
• NONE oznacza, że w przypadku tej operacji nie można propagować podziału.
• Nie można określić BOTH, ponieważ ta operacja byłaby zbędna.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

sdy.reduce_scatter (sdy::ReduceScatterOp)

Wykonuje komunikację reduce-scatter wzdłuż osi

Składnia:

operation ::= `sdy.reduce_scatter` ($reduction_op^)? $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Redukuje fragmenty tensora wzdłuż osi określonych w reduce_scatter_axes, a następnie rozprasza wynik wzdłuż tych samych osi. Ta operacja jest w zasadzie połączeniem sdy.all_reduce i sdy.all_slice wzdłuż tej samej reduce_scatter_axes.

Ograniczenia:

• Musi spełniać ograniczenia wymienione w Sdy_CollectiveOpInterface.
• Elementy w reduce_scatter_axes muszą spełniać ograniczenia wymienione w AxisRefListAttr.
• Zastosowanie reduce_scatter_axes do dzielenia operandów daje out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Przenieś osie replikowane w sposób domyślny lub jawny do osi niezredukowanych.

Składnia:

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

W operandzie należy niejawnie lub jawnie powtórzyć wartość axes. Ta operacja sprawia, że w wyniku nie są one zmniejszane. Mamy te relacje:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

Przykład:

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

Ograniczenia:

• Musi spełniać ograniczenia wymienione w Sdy_CollectiveOpInterface.
• Parametr axes musi spełniać ograniczenia wymienione w AxisRefListAttr.
• axes musi być posortowane względem siatki.
• axes nie są puste.
• Dane wejściowe i wyjściowe muszą mieć te same podziały wymiarów.
• axes musi być replikowany w podziale operandów na fragmenty w sposób pośredni lub bezpośredni.
• inUnreducedAxes + axes = outUnreducedAxes.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.reshard (sdy::ReshardOp)

Ponowne dzielenie tensora na inne części

Składnia:

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

Zmienia podział tensora wejściowego na fragmenty zgodnie z określonym podziałem, który różni się od dotychczasowego podziału tensora wejściowego.

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

1. Przed propagacją podziału użytkownicy dodają ShardingConstraintOp.
2. Propagacja fragmentacji wykorzystuje ShardingConstraintOp. W wynikach propagacji podziału nie ma elementu ShardingConstraintOp. Zamiast tego w razie potrzeby można dodać ReshardOp.
3. Partycjoner przekształca operację ReshardOp w operację zbiorową (lub operację tożsamości). W wynikach partycjonowania nie powinno być żadnych operacji ReshardOp.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), SymbolUserOpInterface

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

sdy.return (sdy::ReturnOp)

Operacja sdy.return kończy działanie regionów dołączonych do sdy operacji opartych na regionach i wszelkich innych operacji opartych na regionach Shardy. Jest to funkcja wieloargumentowa: przyjmuje jako argumenty listę wartości, których typy mogą być dowolne (ale tego samego rodzaju, np. AnyTensor), dzięki czemu można jej używać na różnych poziomach stosu Shardy IR.

Składnia:

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

Cechy: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

Interfejsy: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Skutki: MemoryEffects::Effect{}

Operandy:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

Przenieś niektóre osie operandu z podziałem na fragmenty do osi wyniku bez redukcji.

Składnia:

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

Do podziału operandu należy użyć znaku axes. Ta operacja powoduje, że w wyniku nie są one zredukowane. Łączą nas te relacje:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), gdzie funkcje all-gather, sharded-to-unreduced i all-reduce są stosowane na tych samych osiach.

Przykład:

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

Ograniczenia:

• Musi spełniać ograniczenia wymienione w Sdy_CollectiveOpInterface.
• Elementy w axes muszą spełniać ograniczenia wymienione w AxisRefListAttr.
• Zastosowanie axes do dzielenia operandu daje out_sharding.

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Ogranicza tensor do określonego podziału

Składnia:

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

Dołącza podział do tensora pośredniego (np. wyniku mnożenia macierzy), aby wskazać, w jaki sposób ten tensor lub podzbiór jego zastosowań powinien być podzielony.

Jeśli podział ma otwarte wymiary i nieograniczone osie, oznacza to, że tensor można dalej dzielić wzdłuż otwartych wymiarów.

Ten operator może:

• Nie mają zastosowań (wiszące) – co oznacza, że dołączone dzielenie jest sposobem, w jaki tensor wejściowy powinien być dzielony.
• ma zastosowania – co oznacza, że dołączone dzielenie jest sposobem dzielenia zastosowań operacji ograniczenia dzielenia, podczas gdy inne zastosowania tensora wejściowego mogą mieć inne dzielenie (jeśli tensor wejściowy nie ma innych zastosowań, zachowanie jest takie samo jak w przypadku braku zastosowań);

Cechy: SameOperandsAndResultType

Interfejsy: InferTypeOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.sharding_group (sdy::ShardingGroupOp)

Ogranicza tensory w grupie do tego samego podziału.

Składnia:

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

Ta operacja 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 na fragmenty w dokładnie ten sam sposób. Ta operacja przyjmuje identyfikator grupy argumentów i nie zwraca żadnego wyniku, ale modyfikuje wewnętrzną reprezentację grupy dzielenia na fragmenty, aby dodać tensor wejściowy do grupy o podanym identyfikatorze.

Interfejsy: InferTypeOpInterface

Atrybuty:

Operandy:

Atrybuty

AllToAllParamAttr

Parametr all-to-all

Składnia:

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

Krotka zawierająca osie i wymiary źródłowe/docelowe, na których ma być przeprowadzona operacja all-to-all.

Parametry:

AllToAllParamListAttr

Lista parametrów all-to-all

Składnia:

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

Parametry:

AxisRefAttr

Odniesienie do pełnej osi lub podzielonej osi pomocniczej

Składnia:

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

Ograniczenia:

• W powiązanym MeshAttr musi występować name.
• Jeśli występuje właściwość sub_axis_info, musi spełniać ograniczenia właściwości SubAxisInfoAttr.

Parametry:

AxisRefListAttr

Lista odniesień do osi

Składnia:

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

Ograniczenia:

• Elementy w value muszą spełniać ograniczenia AxisRefAttr.
• Nie ma zduplikowanych odwołań do osi ani podosi, które się ze sobą pokrywają.
• Żadne 2 sąsiadujące ze sobą odwołania do osi nie są kolejnymi podosiami tej samej pełnej osi, tzn. można je scalić w jedną podoś lub pełną oś.

Parametry:

AxisToPropagationDetailsAttr

Szczegóły przepływu krawędzi propagacji dla konkretnej osi i źródła.

Składnia:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

Mapuje odwołanie do wartości źródłowej na listę odwołań do wartości docelowych wzdłuż określonej osi.

Parametry:

DimMappingAttr

Lista indeksów czynników dla wymiaru

Pusta lista oznacza, że jest to mapowanie zerowe (jest ono analizowane/drukowane za pomocą znaku *), tzn. wymiar nie jest mapowany na żadne czynniki.

Ograniczenia:

• Istnieje co najmniej 1 indeks czynnika.
• Indeksy czynników muszą mieścić się w zakresie [0, $factor_sizes).
• Jeśli jest wiele czynników, żaden z nich nie może mieć rozmiaru 1.
• Nie może zawierać zduplikowanych indeksów czynników.

Parametry:

DimensionShardingAttr

Dzielenie wymiarów

Lista nazw osi, na których ma być dzielony wymiar tensora od najważniejszego do najmniej ważnego, wartość logiczna wskazująca, czy wymiar można dalej dzielić, oraz opcjonalna liczba całkowita oznaczająca priorytet tego dzielenia wymiaru, który będzie uwzględniany podczas propagacji dzielenia. Priorytety pochodzą z adnotacji dotyczących podziału użytkowników na grupy, a niższa wartość oznacza wyższy priorytet. Jeśli w adnotacji brakuje priorytetu, przyjmuje się najwyższy priorytet.

Ograniczenia:

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

Parametry:

EdgeValueRefAttr

Odwołanie do konkretnego indeksu krawędzi wartości typu type.

Składnia:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

Parametry:

ListOfAxisRefListsAttr

Lista list referencyjnych osi

Składnia:

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

Parametry:

ManualAxesAttr

Lista osi, w przypadku których operacja ManualComputationOp jest wykonywana ręcznie

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

• Jeśli nie podano właściwości device_ids, jest to pusta siatka.
• Jeśli podasz device_ids, musi to być pojedyncza nieujemna liczba całkowita. Nazywamy ją maksymalną siatką podziału.

Jeśli podano listę osi

• Jeśli podano listę identyfikatorów urządzeń, iloczyn rozmiarów osi powinien być zgodny z liczbą urządzeń.
• Jeśli lista identyfikatorów urządzeń nie jest określona, domyślna lista identyfikatorów urządzeń to iota(product(axes)). Dla uproszczenia zabraniamy też określania listy identyfikatorów urządzeń, która jest taka sama jak iota(product(axes)); w tym przypadku nie należy określać listy identyfikatorów urządzeń.
• Nie jest to siatka z maksymalnym podziałem, nawet jeśli łączny rozmiar osi wynosi 1.

Oto kilka przykładów siatek:

• Pusta siatka reprezentuje siatkę zastępczą, którą można zastąpić podczas propagacji: <[]>
• Siatka bez listy osi i pojedynczy nieujemny identyfikator urządzenia, który jest siatką maksymalnego podziału: <[], device_ids=[3]>
• Siatka z 2 osiami i niejawnymi 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 device_ids powinny być nieujemne.
• Jeśli axes jest puste, rozmiar device_ids może wynosić 0 (pusta siatka) lub 1 (siatka z maksymalnym podziałem).
• Jeśli axes nie jest puste,
  • Elementy w axes nie mogą mieć zduplikowanych nazw.
  • Jeśli podano wartość device_ids, oryginalna wartość device_ids nie jest równa iota(product(axis_sizes)), a posortowana wartość device_ids jest równa iota(product(axis_sizes)).

Parametry:

MeshAxisAttr

Nazwana oś w sieci typu mesh

Składnia:

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

Parametry:

OpShardingRuleAttr

Określa, jak można podzielić operację 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 operację można podzielić na partycje według różnych właściwości operacji – dowolnych 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>

Zezwalamy na czynniki o rozmiarze 1, mimo że nie można ich dzielić na fragmenty. Jest to głównie dla zachowania kompletności, ponieważ wiele operacji, takich jak operacje punktowe, ma wymiary o rozmiarze 1, które odpowiadają operandowi i wynikom.

Typy czynników:

• reduction_factors zawiera indeksy czynników wymagających redukcji, np. wymiarów kurczących się w operacji iloczynu skalarnego. Czynniki te mogą występować w operandach, ale nie w wynikach.
• need_replication_factors zawiera indeksy czynników wymagających pełnej replikacji, np. posortowanego wymiaru w operacji sortowania.
• permutation_factors zawiera indeksy czynników wymagających zbiorczego przestawienia, jeśli są podzielone na fragmenty, np. wymiary dopełnienia w operacji dopełniania.
• 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 we wszystkich tensorach, do których są mapowane.

blocked_propagation_factors zawiera czynniki, w przypadku których nie można propagować podziałów. Jest on ortogonalny w stosunku do typów czynników. Czyli czynnik blokujący propagację może być dowolnego typu.

is_custom_rule określa, czy jest to reguła zdefiniowana przez użytkownika. Użytkownicy mogą definiować reguły dzielenia na fragmenty dla wywołań niestandardowych lub zastępować wstępnie zdefiniowane reguły dzielenia na fragmenty w przypadku operacji standardowych. Reguła niestandardowa 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 ani wyników).
• Ranga każdego TensorMappingAttr odpowiada randze odpowiedniego 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].
  • W każdej grupie i między grupami nie ma zduplikowanych indeksów czynników.

Parametry:

PropagationEdgesAttr

Metadane krawędzi propagacji dla wszystkich etapów propagacji.

Składnia:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

Lista szczegółów propagacji dla poszczególnych osi dla danej wartości, pogrupowana według indeksu kroku.

Parametry:

PropagationOneStepAttr

Metadane propagacji na poszczególnych etapach.

Składnia:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

Szczegóły propagacji dla wszystkich osi w przypadku pojedynczego kroku propagacji.

Parametry:

SubAxisInfoAttr

Informacje o tym, jak ten podzakres jest wyodrębniany z pełnego zakresu

Składnia:

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

Podczas dzielenia pełnej osi na n podosi oś jest przekształcana w [k_1,...,k_n], a i-tą podoś można wyrazić jako iloczyn wszystkich rozmiarów osi po lewej stronie m=prod(k_1,...,k_(i-1)) (czyli rozmiar wstępny) i rozmiar k_i. Dlatego atrybut sub-axis-info zawiera te 2 liczby i jest oznaczony jako (m)k dla rozmiaru wstępnego m i rozmiaru k.

Ograniczenia:

• pre-size wynosi co najmniej 1.
• size jest większe niż 1.
• pre-size musi dzielić rozmiar pełnej osi, tzn. zarówno pre-size, jak i size dzielą rozmiar pełnej osi, a oś podrzędna nie wykracza poza pełną oś.
• Rozmiar osi pomocniczej nie jest równy rozmiarowi odpowiedniej osi głównej. W takim przypadku należy użyć osi głównej.

Parametry:

TensorMappingAttr

Mapowania czynników dla każdego wymiaru tensora.

Składnia:

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

Ograniczenia:

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

Parametry:

TensorShardingAttr

Dzielenie tensora

Składnia:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

Podział tensora jest powiązany z konkretną siatką i może odwoływać się tylko do nazw osi z tej siatki. Podział wymiarów informuje nas, w przypadku każdego wymiaru tensora, wzdłuż których osi (lub podosi) jest on dzielony od głównej do podrzędnej. Wszystkie pozostałe osie, które nie dzielą wymiaru, są replikowane w sposób dorozumiany lub jawny (jeśli pojawiają się na liście replikowanych osi).

Pamiętaj, że brak atrybutu podziału tensora jest równoznaczny z pełnym podziałem tensora.

Siatkę, do której jest przypisane to dzielenie, można określić za pomocą nazwy symbolu, która odwołuje się do odpowiedniego symbolu MeshOp, lub wstawionego symbolu MeshAttr.

Partycjonowanie może mieć osie niepodlegające redukcji (oznaczone symbolem unreduced_axes), co oznacza, że tensor nie podlega redukcji wzdłuż tych osi. Jeśli na przykład wymiar kontraktowy operacji mnożenia macierzy jest podzielony wzdłuż osi x zarówno w przypadku lewej, jak i prawej strony, wynik nie jest zredukowany wzdłuż osi x. Zastosowanie operacji all-reduce na tensorze wzdłuż osi, które nie zostały zredukowane, spowoduje replikację tensora wzdłuż tych osi. Tensor z niezredukowanymi osiami nie musi jednak być od razu w pełni zredukowany. Może pozostać niezredukowany, gdy jest przekazywany do operacji liniowych, takich jak stablehlo.add (o ile zarówno lewa, jak i prawa strona są niezredukowane), a następnie może być w pełni zredukowany. Zakładamy, że typem obniżki jest suma. W przyszłości możemy obsługiwać inne obniżki.

Ograniczenia:

• Elementy w dim_shardings muszą spełniać ograniczenia wymienione w DimensionShardingAttr.
• Elementy w replicated_axes muszą spełniać ograniczenia wymienione w AxisRefListAttr.
• Elementy w unreduced_axes muszą spełniać ograniczenia wymienione w AxisRefListAttr.
• Jeśli odpowiedni typ tensora nie jest ShapedType, podział musi mieć rangę 0 i nie może zawierać replikowanych osi.
• Jeśli jest to ShapedType, wykonaj te czynności:
  • Tensor powinien mieć rangę.
  • Liczba podziałów wymiarów jest równa randze tensora.
  • Wymiary o rozmiarze 0 nie są dzielone.
• Nie ma zduplikowanych odniesień do osi ani podosi, które nakładają się na siebie w przypadku dim_shardings, replicated_axes i unreduced_axes.
• Elementy w replicated_axes i unreduced_axes są uporządkowane względem mesh_or_ref (patrz AxisRefAttr::getMeshComparator).

Parametry:

TensorShardingPerValueAttr

Dzielenie tensora na fragmenty według operandu lub wyniku operacji

Składnia:

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

Lista TensorShardingAttr, po jednym dla każdego operandu lub wyniku operacji.

Ograniczenia:

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

Parametry:

Wartości w polu enum

EdgeNodeType

Wyliczenie typu węzła brzegowego

Przypadki:

PropagationDirection

Wyliczenie kierunku propagacji

Przypadki:

ReductionOp

Wyliczenie Reduction op

Przypadki: