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.