Definições de OpInterface

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Interface para todas as operações coletivas. Encapsula o get/set comum para o atributo outSharding.

Restrições:

• O operando precisa ter um divisão ou allowMissingInputSharding() retorna "true".
• out_sharding é válido em relação ao tipo correspondente.
• O particionamento de operandos e resultados precisa ter a mesma malha se allowDifferentMeshes() retornar falso.
• Mesma classificação para o operador e o fragmentação do resultado.

Métodos:

getOutSharding

::mlir::sdy::TensorShardingAttr getOutSharding();

Retorna o fragmentação de tensor de saída da operação coletiva.

setOutShardingAttr

void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);

Define o fragmentação do tensor de saída da operação coletiva.

getTensor

::mlir::TypedValue<::mlir::TensorType> getTensor();

Receba o operando de tensor da operação coletiva.

getType

::mlir::Type getType();

Receba o tipo do resultado da operação coletiva.

allowDifferentMeshes

bool allowDifferentMeshes();

Indica se a operação coletiva permite que o particionamento de entrada e saída tenha malhas diferentes.

allowMissingInputSharding

bool allowMissingInputSharding();

Indica se a operação coletiva permite que a entrada não tenha fragmentação, ou seja, se ela é implicitamente totalmente replicada.

ShardableDataFlowOpInterface (Sdy_ShardableDataFlowOpInterface)

Uma interface de operação que permite que o shardy propague particionamentos pelas bordas de fluxo de dados de 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 de borda são os mesmos que os destinos.

Confira um exemplo de 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::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Retorna os shardings de todos os proprietários de aresta de fluxo de dados de argumento de bloco.

setBlockArgumentEdgeOwnerShardings

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

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

getOpResultEdgeOwnerShardings

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

Retorna os particionamentos de todos os proprietários de borda do fluxo de dados de resultados de operações.

setOpResultEdgeOwnerShardings

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

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

transformTargetSharding

mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);

Transforma o sharding do destino dependendo de transformType

Consulte DataFlowShardingTransformType para mais informações.

getBlockArgumentEdgeOwners

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

Recebe todos os proprietários de arestas de argumentos de bloco.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

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

getEdgeSources

mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);

Extrai as origens de borda do fluxo de dados com base na borda owner.

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 de fluxo de dados com um source.

getNonEdgeOwnerTargets

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

Recebe as metas que não são do proprietário de uma borda de fluxo de dados com base na borda owner.

ShardingRuleOpInterface (Sdy_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 atributos, 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.