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.