Dialekt „mpmd”

Operacje

mpmd.assign (mpmd::AssignOp)

Przypisz operację

Składnia:

operation ::= `mpmd.assign` attr-dict $tensor `:` functional-type(operands, results)

Przypisuje tensor lokalny do siatki jako w pełni zreplikowany w tej siatce.

Jest to tymczasowe działanie, które jest wprowadzane podczas obniżania liczby operacji Jax, aby przejść z typów lokalnych na typy siatkowe. Te operacje zostaną wyeliminowane podczas importowania, gdy dane wejściowe i wyniki operacji funkcji staną się tensorami siatki.

Nazwa siatki typu wyniku powinna odpowiadać siatce w topologii, a jej typ globalny powinien być identyczny z typem operandu.

Pochodzenie operacji przypisania to pochodzenie siatki, np. named_computation, mesh inference itp.

Cechy: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>

Interfejsy: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

mpmd.broadcast (mpmd::BroadcastOp)

Operacja transmisji

Składnia:

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

Umożliwia przeniesienie (lub replikację) tensora w dowolnej siatce, w której jest używany. Za każdym razem, gdy operand jest przenoszony, źródłem przeniesienia jest jego bieżąca lokalizacja.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Operandy:

Wyniki:

mpmd.call (mpmd::CallOp)

Funkcja połączenia specyficzna dla MPMD

Składnia:

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

Operacja wywołania funkcji. Przydatne np. do umieszczania treści pętli w deklaracjach funkcji, aby zmniejszyć rozmiar kodu.

Interfejsy: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

mpmd.for (mpmd::ForOp)

Dla operatora

Zwraca wynik wykonania funkcji treści dla stałej liczby iteracji, przy czym indeks iteracji jest dostępny w treści.

Można określić opcjonalny współczynnik rozwinięcia, który musi dzielić liczbę iteracji.Umożliwia on rozwinięcie treści operacji o ten współczynnik, tzn. w przypadku współczynnika rozwinięcia N treść jest replikowana w celu utworzenia N kopii, a liczba iteracji jest zmniejszana o współczynnik 1/N. Każda kopia z wyjątkiem pierwszej używa wyników poprzedniej kopii zamiast argumentów bloku, a indeks iteracji jest mnożony przez współczynnik rozwinięcia i zwiększany po każdej kopii.

Operator „for” może akceptować i zwracać dowolne typy, ale ich TypeID musi być taki sam, np. wszystkie typy tensorów lub wszystkie typy siatek MPMD itp. Dzięki temu możemy używać operatora na różnych poziomach, udostępniając implementację i przekształcenia.

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

Interfejsy: ConditionallySpeculatable, LoopLikeOpInterface, OpAsmOpInterface, ShardableDataFlowOpInterface

Atrybuty:

Operandy:

Wyniki:

mpmd.fragment (mpmd::FragmentOp)

Operacja na fragmencie

Przypisuje obliczenia, czyli blok operacji, do określonej siatki w topologii MPMD, która ma być wykonywana jako fragment indywidualnego programu SPMD.

Fragment przyjmuje i zwraca tylko tensory siatki przypisane do tej samej siatki co fragment.

Nazwa siatki fragmentu powinna odpowiadać siatce w topologii.

Fragment zawiera listę źródeł, czyli metadanych z informacjami o oryginalnych obliczeniach o nazwach, które utworzyły ten fragment, oraz identyfikator staged_id zdefiniowany iff, jeśli jest to fragment zdefiniowany przez użytkownika, czyli ma niepustą listę źródeł. Opcjonalny argument in_shardings określa podział argumentów bloku fragmentu, które odpowiadają operandowi. Opcjonalny parametr out_shardings określa podział wyników na fragmenty.

Region fragmentu nie powinien zawierać żadnych wolnych zmiennych, a typ każdego argumentu bloku i wartości zwracanej w regionie to globalny typ tensora odpowiadającego tensora siatki.

Cechy: HasParent<::mlir::func::FuncOp, ForOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfejsy: ConditionallySpeculatable, ShardableDataFlowOpInterface

Atrybuty:

Operandy:

Wyniki:

mpmd.fragment_call (mpmd::FragmentCallOp)

Operacja wywołania fragmentu

Reprezentuje wywołanie funkcji, która zawiera fragment MPMD, czyli obliczenia przypisane do konkretnej siatki w topologii MPMD, które mają być wykonywane jako pojedynczy fragment programu SPMD.

Nazwa siatki fragmentu powinna odpowiadać siatce w topologii funkcji otaczającej, a kształt siatki powinien być zgodny z kształtem wywoływanej funkcji.

Pochodzenie określa obliczenia nazwane przez użytkownika, które przyczyniły się do tego wywołania fragmentu, np. przez scalanie.

Typy danych wejściowych i wyników funkcji wywoływanej muszą być lokalnymi typami tensorów odpowiednich tensorów siatki operandów i wyników tej operacji odpowiednio.

Przykład:

%2 = mpmd.fragment_call<mesh="m1",> origin=[] @my_fragment(%0, %1) :
  (m<esh>_tensor..., m<esh>_te>nsor...) - m<esh>_tensor...

Cechy: HasParent<::mlir::func::FuncOp>, MemRefsNormalizable

Interfejsy: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Atrybuty:

Operandy:

Wyniki:

mpmd.named_computation (mpmd::NamedComputationOp)

Operacja w zakresie nazwanym

Grupuje obliczenia, czyli blok operacji, i nadaje im nazwę oraz liczbę transpozycji za pomocą atrybutu UserOrigin. To obliczenie nazwane może służyć do przypisywania siatki do obliczeń w MPMD lub do optymalizacji.

Liczba transpozycji (domyślnie 0) określa, czy nazwane obliczenie zostało uzyskane w wyniku określonej liczby przekształceń transpozycji JAX AD.

Region operacji nie powinien zawierać żadnych wolnych zmiennych, a typ każdego argumentu bloku i zwracanych wartości w regionie musi być taki sam jak typ danych wejściowych i zwracany typ operacji.

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

Interfejsy: ConditionallySpeculatable

Atrybuty:

Operandy:

Wyniki:

mpmd.named_tensor (mpmd::NamedTensorOp)

Przypisywanie tensora do siatki typu mesh

Składnia:

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

Operacja tożsamości, która przypisuje wynik tensora do danej nazwy. Ten obiekt NamedTensor może służyć do przypisywania siatki do tensora w MPMD.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

mpmd.reduce (mpmd::ReduceOp)

Operacja zmniejszania w wielu siatkach

Składnia:

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

Umożliwia zmniejszenie tensora w różnych siatkach, a następnie rozgłaszanie go w miejscach, w których ma być używany.

Cechy: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfejsy: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

mpmd.return (mpmd::ReturnOp)

Operacja mpmd.return kończy regiony dołączone do operacji opartych na regionach mpmd. Jest to funkcja o zmiennej liczbie argumentów: przyjmuje jako argumenty listę wartości, których typy mogą być dowolne (ale tego samego rodzaju, np. AnyTensor), i dlatego można jej używać na różnych poziomach stosu MPMD IR.

Składnia:

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

Cechy: AlwaysSpeculatableImplTrait, Terminator

Interfejsy: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Skutki: MemoryEffects::Effect{}

Operandy:

mpmd.transfer (mpmd::TransferOp)

Operacja przenoszenia

Składnia:

operation ::= `mpmd.transfer` attr-dict $tensor `:` functional-type(operands, results)

Przenosi rozproszony tensor z jednej siatki do drugiej.

Nazwy siatek typów operandów i wyników powinny odpowiadać siatkom w topologii, a ich typy globalne powinny być identyczne.

Cechy: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp>

Interfejsy: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ShardingRuleOpInterface

Skutki: MemoryEffects::Effect{}

Operandy:

Wyniki:

mpmd.unassign (mpmd::UnassignOp)

Operacja anulowania przypisania

Składnia:

operation ::= `mpmd.unassign` attr-dict $tensor `:` functional-type(operands, results)

Usuwa w pełni zreplikowany tensor z siatki.

Jest to tymczasowe działanie wprowadzane podczas obniżania poziomu operacji JAX, aby przejść z typów lokalnych na typy siatkowe. Te operacje zostaną wyeliminowane podczas importowania, gdy dane wejściowe i wyniki operacji funkcji staną się tensorami siatki.

Nazwa siatki typu operandu powinna odpowiadać siatce w topologii, a jej typ globalny powinien być identyczny z typem wyniku.

Cechy: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>, InferTensorType

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

Skutki: MemoryEffects::Effect{}

Atrybuty:

Operandy:

Wyniki:

Atrybuty

MeshWithOriginsAttr

Pasuje do swoich korzeni.

Składnia:

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

Parametry:

MeshesWithOriginsAttr

Lista siatek z ich źródłami.

Składnia:

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

Parametry:

NamedMeshAttr

Para z nazwą i siatką.

Składnia:

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

Parametry:

OriginAttr

Źródło przypisania siatki.

Składnia:

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

Pochodzenie przypisania do sieci mesh.

origin_label to zrozumiała dla człowieka etykieta źródła. Jest przeznaczona do debugowania.

Parametry:

ReductionAttr

Oznacza spadek.

Składnia:

#mpmd.reduction<
  `none` | `add` | `max` | `min` | `mul` | `or` | `and`   # reduction_type
>

Parametry:

TopologyAttr

Topologia nazwanych siatek.

Składnia:

#mpmd.topology<
  ::llvm::ArrayRef<NamedMeshAttr>   # meshes
>

Parametry:

UserOriginAttr

Pochodzenie obliczeń określonych przez użytkownika.

Składnia:

#mpmd.user_origin<
  ::mlir::StringAttr,   # userName
  int64_t   # transposeCount
>

Parametry:

Typy

MeshTensorType

Typ tensora siatki

Przypisuje typ RankedTensorType do określonej siatki SPMD w topologii siatek MPMD programu. Typ zawiera opcjonalne dzielenie na fragmenty, które określa, jak tensor jest dzielony na fragmenty w odniesieniu do siatki SPMD. Jeśli nie ma podziału, tensor jest w pełni replikowany.

Parametry:

Wartości w polu enum

ReductionType

Atrybut typu obniżki

Przypadki:

EdgeNodeType

Wyliczenie typu węzła brzegowego

Przypadki:

PropagationDirection

Wyliczenie kierunku propagacji

Przypadki:

ReductionOp

Wyliczenie Reduction op

Przypadki: