Définitions d'OpInterface

ShardableDataFlowOpInterface (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 de l'arête du flux de données utilisée par la propagation de Segment. 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::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

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

setBlockArgumentEdgeOwnerSharding

void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

Définit le sharding du propriétaire du bord de l'argument de bloc avec le index donné.

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::ArrayRef<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

Renvoie les segmentations de tous les propriétaires en périphérie du flux de données de résultat d'opérations.

setOpResultEdgeOwnerSharding

void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

Définit le sharding du propriétaire du bord du résultat de l'opération avec l'index donné.

setOpResultEdgeOwnerShardings

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

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

getBlockArgumentEdgeOwners

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

Récupère tous les propriétaires des 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::Value> getEdgeSources(mlir::Value target);

Récupère les sources de bord du flux de données à partir d'une valeur target.

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 non propriétaires d'un périmètre de flux de données en fonction de l'arête owner.

ShardingRuleOpInterface (ShardingRuleOpInterface)

Interface d'opération qui permet à l'opération de définir sa propre règle de fractionnement. 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 segmentation de l'opération.