Definizioni di OpInterface

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

Un'interfaccia operativa che consente allo shardy di propagare gli sharding attraverso i margini del flusso di dati delle operazioni che estendono l'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 blocchi di X) e un insieme di destinazioni (ognuno è il risultato di X o un argomento di blocco di X), in modo che tutte le origini e tutte le destinazioni debbano essere sottoposte a sharding nello stesso modo. Un'operazione può avere più spigoli di flusso di dati ortogonali tra loro.

Un proprietario è un target specificato dall'utente dell'edge del 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
                  })

Questo custom_op ha due tipi di lati del flusso di dati, n lati ciascuno tra return_value_i (origini) e y_i (target) e n lati tra x_i(origini) e body_arg_i(target). In questo caso, i proprietari del perimetro corrispondono alle destinazioni.

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

Restituisce lo sharding di tutti i proprietari periferici del flusso di dati con argomento a blocchi.

setBlockArgumentEdgeOwnerSharding

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

Imposta sharding dell'edge owner dell'argomento del blocco con il valore index specificato.

setBlockArgumentEdgeOwnerShardings

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

Imposta shardings di tutti i proprietari di edge dell'argomento blocco.

getOpResultEdgeOwnerShardings

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

Restituisce gli shard di tutti i proprietari di edge del flusso di dati dei risultati dell'operazione.

setOpResultEdgeOwnerSharding

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

Imposta il sharding del proprietario dell'edge del risultato dell'operazione con il index specificato.

setOpResultEdgeOwnerShardings

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

Imposta shardings di tutti i proprietari di edge dei risultati dell'operazione.

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::Value> getEdgeSources(mlir::Value target);

Recupera le origini esterne del flusso di dati in base a un valore target.

getEdgeOwnerFromTarget

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

Recupera il proprietario target di un bordo del flusso di dati dato un target che può 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 (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.