Définitions d'OpInterface

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Interface pour toutes les opérations collectives. Encapsule les opérations get/set courantes pour l'attribut de fractionnement externe.

Contraintes :

  • L'opérande doit avoir un fractionnement, ou allowMissingInputSharding() doit renvoyer "true".
  • out_sharding est valide par rapport au type correspondant.
  • Le fractionnement de l'opérande et du résultat doit avoir le même maillage si allowDifferentMeshes() renvoie la valeur "false".
  • Même rang pour le fractionnement de l'opérande et du résultat.

Méthodes :

getOutSharding

::mlir::sdy::TensorShardingAttr getOutSharding();

Renvoie le fractionnement du tenseur de sortie de l'opération collective.

setOutShardingAttr

void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);

Définit le fractionnement du tenseur de sortie de l'opération collective.

getTensor

::mlir::TypedValue<::mlir::TensorType> getTensor();

Obtenez l'opérande de tenseur de l'opération collective.

getType

::mlir::Type getType();

Obtenez le type du résultat de l'opération collective.

allowDifferentMeshes

bool allowDifferentMeshes();

Indique si l'opération collective permet au fractionnement des entrées et des sorties d'avoir des maillages différents.

allowMissingInputSharding

bool allowMissingInputSharding();

Indique si l'opération collective permet à l'entrée de ne pas avoir de fractionnement, c'est-à-dire de la répliquer entièrement de manière implicite.

ShardableDataFlowOpInterface (Sdy_ShardableDataFlowOpInterface)

Interface d'opération qui permet à shardy de propager des fractionnements via les bords de flux de données des opérations qui étendent cette interface.

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.

Un propriétaire est une cible spécifiée par l'utilisateur du bord de flux de données utilisé par la propagation de shardy. L'utilisateur peut le choisir arbitrairement, mais il doit être statique.

Exemple :

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

Cette custom_op comporte deux types d'arêtes de flux de données, n'arêtes chacune entre return_value_i (sources) et y_i (cibles) et n'arêtes entre x_i(sources) et body_arg_i(cibles). Dans ce cas, les propriétaires des bords sont les mêmes que les cibles.

Voici un exemple d'opération avec plusieurs cibles:

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

Cette opération while comporte n arcs de flux de données, les arcs de flux de données i se situant entre les sources x_i, return_value_i et les cibles y_i, pred_arg_i, body_arg_i.

Méthodes :

getBlockArgumentEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Renvoie les fractionnements de tous les propriétaires de bords de flux de données d'arguments de bloc.

setBlockArgumentEdgeOwnerShardings

void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Définit shardings de tous les propriétaires de bords d'arguments de bloc.

getOpResultEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

Renvoie les fractionnements de tous les propriétaires de bords de flux de données de résultats d'opération.

setOpResultEdgeOwnerShardings

void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Définit shardings de tous les propriétaires de bords de résultat d'opération.

transformTargetSharding

mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);

Transforme le sharding de la cible en fonction de transformType

Pour en savoir plus, consultez DataFlowShardingTransformType.

getBlockArgumentEdgeOwners

mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();

Récupère tous les propriétaires de bords d'argument de bloc.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Récupère tous les propriétaires de bords de résultats d'opération.

getEdgeSources

mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);

Récupère les sources de bord de flux de données en fonction du bord owner.

getEdgeOwnerFromTarget

mlir::Value getEdgeOwnerFromTarget(mlir::Value target);

Récupère le propriétaire target d'un bord de flux de données donné un target qui peut ou non être le propriétaire.

getEdgeOwnerFromSource

mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);

Récupère la cible propriétaire d'un bord de flux de données à partir d'un source.

getNonEdgeOwnerTargets

mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);

Récupère les cibles autres que le propriétaire d'un bord de flux de données donné le bord owner.

ShardingRuleOpInterface (Sdy_ShardingRuleOpInterface)

Interface d'opération qui permet à l'opération de définir sa propre règle de partitionnement. Une règle de fractionnement spécifie comment une opération peut être partitionnée en fonction de diverses propriétés de l'opération (attributs, forme des opérandes, forme des résultats, etc.). Pour en savoir plus, consultez OpShardingRuleAttr.

Méthodes :

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Renvoie la règle de partitionnement de l'opération.