-sdy-close-shardings

Fecha partições de tensor e exclui eixos replicados.

-sdy-drop-sharding-rules

Exclui OpShardingRuleAttr de todas as operações registradas.

-sdy-insert-explicit-reshards

Insere reshards explícitos para que todas as operações tenham particionamentos compatíveis.

Um particionamento compatível significa que a operação pode aceitar os operandos particionados e produzir um resultado particionado sem exigir nenhuma comunicação de resparticionamento. No entanto, a operação ainda pode exigir comunicação, como redução total ou trocas de halo.

Após a propagação, algumas operações ainda podem ter particionamentos incompatíveis.

Quando um eixo (ou subeixo) é usado para dividir dimensões não correspondentes (por exemplo, dimensões não contratuais em matmul) em vários tensores ou quando um eixo divide uma dimensão em um tensor, mas não a dimensão correspondente no outro tensor, diz-se que a operação tem um conflito de fragmentação. Portanto, depois dessa passagem, as operações ficam sem conflitos.

Esse cartão injeta operações de reshard de forma explícita para que, em cada operação, as dimensões correspondentes sejam fragmentadas da mesma maneira em todos os operandos e resultados, e cada eixo (ou subeixo) só possa ser usado para dividir um único tipo de dimensão.

Um exemplo para esclarecer:

Entrada:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"y"},{"x"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

Saída:

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

No exemplo acima, há um conflito, já que os tensores lhs e rhs são divididos no eixo "x" nas dimensões não contratuais. Aqui, o tensor rhs é dividido novamente, antes da operação de ponto, explicitamente para ser dividido apenas na primeira dimensão e no eixo "x". Dessa forma, a operação de ponto passa a ser compatível.

-sdy-remove-sharding-groups

Remove ShardingGroupOps após a propagação.

-sdy-sharding-constraint-to-reshard

Converte ShardingConstraintOp em ReshardOp.

-sdy-sink-data-flow-edges

Acumula todos os DataFlowEdgeOp na entrada.

Move o fragmentação de cada DataFlowEdgeOp para a entrada (o destino raiz da borda) e substitui a operação pela entrada.

TODO(tomnatan): considere mover o particionamento para todos os destinos que podem ter um particionamento anexado.

-sdy-update-non-divisible-input-output-shardings

Faz com que as entradas/saídas de FuncOp sejam divididas de maneira uniforme, removendo a necessidade de preenchimento devido a particionamentos não divisíveis.

Os usuários do Shardy esperam que as entradas/saídas da função sejam divisíveis/fragmentáveis de maneira uniforme para evitar o preenchimento dos tensores. A propagação pode fazer com que as entradas/saídas tenham particionamentos não divisíveis. Portanto, esse processo as atualiza para o maior prefixo de particionamento de dimensão do particionamento original que é dividido de maneira uniforme.