Definições de OpInterface

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

Uma interface op que permite a fragmentação de propagar fragmentos pelas bordas do fluxo de dados das operações que estendem essa interface.

Uma borda de fluxo de dados de alguma operação X define uma ponte entre um conjunto de origens (cada uma é um operando de X ou um operando do terminador de bloco de X) e um conjunto de destinos (cada um é um resultado de X ou um argumento de bloco de X), de modo que todas as origens e destinos precisam ser particionados da mesma maneira. Uma operação pode ter várias arestas de fluxo de dados ortogonais entre si.

Um proprietário é um destino especificado pelo usuário da borda do fluxo de dados usada pela propagação do shardy. O usuário pode escolher de forma arbitrária, mas precisa ser estático.

Exemplo:

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

Esse custom_op tem dois tipos de arestas de fluxo de dados, n arestas entre return_value_i (origens) e y_i (destinos) e n arestas entre x_i(origens) e body_arg_i(destinos). Nesse caso, os proprietários das bordas são os mesmos que os destinos.

Veja aqui um exemplo de uma operação com vários destinos:

  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
                  })

Essa operação while tem n bordas de fluxo de dados, e a i-ésima borda de fluxo de dados está entre as fontes x_i, return_value_i e os destinos y_i, pred_arg_i e body_arg_i.

Métodos:

getBlockArgumentEdgeOwnerShardings

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

Retorna as fragmentações de todos os proprietários de borda do fluxo de dados do argumento de bloco.

setBlockArgumentEdgeOwnerSharding

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

Define o sharding do proprietário da borda do argumento de bloco com o index fornecido.

setBlockArgumentEdgeOwnerShardings

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

Define shardings de todos os proprietários de aresta de argumento de bloco.

getOpResultEdgeOwnerShardings

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

Retorna as fragmentações de todos os proprietários de borda do fluxo de dados do resultado da operação.

setOpResultEdgeOwnerSharding

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

Define o sharding do proprietário da borda do resultado da operação com o index fornecido.

setOpResultEdgeOwnerShardings

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

Define shardings de todos os proprietários de borda de resultados de operação.

getBlockArgumentEdgeOwners

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

Recebe todos os proprietários de borda de argumento de bloco.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Mostra todos os proprietários de borda de resultados de operações.

getEdgeSources

mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);

Recebe as origens de borda do fluxo de dados com um valor target.

getEdgeOwnerFromTarget

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

Recebe o target do proprietário de uma borda de fluxo de dados com um target que pode ou não ser o proprietário.

getEdgeOwnerFromSource

mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);

Recebe o destino do proprietário de uma borda do fluxo de dados com um source.

getNonEdgeOwnerTargets

mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);

Recebe os destinos não proprietários de uma borda de fluxo de dados conforme a borda owner.

ShardingRuleOpInterface (ShardingRuleOpInterface)

Uma interface de operação que permite que a operação defina a própria regra de fragmentação. Uma regra de fragmentação especifica como uma operação pode ser particionada de acordo com várias propriedades na operação, como qualquer atributo, a forma dos operandos, a forma dos resultados etc. Consulte OpShardingRuleAttr para mais detalhes.

Métodos:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Retorna a regra de fragmentação da operação.