-sdy-close-shardings

Chiude le suddivisioni dei tensori ed elimina gli assi replicati.

-sdy-drop-sharding-rules

Abbassa OpShardingRuleAttr da tutte le opzioni registrate.

-sdy-insert-explicit-reshards

Inserisce reshard espliciti per fare in modo che tutte le operazioni abbiano suddivisioni compatibili.

Uno sharding compatibile significa essenzialmente che l'operazione può accettare gli operatori suddivisi in parti e produrre un risultato suddiviso in parti senza richiedere comunicazioni di suddivisione in parti (tieni presente che l'operazione potrebbe comunque richiedere comunicazioni come all-reduce o halo-swap).

Dopo la propagazione, alcune operazioni potrebbero avere ancora partizioni incompatibili.

Tieni presente che quando un asse (o un asse secondario) viene utilizzato per suddividere dimensioni non corrispondenti (ad es. dimensioni non contrattabili in matmul) in più tensori o quando un asse suddivide una dimensione in un tensore, ma non la dimensione corrispondente nell'altro tensore, si dice che l'operazione presenta un conflitto di suddivisione. Di conseguenza, dopo questo passaggio, le operazioni diventano senza conflitti.

Questo passaggio inserisce esplicitamente le operazioni di suddivisione in modo che, per ogni operazione, le dimensioni corrispondenti vengano suddivise nello stesso modo in tutti gli operandi e i risultati e ogni asse (o asse secondario) possa essere utilizzato solo per suddividere un singolo tipo di dimensione.

Un esempio per chiarire:

Input:

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>

Output:

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>

Nell'esempio riportato sopra, si verifica un conflitto poiché i tensori lhs e rhs sono entrambi suddivisi in parti sull'asse "x" nelle dimensioni non contrattabili. Qui, rhs viene suddiviso nuovamente, prima dell'operazione di moltiplicazione, in modo esplicito per essere suddiviso solo per la prima dimensione e per l'asse "x". In questo modo, l'operazione con il punto diventa compatibile.

-sdy-remove-sharding-groups

Rimuove ShardingGroupOps dopo la propagazione.

-sdy-sharding-constraint-to-reshard

Converte ShardingConstraintOp in ReshardOp.

-sdy-sink-data-flow-edges

Incanala tutto il DataFlowEdgeOp nell'input.

Sposta lo sharding di ogni DataFlowEdgeOp al relativo input (il target principale dell'edge) e sostituisce l'operazione con il relativo input.

DA FARE(tomnatan): valuta la possibilità di spostare lo sharding su tutti i target a cui può essere associato un sharding.

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

Esegue lo sharding uniforme degli input/output di FuncOp, eliminando la necessità di spaziatura aggiuntiva a causa di sharding non divisibile.

Gli utenti di Shardy si aspettano che gli input/output della funzione siano equamente divisibile/suddividibili per evitare di dover aggiungere spazi aggiuntivi ai tensori. La propagazione può fare in modo che gli input/output abbiano suddivisioni non divisibile, quindi questo passaggio li aggiorna al prefisso di suddivisione delle dimensioni più grande del suddivisione originale che è suddiviso in modo uniforme.