Dialecte 'mpmd'

Opérations

mpmd.assign (mpmd::AssignOp)

Opération d'attribution

Syntaxe :

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

Attribue un Tensor local à un maillage en tant que répliqué dans ce maillage.

Il s'agit d'une opération temporaire introduite lors de la réduction des opérations JAX, pour passer des types locaux aux types de maillage. Ces opérations seront éliminées lors de l'importation, lorsque les entrées et les résultats de l'opération de fonction deviendront des tenseurs maillés.

Le nom de maillage du type de résultat doit correspondre à un maillage dans la topologie, et son type global doit être identique au type d'opérande.

L'origine de l'opération d'attribution est l'origine du maillage, par exemple named_computation, inférence de maillage, etc.

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

Interfaces : ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effets : MemoryEffects::Effect{}

Attributs :

Opérandes :

Résultats :

mpmd.broadcast (mpmd::BroadcastOp)

Opération de diffusion

Syntaxe :

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

Permet de transférer (ou de répliquer) un Tensor dans n'importe quel maillage où il est utilisé. Chaque fois qu'un transfert est effectué, l'origine du transfert correspond à l'emplacement actuel de l'opérande.

Traits : AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfaces : ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effets : MemoryEffects::Effect{}

Opérandes :

Résultats :

mpmd.call (mpmd::CallOp)

Fonction d'appel spécifique à MPMD

Syntaxe :

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

Opération d'appel de fonction. Utile pour encapsuler le corps des boucles dans des déclarations de fonctions afin de réduire la taille du code, par exemple.

Interfaces : ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Attributs :

Opérandes :

Résultats :

mpmd.for (mpmd::ForOp)

Pour l'opérateur

Renvoie le résultat de l'exécution d'une fonction de corps pour un nombre fixe d'itérations, avec l'index d'itération disponible dans le corps.

Un facteur de déploiement facultatif, qui doit diviser le nombre d'itérations, peut être spécifié pour déployer le corps de l'opération par ce facteur, c'est-à-dire que pour un facteur de déploiement N, le corps est répliqué pour créer N copies et le nombre d'itérations est réduit d'un facteur 1/N. Chaque copie, à l'exception de la première, utilise les résultats de la copie précédente au lieu des arguments de bloc. L'index d'itération est multiplié par le facteur de déploiement et incrémenté après chaque copie.

Un opérateur for peut accepter et renvoyer n'importe quel type, mais le TypeID de ceux-ci doit être le même (par exemple, tous les types de Tensor ou tous les types de maillage MPMD, etc.). Cela nous permet d'utiliser l'opérateur à différents niveaux, en partageant l'implémentation et les transformations.

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

Interfaces : ConditionallySpeculatable, LoopLikeOpInterface, OpAsmOpInterface, ShardableDataFlowOpInterface

Attributs :

Opérandes :

Résultats :

mpmd.fragment (mpmd::FragmentOp)

Opération de fragment

Attribue un calcul, c'est-à-dire un bloc d'opérations, à un maillage spécifique dans une topologie MPMD, qui est destiné à être exécuté en tant que fragment de programme SPMD individuel.

Le fragment ne prend et ne renvoie que les Tensors de maillage qui sont attribués au même maillage que le fragment.

Le nom du maillage du fragment doit correspondre à un maillage de la topologie.

Le fragment inclut une liste d'origines, c'est-à-dire des métadonnées contenant des informations sur les named_computations d'origine qui ont formé ce fragment, et un staged_id défini iff s'il s'agit d'un fragment défini par l'utilisateur, c'est-à-dire s'il comporte une liste d'origines non vide. Le paramètre facultatif in_shardings spécifie le partitionnement des arguments de bloc d'un fragment, qui correspondent aux opérandes. Le paramètre facultatif out_shardings spécifie le partitionnement des résultats.

La région du fragment ne doit pas comporter de variables libres, et le type de chaque argument de bloc et des valeurs renvoyées dans la région correspond au type de Tensor global du Tensor de maillage correspondant.

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

Interfaces : ConditionallySpeculatable, ShardableDataFlowOpInterface

Attributs :

Opérandes :

Résultats :

mpmd.fragment_call (mpmd::FragmentCallOp)

Opération d'appel de fragment

Représente un appel à une fonction qui contient un corps de fragment MPMD, c'est-à-dire un calcul attribué à un maillage spécifique dans une topologie MPMD, qui est destiné à être exécuté en tant que fragment de programme SPMD individuel.

Le nom du maillage du fragment doit correspondre à un maillage dans la topologie de la fonction englobante, et la forme de ce maillage doit correspondre à celle de l'appelé.

L'origine spécifie les calculs nommés par l'utilisateur qui ont contribué à cet appel de fragment, par exemple par fusion.

Les types d'entrée et de résultat de la fonction de l'appelé doivent être les types de tenseur locaux des tenseurs de maillage correspondants des opérandes et des résultats de cette opération, respectivement.

Exemple :

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

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

Interfaces : ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Attributs :

Opérandes :

Résultats :

mpmd.named_computation (mpmd::NamedComputationOp)

Opération de portée nommée

Regroupe un calcul (c'est-à-dire un bloc d'opérations), et lui attribue un nom et un nombre de transpositions via l'attribut UserOrigin. Cette NamedComputation peut être utilisée pour attribuer un maillage au calcul dans MPMD ou pour les optimisations.

Le nombre de transpositions (par défaut=0) indique si le calcul nommé a été produit par un certain nombre de transformations de transposition JAX AD.

La région de l'opération ne doit pas comporter de variables libres, et le type de chaque argument de bloc et des valeurs renvoyées dans la région doit être le même que le type des entrées et le type renvoyé de l'opération.

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

Interfaces : ConditionallySpeculatable

Attributs :

Opérandes :

Résultats :

mpmd.named_tensor (mpmd::NamedTensorOp)

Attribuer un Tensor à un maillage

Syntaxe :

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

Op d'identité qui associe le résultat du Tensor à un nom donné. Ce NamedTensor peut être utilisé pour attribuer un maillage au Tensor dans MPMD.

Traits : AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfaces : ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effets : MemoryEffects::Effect{}

Attributs :

Opérandes :

Résultats :

mpmd.reduce (mpmd::ReduceOp)

Opération de réduction cross-mesh

Syntaxe :

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

Permet de réduire un Tensor sur différents maillages, puis de le diffuser là où il doit être utilisé.

Traits : AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfaces : ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effets : MemoryEffects::Effect{}

Attributs :

Opérandes :

Résultats :

mpmd.return (mpmd::ReturnOp)

L'opération mpmd.return met fin aux régions associées aux opérations basées sur les régions mpmd. Il est variadique : il prend comme arguments une liste de valeurs dont les types peuvent être quelconques (mais du même genre, par exemple AnyTensor) et peut donc être réutilisé à différents niveaux de la pile IR MPMD.

Syntaxe :

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

Traits : AlwaysSpeculatableImplTrait, Terminator

Interfaces : ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effets : MemoryEffects::Effect{}

Opérandes :

mpmd.transfer (mpmd::TransferOp)

Opération de transfert

Syntaxe :

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

Transfère un Tensor distribué d'un maillage à un autre.

Les noms de maillage des types d'opérande et de résultat doivent correspondre aux maillages de la topologie, et leurs types globaux doivent être identiques.

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

Interfaces : ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ShardingRuleOpInterface

Effets : MemoryEffects::Effect{}

Opérandes :

Résultats :

mpmd.unassign (mpmd::UnassignOp)

Opération de désattribution

Syntaxe :

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

Désattribue un tenseur entièrement répliqué d'un réseau maillé.

Il s'agit d'une opération temporaire introduite lors de la réduction des opérations JAX, pour passer des types locaux aux types de maillage. Ces opérations seront éliminées lors de l'importation, lorsque les entrées et les résultats de l'opération de fonction deviendront des tenseurs maillés.

Le nom de maillage du type d'opérande doit correspondre à un maillage dans la topologie, et son type global doit être identique au type de résultat.

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

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

Effets : MemoryEffects::Effect{}

Attributs :

Opérandes :

Résultats :

Attributs

MeshWithOriginsAttr

Mesh avec ses origines.

Syntaxe :

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

Paramètres :

MeshesWithOriginsAttr

Liste des maillages et de leurs origines.

Syntaxe :

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

Paramètres :

NamedMeshAttr

Une paire avec un nom et un maillage.

Syntaxe :

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

Paramètres :

OriginAttr

Origine de l'attribution du maillage.

Syntaxe :

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

Origine d'une attribution de maillage.

origin_label est un libellé lisible pour l'origine. Il est destiné à être utilisé à des fins de débogage.

Paramètres :

ReductionAttr

Indique une réduction.

Syntaxe :

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

Paramètres :

TopologyAttr

Topologie des maillages nommés.

Syntaxe :

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

Paramètres :

UserOriginAttr

Origine du calcul spécifié par l'utilisateur.

Syntaxe :

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

Paramètres :

Types

MeshTensorType

Type de tenseur de maillage

Attribue un RankedTensorType à un maillage SPMD spécifique dans la topologie MPMD des maillages du programme. Le type contient un sharding facultatif qui spécifie comment le Tensor est shardé par rapport au maillage SPMD. Si le partitionnement n'est pas présent, le Tensor est entièrement répliqué.

Paramètres :

Enums

ReductionType

Attribut du type de réduction

Étuis :

EdgeNodeType

Énumération du type de nœud Edge

Étuis :

PropagationDirection

Énumération de la direction de propagation

Étuis :

ReductionOp

Énumération de l'opération de réduction

Étuis :