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.constant (sdy::ConstantOp)

Działanie ciągłe

Generuje tensor output ze stałej wartości 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)

Funkcja brzegu przepływu danych

Składnia:

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

Brzeg przepływu danych niektórych działań X określa most między zbiorem źródeł (każdy jest operandem X lub operandem zakończenia bloku X) a zbiorem celów (każdy jest wynikiem wyniku X lub argumentu blokowego X). W ten sposób wszystkie źródła i cele powinny być podzielone na fragmenty w ten sam sposób.

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 wierzchołek korzenia krawędzi (może to być dowolny cel, ale najlepiej wynik funkcji op zamiast argumentu block), który nie powinien mieć żadnych innych zastosowań. 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 również opcjonalne fragmentację dla wszystkich środowisk docelowych brzegu, która powinna być aktualizowana zamiast fragmentacji celów (jeśli można dołączyć) podczas propagacji. Jest to przydatne, gdy opcja ma wiele krawędzi, ponieważ znacznie wydajniej jest:

• rozprzestrzeniać się po każdej krawędzi osobno.
• zaktualizuj fragmentację każdej krawędzi oddzielnie zamiast wszystkich celów jednocześnie (np. operacja ma 1 stały element TensorShardingPerValueAttr na potrzeby fragmentacji wyników).
• dodawać poszczególne krawędzie do listy zadań osobno, gdy zmieni się podział na fragmenty ź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 fizycznych buforów na urządzeniu, a zbiory dokładnie odpowiadają fizycznej komunikacji między urządzeniami.

Ciało jest lokalne względem manual_axes. Rozchodzenie się będzie występować w całej objętości w przypadku wszystkich wolnych osi – tych, które nie znajdują się na liście manual_axes.

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

Atrybuty:

Argumenty:

Wyniki:

sdy.mesh (sdy::MeshOp)

Nazwa siatki

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ę. Rozprzestrzenianie się informacji do i z regionu będzie wyglądać tak, jakby wszystko było wstawione.

Można go użyć do propagowania instrukcji wywołania do innych funkcji. Wszyscy użytkownicy Shardy powinni napisać funkcję importowania/eksportowania, która konwertuje ich wywołania funkcji na operacje sdy.named_computation, duplikując lub kopiując ciało wywoływanej funkcji do ciała funkcji 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, 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 dane wejściowe. Jednak w przypadku propagacji będzie to oznaczać, że propagacja będzie mogła odbywać się tylko w określonym kierunku.

Zapobiega to rozpowszechnianiu fragmentacji między zastosowaniami wyniku operacji bariery i jej operandu.

• FORWARD oznacza, że elementy mogą być dzielone tylko od operandu do wyniku.
• BACKWARD oznacza, że elementy mogą być dzielone tylko z wyniku na operand.
• NONE oznacza, że fragmentowanie nie może być rozpowszechniane przez tę operację.
• Nie można określić BOTH, ponieważ ta operacja byłaby nadmiarowa.

Cechy: AlwaysSpeculatableImplTrait, Elementwise, 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 tensor wejściowy za pomocą określonego podziału, który różni się od dotychczasowego podziału tensora wejściowego.

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

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. partycjoner przekształca ReshardOp w zbiorową operację tożsamości (lub operację tożsamości). W wynikach partycjonowania nie powinno być operacji ReshardOp.

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

Cechy: AlwaysSpeculatableImplTrait, Elementwise, 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 i wszystkich innych operacji opartych na regionie Shardy. Jest zmiennoargumentowa: 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 modułu 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 fragmentu fragmentacji

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ą żadnych zastosowań (wiszące) – oznacza to, że dołączone shardingi to sposób, w jaki tensor wejściowy powinien być podzielony na części;
• Have uses – oznacza, że dołączone podziały to sposób, w jaki należy dzielić elementy operacji 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: Elementwise, SameOperandsAndResultType

Interfejsy: InferTypeOpInterface

Atrybuty:

Operandy:

Wyniki:

sdy.sharding_group (sdy::ShardingGroupOp)

Operacja grupowania fragmentów

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). Gdy podczas propagacji tylko jeden element grupy zostanie podzielony na fragmenty, wszystkie pozostałe elementy zostaną podzielone na fragmenty w taki sam sposób. Ta operacja przyjmuje argument group_id 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.

Atrybuty:

Operandy:

Atrybuty

AxisRefAttr

Odniesienie do pełnej osi lub podzielonej podosi

Składnia:

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

Parametry:

DimMappingAttr

Lista indeksów czynników wymiaru

Wszystkie indeksy czynników muszą mieścić się w zakresie [0, num_factors), a pusta lista wskazuje, że jest to mapowanie puste (jest ono analizowane i drukowane za pomocą funkcji *), tzn. wymiar nie jest zmapowany na żadne czynniki.

Parametry:

DimensionShardingAttr

Fragmentacja wymiarów

Lista nazw Os 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.

Parametry:

ManualAxesAttr

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 zawiera niejawną nienazwaną oś rozmiaru 1. Jeśli w tym przypadku nie podano listy identyfikatorów urządzeń, niejawna lista identyfikatorów urządzeń ma postać [0]. Jeśli lista identyfikatorów urządzeń jest podana, musi ona zawierać jedną liczbę całkowitą dowolną nieujemną wartość. Nazywamy to maksymalnym dzieleniem.

W przypadku wszystkich przypadków niemaksymalizowanego dzielenia na części, 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 sieci typu mesh:

• Pusta siatka to zastępcza siatka, którą można zastąpić podczas propagacji: <[]>
• Siatka z osią bez nazwy i wyraźnym identyfikatorem urządzenia, który zwykle reprezentuje maksymalne wykorzystanie do fragmentowania: <[], 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]>

Parametry:

MeshAxisAttr

Nazwana oś w siatce

Składnia:

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

Parametry:

OpShardingRuleAttr

Określa sposób partycjonowania operacji.

Składnia:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  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.

is_custom_rule określa, czy jest to reguła zdefiniowana przez użytkownika dla operacji stablehlo.custom_call. Usługa partycjonowania nie wie, jak partycjonować te operacje, więc użytkownik musi podać sposób. Jeśli jest to reguła niestandardowa, jest ona zawsze zachowywana i nigdy nie jest usuwana. Wartość is_custom_rule może być ustawiona na Prawda tylko w przypadku operacji stablehlo.custom_call.

Parametry:

SubAxisInfoAttr

Informacje o tym, jak ta oś podrzędna jest oddalona od pełnej osi

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.

Parametry:

TensorMappingAttr

Mapowania czynników dla każdego wymiaru tensora.

Składnia:

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

Parametry:

TensorShardingAttr

dzielenie tensorów,

Składnia:

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

Fragmentacja tensora jest powiązana z konkretną siatką i może odwoływać się tylko do nazw osi z tej siatki. Fragmenty wymiarów informują nas o każdym wymiarze tensora, według których osie (lub osi podrzędnych) jest podzielony na mniejsze i dużej osi. Wszystkie inne osie, które nie dzielą wymiaru, są powielane (lub powielane w sposób domyślny lub jawny, 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 symbolu MeshOp, lub za pomocą wbudowanego elementu MeshAttr.

Parametry:

TensorShardingPerValueAttr

Dzielenie tensora na operandy/wyniki operacji

Składnia:

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

Parametry:

Wartości w polu enum

PropagationDirection

kierunek propagacji – typ wyliczeniowy

Przypadki: