Dialecte

Le dialecte Shardy (SDY) définit une représentation de fractionnement de tenseur basée sur les axes et des composants d'API supplémentaires pour associer des fractionnements à des tenseurs.

Opérations

sdy.constant (sdy::ConstantOp)

Opération constante

Génère un tenseur output à partir d'une constante value.

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

Exemple :

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

Caractéristiques: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effets: MemoryEffects::Effect{}

Attributs :

Résultats :

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Opération de bord de flux de données.

Syntaxe :

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

Un bord de flux de données d'une opération X définit un pont entre un ensemble de sources (chacune est un opérande de X ou un opérande du terminateur de bloc de X) et un ensemble de cibles (chacune est un résultat de X ou un argument de bloc de X), de sorte que toutes les sources et cibles doivent être fractionnées de la même manière.

Une opération peut comporter plusieurs arêtes de flux de données qui sont orthogonales les unes aux autres.

Exemple :

  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
                  })

Alors que l'opération comporte n arêtes de flux de données, la première arête du flux de données se situe entre les sources x_i, return_value_i et les cibles y_i, pred_arg_i, body_arg_i.

Un sdy.data_flow_edge prend en entrée la cible racine d'un bord (il peut s'agir de n'importe quelle cible, mais de préférence d'un résultat d'opération plutôt que d'un argument de bloc), qui ne doit pas avoir d'autre utilité. Cette opération n'est pas pure, car elle peut accepter une entrée qui n'avait initialement aucune utilité.

sdy.data_flow_edge contient également un fractionnement facultatif pour toutes les cibles du bord. Ce fractionnement doit être mis à jour au lieu du fractionnement des cibles (s'il peut être associé) lors de la propagation. Cela est utile lorsqu'une opération comporte de nombreuses arêtes, car il est beaucoup plus efficace de:

• se propagent séparément dans chaque arête.
• mettez à jour le fractionnement de chaque arête séparément au lieu de toutes les cibles à la fois (par exemple, une opération ne comporte qu'un seul TensorShardingPerValueAttr immuable pour les fractionnements de résultats).
• ajoutez chaque arête à la liste de travail séparément lorsque le fractionnement d'une source a changé.

La propagation propage les fractionnements entre toutes les sources et cibles d'un sdy.data_flow_edge comme s'il s'agissait d'une opération régulière avec les sources comme opérandes et les cibles comme résultats, et une identité sdy.op_sharding_rule. Cela signifie que la propagation avant va des sources aux cibles, et que la rétropropagation va des cibles aux sources.

Nous n'autorisons pas l'entrée d'un sdy.data_flow_edge à être définie par une opération SdyDialect. Nous pouvons donc supposer qu'elle est définie par une opération dont l'attribut sdy.sharding n'est pas enregistré.

Caractéristiques: SameOperandsAndResultType

Interfaces: InferTypeOpInterface

Attributs :

Opérandes:

Résultats :

sdy.manual_computation (sdy::ManualComputationOp)

Opération de parallélisme multi-appareil avec des collectifs manuels

Syntaxe :

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)

Accédez à une région écrite en termes de code local par appareil avec des collectifs explicites, où les formes logiques correspondent aux formes de tampon physiques locales par appareil et les collectifs correspondent exactement à la communication physique entre les appareils.

Le corps est localisé par les champs manuels_axes. La propagation se produira à travers le corps sur toutes les axes libres (celles qui ne figurent pas dans la liste "manual_axes").

Caractéristiques: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Attributs :

Opérandes:

Résultats :

sdy.mesh (sdy::MeshOp)

Maillage nommé

Syntaxe :

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

Définit un nouveau maillage nommé. Tous les maillages d'un module doivent avoir le même nombre d'appareils (à l'exception des maillages avec un seul device_id). Le maillage est une opération Symbol qui apparaît dans le SymbolTable du module et peut être référencée par son name.

Caractéristiques: HasParent<ModuleOp>

Interfaces: Symbol

Attributs :

sdy.named_computation (sdy::NamedComputationOp)

Opération de calcul nommée

Syntaxe :

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)

Regroupe un calcul, c'est-à-dire un bloc d'opérations, et lui donne un nom. La propagation s'effectue dans/hors de la région comme si tout était intégré.

Cela peut être utilisé pour gérer la propagation via des instructions d'appel vers d'autres fonctions. Tous les utilisateurs de Shardy doivent écrire un passage d'importation/d'exportation qui convertit leurs opérations d'appel en opérations sdy.named_computation, en dupliquant/copiant le corps de la fonction appelée dans le corps de la named_computation.

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 opérandes et le type de résultats de l'opération.

Exemple :

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

Caractéristiques: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp> et SingleBlock

Interfaces: ConditionallySpeculatable, ShardableDataFlowOpInterface

Attributs :

Opérandes:

Résultats :

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Opération de barrière de propagation

Syntaxe :

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

Cette opération fonctionne comme une opération d'identité, en affichant la même valeur qu'elle a prise en entrée. Toutefois, en termes de propagation, cela ne permet qu'à la propagation de s'écouler dans une certaine direction.

Cela empêche la propagation des segmentations entre les utilisations du résultat de l'opération de barrière et son opérande.

• FORWARD signifie que les fractionnements ne peuvent passer que de l'opérande au résultat.
• BACKWARD signifie que les segmentations ne peuvent circuler que du résultat vers l'opérande.
• NONE signifie qu'aucun fractionnement ne peut se propager via cette opération.
• Vous ne pouvez pas spécifier BOTH, car cette opération serait redondante.

Caractéristiques: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effets: MemoryEffects::Effect{}

Attributs :

Opérandes:

Résultats :

sdy.reshard (sdy::ReshardOp)

Répartitionne un tenseur dans un autre schéma de partitionnement.

Syntaxe :

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

Répartitionne le tenseur d'entrée avec la répartition spécifiée, qui est différente de la répartition existante du tenseur d'entrée.

ShardingConstraintOp et ReshardOp associent un fractionnement à un tenseur. Leur durée de vie est la suivante:

1. Avant la propagation du fractionnement, ShardingConstraintOp est ajouté par les utilisateurs.
2. La propagation de la segmentation consomme ShardingConstraintOp. Il n'y a pas de ShardingConstraintOp dans les résultats de la propagation de segmentation. À la place, ReshardOp peut être ajouté si nécessaire.
3. Un partitionneur convertit un ReshardOp en opération collective (ou opération d'identité). Les résultats du partitionneur ne doivent pas contenir de ReshardOp.

// TODO(b/331680067). Ajout d'un modèle de canonisation pour supprimer les opérations de reshard redondantes.

Caractéristiques: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effets: MemoryEffects::Effect{}

Attributs :

Opérandes:

Résultats :

sdy.return (sdy::ReturnOp)

L'opération sdy.return met fin aux régions associées aux opérations basées sur la région sdy et à toute autre opération basée sur la région Shardy. Elle est variable: elle prend comme arguments une liste de valeurs dont les types peuvent être n'importe quel type (mais du même genre, par exemple AnyTensor) et peut donc être réutilisée à différents niveaux de la pile IR de Shardy.

Syntaxe :

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

Caractéristiques: AlwaysSpeculatableImplTrait, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effets: MemoryEffects::Effect{}

Opérandes:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Contraint un tenseur au fractionnement spécifié

Syntaxe :

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

Associe une segmentation à un Tensor intermédiaire (par exemple, le résultat d'une fonction matmul) pour indiquer que c'est ainsi que ce Tensor, ou un sous-ensemble de ses utilisations, doit être segmenté.

Si le fractionnement comporte des dimensions ouvertes et des axes non contraints, cela signifie que le tenseur peut être fractionné davantage selon les dimensions ouvertes.

Cette opération peut soit:

• Ne pas être utilisé (dangling) : cela signifie que le fractionnement associé est la façon dont le tenseur d'entrée lui-même doit être fractionné.
• Avoir des utilisations : cela signifie que le partitionnement associé est la façon dont les utilisations de l'opération de contrainte de partitionnement doivent être partitionnées, tandis que d'autres utilisations du tenseur d'entrée peuvent avoir un partitionnement différent (si le tenseur d'entrée n'a pas d'autres utilisations, le comportement est le même que dans le cas où il n'y a pas d'utilisations).

Caractéristiques: Elementwise, SameOperandsAndResultType

Interfaces: InferTypeOpInterface

Attributs :

Opérandes:

Résultats :

sdy.sharding_group (sdy::ShardingGroupOp)

Opération de groupe de fractionnement

Syntaxe :

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

Cette opération fournit une interface permettant d'attribuer des tenseurs à des groupes de fractionnement (groupes de tenseurs qui devront avoir des fractionnements identiques). Lors de la propagation, dès qu'un élément de groupe est partitionné, tous les autres membres sont partitionnés exactement de la même manière. Cette opération prend l'ID de groupe d'arguments et ne renvoie aucun résultat, mais modifie plutôt la représentation du groupe de fractionnement interne pour ajouter le tenseur d'entrée au groupe avec l'ID donné.

Attributs :

Opérandes:

Attributs

AxisRefAttr

Référence à un axe complet ou à un sous-axe fractionné

Syntaxe :

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

Paramètres :

DimMappingAttr

Liste des indices de facteur pour une dimension

Tous les indices de facteur doivent être compris dans la plage [0, num_factors) et une liste vide indique qu'il s'agit d'un mappage nul (analysé/imprimé avec *), c'est-à-dire que la dimension n'est mappée sur aucun facteur.

Paramètres :

DimensionShardingAttr

Division des dimensions

Liste des noms d'axes sur lesquels diviser une dimension de tenseur, de l'axe principal à l'axe secondaire, un booléen indiquant si la dimension peut être divisée davantage et un entier facultatif indiquant la priorité de ce fractionnement de dimension, qui sera respecté lors de la propagation du fractionnement. Les priorités proviennent des annotations de fractionnement des utilisateurs. Une valeur inférieure indique une priorité plus élevée. La priorité la plus élevée est utilisée lorsque la priorité est manquante dans l'annotation.

Paramètres :

ManualAxesAttr

Syntaxe :

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

Paramètres :

MeshAttr

Maillage des axes et liste des appareils

Syntaxe :

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

Une maille est une liste d'axes et une liste facultative d'ID d'appareils spécifiant l'ordre des appareils.

Si la liste des axes est vide, le maillage comporte une taille d'axe implicite non nommée de 1. Dans ce cas, si aucune liste d'ID d'appareils n'est fournie, la liste implicite d'ID d'appareils est [0]. Si une liste d'ID d'appareils est fournie, elle doit contenir un entier unique de n'importe quelle valeur non négative. C'est ce que nous appelons le cas de segmentation maximale.

Pour tous les cas de partitionnement non maximal, si une liste d'ID d'appareil est spécifiée, le produit des tailles d'axe doit correspondre au nombre d'appareils. Si aucune liste d'ID d'appareil n'est spécifiée, la liste d'ID d'appareil implicite est iota(product(axes)). Pour simplifier, nous n'autorisons pas non plus de spécifier une liste d'ID d'appareil identique à iota(product(axes)). Dans ce cas, une liste d'ID d'appareil ne doit pas être spécifiée.

Voici quelques exemples de maillages:

• Un maillage vide représente un maillage d'espace réservé pouvant être remplacé lors de la propagation: <[]>
• Une maille avec une axe sans nom et un ID d'appareil explicite, qui est généralement utilisé pour représenter le fractionnement maximal: <[], device_ids=[3]>
• Maillage à deux axes et ID d'appareil implicites iota(6) : <["a"=2, "b"=3]>
• Une maille avec deux axes et des ID d'appareil explicites spécifiant l'ordre des appareils: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Paramètres :

MeshAxisAttr

Axe nommé dans un maillage

Syntaxe :

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

Paramètres :

OpShardingRuleAttr

Spécifie comment une opération peut être partitionnée.

Syntaxe :

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

Une règle de segmentation spécifie comment une opération peut être partitionnée en fonction de différentes propriétés de l'opération : attributs, forme des opérandes, forme des résultats, etc. Par exemple :

%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>

Notez que nous autorisons les facteurs de taille 1, même s'ils ne peuvent pas être partitionnés. Cela est principalement pour des raisons de complétude, car de nombreuses opérations, telles que les opérations ponctuelles, ont des dimensions de taille 1 qui correspondent entre les opérandes et les résultats.

is_custom_rule indique s'il s'agit d'une règle définie par un utilisateur pour une opération stablehlo.custom_call. Le partitionneur ne sait pas comment partitionner ces opérations. Un utilisateur doit donc lui indiquer comment procéder. S'il s'agit d'une règle personnalisée, elle est toujours conservée et jamais supprimée. is_custom_rule ne peut être défini sur "true" que pour les opérations stablehlo.custom_call.

Paramètres :

SubAxisInfoAttr

Informations sur la dérivée de ce sous-axe de l'axe complet

Syntaxe :

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

Lors de la division d'un axe complet en n sous-axes, celui-ci est remodelé en [k_1,...,k_n], et le n-ième sous-axe peut être exprimé par le produit de toutes les tailles d'axe situées à sa gauche m=prod(k_1,...,k_(i-1)) (également appelée prétaille) et de la taille k_i. Par conséquent, l'attribut sub-axis-info contient ces deux nombres et est indiqué comme suit: (m)k pour la pré-taille m et la taille k.

Paramètres :

TensorMappingAttr

Mappages de facteurs pour chaque dimension d'un tenseur.

Syntaxe :

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

Paramètres :

TensorShardingAttr

Segmentation de tenseurs

Syntaxe :

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

Une segmentation de Tensor est liée à un maillage spécifique et ne peut référencer que les noms d'axe de ce maillage. Les fractionnements de dimension nous indiquent, pour chaque dimension du tenseur, le long desquels axes (ou sous-axes) il est fractionné de majeur à mineur. Tous les autres axes qui ne fractionnent pas une dimension sont répliqués implicitement ou explicitement (s'ils figurent dans la liste des axes répliqués).

Le maillage auquel ce fractionnement est lié peut être spécifié par un nom de symbole, en référence à un symbole MeshOp correspondant ou à un MeshAttr intégré.

Paramètres :

TensorShardingPerValueAttr

Segmentation des Tensors par opérande/résultat d'une opération

Syntaxe :

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

Paramètres :

Enums

PropagationDirection

énumération de la direction de propagation

Étuis :