Definizioni di OpInterface

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.