CollectiveOpInterface (Sdy_CollectiveOpInterface
)
Interfaccia per tutte le operazioni collettive. Incapsula get/set comuni per l'attributo outSharding.
Vincoli:
- L'operando deve avere un'operazione di suddivisione o
allowMissingInputSharding()
deve restituire true. out_sharding
è valido rispetto al tipo corrispondente.- Lo sharding dell'operando e del risultato deve avere la stessa mesh se
allowDifferentMeshes()
restituisce false. - Stesso ranking per lo sharding dell'operando e del risultato.
Metodi:
getOutSharding
::mlir::sdy::TensorShardingAttr getOutSharding();
Restituisce la suddivisione in parti del tensore di output dell'operazione collettiva.
setOutShardingAttr
void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);
Imposta lo sharding del tensore di output dell'operazione collettiva.
getTensor
::mlir::TypedValue<::mlir::TensorType> getTensor();
Recupera l'operando tensore dell'operazione collettiva.
getType
::mlir::Type getType();
Recupera il tipo di risultato dell'operazione collettiva.
allowDifferentMeshes
bool allowDifferentMeshes();
Indica se l'operazione collettiva consente lo sharding di input e output di avere mesh diversi.
allowMissingInputSharding
bool allowMissingInputSharding();
Indica se l'operazione collettiva consente di non eseguire lo sharding sull'input, ovvero se è implicitamente completamente replicata.
ShardableDataFlowOpInterface (Sdy_ShardableDataFlowOpInterface
)
Un'interfaccia op che consente a Shardy di propagare gli shard tramite i bordi del flusso di dati delle op che estendono questa interfaccia.
Un bordo del flusso di dati di un'operazione X definisce un ponte tra un insieme di origini (ognuna è un operando di X o un operando del terminatore di blocco di X) e un insieme di destinazioni (ognuna è un risultato di X o un argomento del blocco di X), in modo che tutte le origini e le destinazioni debbano essere suddivise in modo uguale. Un'operazione può avere più bordi di flusso di dati ortogonali tra loro.
Un proprietario è un target specificato dall'utente dell'elemento di flusso di dati utilizzato dalla propagazione di Shardy. L'utente può sceglierlo in modo arbitrario, ma deve essere statico.
Ad esempio:
y_1, ..., y_n = custom_op (x_1, ..., x_n)
((body_arg_1,..., body_arg_n) {
...
return return_value_1, ..., return_value_n
})
Questa custom_op ha due tipi di archi di flusso di dati, ciascuno con n archi tra
return_value_i
(origini) e y_i
(destinazioni) e n archi tra
x_i
(origini) e body_arg_i
(destinazioni). In questo caso, i proprietari degli edge sono uguali ai target.
Ecco un esempio di un'operazione con più target:
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
})
Questa operazione while ha n bordi di flusso di dati, il bordo di flusso di dati i-esimo è tra le origini x_i
, return_value_i
e le destinazioni y_i
, pred_arg_i
,
body_arg_i
.
Metodi:
getBlockArgumentEdgeOwnerShardings
mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();
Restituisce gli shard di tutti i proprietari di edge del flusso di dati degli argomenti del blocco.
setBlockArgumentEdgeOwnerShardings
void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
Imposta shardings
di tutti i proprietari di edge dell'argomento blocco.
getOpResultEdgeOwnerShardings
mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();
Restituisce gli shard di tutti i proprietari di edge del flusso di dati dei risultati dell'operazione.
setOpResultEdgeOwnerShardings
void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
Imposta shardings
di tutti i proprietari di edge dei risultati dell'operazione.
transformTargetSharding
mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);
Trasforma il sharding
del target in base a transformType
Consulta DataFlowShardingTransformType
per saperne di più.
getBlockArgumentEdgeOwners
mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();
Recupera tutti i proprietari di edge degli argomenti del blocco.
getOpResultEdgeOwners
mlir::ResultRange getOpResultEdgeOwners();
Recupera tutti i proprietari di edge dei risultati dell'operazione.
getEdgeSources
mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);
Recupera le origini dell'elemento laterale del flusso di dati in base all'elemento laterale owner
.
getEdgeOwnerFromTarget
mlir::Value getEdgeOwnerFromTarget(mlir::Value target);
Recupera il proprietario target
di un bordo del flusso di dati in base a un target
che potrebbe o meno essere il proprietario.
getEdgeOwnerFromSource
mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);
Recupera il proprietario target di un bordo del flusso di dati in base a un source
.
getNonEdgeOwnerTargets
mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);
Recupera i target non proprietari di un bordo del flusso di dati in base al bordo owner
.
ShardingRuleOpInterface (Sdy_ShardingRuleOpInterface
)
Un'interfaccia op che consente all'operatore di definire la propria regola di suddivisione.
Una regola di suddivisione in parti specifica in che modo un'operazione può essere suddivisa in base a varie proprietà dell'operazione, ad esempio qualsiasi attributo, la forma degli operandi, la forma dei risultati e così via. Per maggiori dettagli, consulta OpShardingRuleAttr
.
Metodi:
getShardingRule
mlir::sdy::OpShardingRuleAttr getShardingRule();
Restituisce la regola di sharding dell'operazione.