Dialetto 'sdy'

Il dialetto Shardy (SDY)

Il dialetto Shardy (SDY) definisce una rappresentazione dello sharding dei tensori basata sugli assi e componenti API aggiuntivi per collegare gli sharding ai tensori.

Log delle versioni: 0.0.1: Aggiungi assi non ridotti a TensorShardingAttr.

Operazioni

sdy.all_gather (sdy::AllGatherOp)

Esegue una comunicazione all-gather lungo gli assi

Sintassi:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Raccoglie blocchi di un tensore lungo gli assi specificati in gathering_axes.

gathering_axes è un elenco di elenchi di assi. L'elenco esterno supera le dimensioni del tensore. Ogni elenco interno specifica gli assi lungo i quali deve essere eseguita una raccolta separata sulla rispettiva dimensione. Verrà applicato allo sharding dell'operando (tensor) per ottenere lo sharding del risultato (out_sharding).

Tieni presente che out_sharding non viene utilizzato per determinare lo sharding del risultato. Lo sharding del risultato è invece determinato dallo sharding dell'operando e di gathering_axes e out_sharding deve corrispondere a questo sharding dedotto.

Esempio:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• Gli elementi in gathering_axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• L'applicazione di gathering_axes allo sharding dell'operando restituisce out_sharding.

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.all_reduce (sdy::AllReduceOp)

Esegui una comunicazione all-reduce lungo gli assi

Sintassi:

operation ::= `sdy.all_reduce` ($reduction_op^)? $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Riduce i blocchi di un tensore lungo gli assi specificati in reduction_axes. L'ordine di reduction_axes non è importante per il risultato, ma può influire sull'ordine dei gruppi di repliche corrispondenti.

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• reduction_axes deve soddisfare i vincoli elencati in AxisRefListAttr.
• reduction_axes deve essere ordinato rispetto alla mesh.
• Lo sharding degli operandi e out_sharding devono avere sharding delle dimensioni equivalenti.
• reduction_axes non deve sovrapporsi alla distribuzione delle dimensioni dell'operando e agli assi replicati (può sovrapporsi agli assi non ridotti).
• reduction_axes non deve sovrapporsi agli assi non ridotti di out_sharding. In altre parole, out_sharding deve essere replicato lungo reduction_axes (in modo implicito o esplicito).

Tratti: SameOperandsAndResultType

Interfacce: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.all_slice (sdy::AllSliceOp)

Esegue un'operazione di sezionamento dinamico lungo gli assi

Sintassi:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Suddivide i blocchi di un tensore lungo gli assi specificati in slicing_axes. Esiste una dualità algebrica tra sdy.all_slice e sdy.all_gather.

slicing_axes è un elenco di elenchi di assi. L'elenco esterno supera le dimensioni del tensore. Ogni elenco interno specifica gli assi lungo i quali deve essere eseguita una sezione sulla dimensione corrispondente. Verrà applicato allo sharding dell'operando (tensor) per ottenere lo sharding del risultato (out_sharding).

Tieni presente che out_sharding non viene utilizzato per determinare lo sharding del risultato. Lo sharding del risultato è invece determinato dallo sharding dell'operando e di slicing_axes e out_sharding deve corrispondere a questo sharding dedotto.

Esempio:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• Gli elementi in slicing_axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• L'applicazione di slicing_axes allo sharding dell'operando restituisce out_sharding.

Tratti: SameOperandsAndResultType

Interfacce: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.all_to_all (sdy::AllToAllOp)

Esegue una comunicazione all-to-all lungo gli assi

Sintassi:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Per ogni tupla (axes, src_dim, tgt_dim) nell'elenco dei parametri, questa operazione suddivide i blocchi di un tensore lungo la dimensione tgt_dim e gli assi specificati in axes, li distribuisce lungo gli assi e li concatena lungo la dimensione src_dim.

Questa operazione è essenzialmente una combinazione di un'operazione all-gather lungo src_dim e axes, seguita da un'operazione all-slice lungo tgt_dim e axes, ovvero un suffisso della dimensione di sharding degli assi src_dim sul tensore di input viene aggiunto alla dimensione di sharding degli assi tgt_dim sul tensore di output.

L'operazione all-to-all verrà applicata allo sharding dell'operando (tensor) per ottenere lo sharding del risultato (out_sharding).

Tieni presente che out_sharding non viene utilizzato per determinare lo sharding del risultato. Al contrario, lo sharding del risultato è determinato dallo sharding dell'operando, src_dim, tgt_dim e axes, e out_sharding deve corrispondere a questo sharding dedotto.

Esempio:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• L'elenco dei parametri non deve essere vuoto.
• Per ogni parametro in params:
  • Gli elementi in axes devono soddisfare i vincoli di AxisRefAttr.
  • src_dim e tgt_dim devono essere dimensioni valide (non negative e inferiori al rango del tensore).
  • Qualsiasi src_dim o tgt_dim deve essere univoco in tutti i parametri.
  • src_dim deve essere ordinato in ordine crescente in tutti i parametri.
• Lo spostamento di axes da src_dim a tgt_dim nello sharding degli operandi out_sharding.

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.collective_permute (sdy::CollectivePermuteOp)

Esegue una comunicazione collective-permute per sostituire gli assi

Sintassi:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Invia un blocco del tensore di input da un dispositivo all'altro per riordinare/sostituire gli assi che suddividono il tensore.

Una permutazione collettiva può trasformare lo sharding di input in modo che ogni dimensione deve essere suddivisa in shard come prima, ovvero deve essere suddivisa in shard lungo gli assi il cui prodotto delle dimensioni corrisponde a quello degli assi che in precedenza hanno suddiviso il tensore in shard.

Questa funzionalità è utile per riordinare gli assi in una singola dimensione o in dimensioni diverse e per scambiare gli assi suddivisi con quelli replicati.

Nell'esempio riportato di seguito, la dimensione del tensore suddiviso è tensor<1x4x2xf32>, che viene mantenuta dalla permutazione collettiva.

Esempio:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• Se lo sharding di input e output ha mesh diverse, queste devono avere esattamente gli stessi assi e un ordine diverso degli ID dispositivo.
• Per ogni dimensione, il prodotto delle dimensioni dell'asse di sharding in out_sharding deve corrispondere a quello dello sharding della dimensione dell'operando corrispondente.

Tratti: SameOperandsAndResultType

Interfacce: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.constant (sdy::ConstantOp)

Funzionamento costante

Produce un tensore output da una costante value.

Vedi: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Esempio:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

Tratti: AlwaysSpeculatableImplTrait

Interfacce: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effetti: MemoryEffects::Effect{}

Attributi:

Risultati:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Operazione perimetrale del flusso di dati

Sintassi:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

Un bordo del flusso di dati di un'operazione X definisce un ponte tra un insieme di origini (ognuna è un operando di X o un operando del terminatore di blocchi di X) e un insieme di destinazioni (ognuna è un risultato di X o un argomento di blocco di X), in modo che tutte le origini e le destinazioni vengano suddivise nello stesso modo.

Un'operazione può avere più bordi del flusso di dati ortogonali tra loro.

Ad esempio:

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

Questo ciclo while ha n archi di flusso di dati, l'i-esimo arco di flusso di dati si trova tra le origini x_i, return_value_i e le destinazioni y_i, pred_arg_i, body_arg_i.

Un sdy.data_flow_edge prende come input il proprietario di un bordo (può essere uno qualsiasi dei target, ma preferibilmente un risultato dell'operazione anziché un argomento del blocco), che non deve avere altri utilizzi. Questa operazione non è pura perché può accettare un input che originariamente non aveva utilizzi.

sdy.data_flow_edge contiene anche uno sharding facoltativo per tutti i target dell'edge e questo sharding deve essere aggiornato anziché lo sharding dei target (se può essere allegato) durante la propagazione. Questo è utile quando un'operazione ha molti bordi, in quanto è molto più efficiente:

• si propagano attraverso ogni arco separatamente.
• aggiorna lo sharding di ogni edge separatamente anziché di tutti i target contemporaneamente (ad es. un'operazione ha un singolo TensorShardingPerValueAttr immutabile per gli sharding dei risultati).
• aggiungi ogni limite all'elenco di lavoro separatamente quando lo sharding di un'origine è cambiato.

La propagazione propaga le suddivisioni tra tutte le origini e le destinazioni di un sdy.data_flow_edge come se fosse un'operazione normale con le origini come operandi e le destinazioni come risultati e un'identità sdy.op_sharding_rule. Ciò significa che la propagazione in avanti va dalle origini alle destinazioni e la propagazione all'indietro va dalle destinazioni alle origini.

Non consentiamo che l'input di un sdy.data_flow_edge sia definito da un'operazione SdyDialect, quindi possiamo presumere che sia definito da un'operazione con l'attributo sdy.sharding non registrato.

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

Operazione di edge del flusso di dati di input/output della funzione

Sintassi:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

Un'operazione di edge del flusso di dati, ma per gli argomenti della funzione o i risultati della chiamata. Quando il suo operando è un BlockArgument, funge da ponte dall'argomento callOp del chiamante agli utenti dell'argomento func. Esiste un bordo del flusso di dati func per ogni argomento func. Quando il suo operando è un OpResult, è un ponte dal valore restituito di funcOp chiamato agli utenti del risultato della chiamata. Esiste un bordo del flusso di dati func per ogni risultato della chiamata.

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, SymbolUserOpInterface

Operandi:

Risultati:

sdy.manual_computation (sdy::ManualComputationOp)

Operazione di parallelismo multi-dispositivo con collettivi manuali

Sintassi:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

Passa a una regione scritta in termini di codice locale per dispositivo con collettivi espliciti, in cui le forme logiche corrispondono alle forme e ai collettivi del buffer fisico locale per dispositivo corrispondono esattamente alla comunicazione fisica cross-device.

Il corpo è locale rispetto a manual_axes. La propagazione avverrà attraverso il corpo su tutti gli assi liberi, ovvero quelli non presenti nell'elenco manual_axes.

Tieni presente che i tensori senza classificazione devono avere uno sharding con rango 0, ovvero essere completamente replicati.

Vincoli:

• Gli elementi in in_shardings e out_shardings devono soddisfare i vincoli elencati in TensorShardingAttr.
• Il numero di input/output tensore globali e locali della regione dell'operazione deve corrispondere.
• Gli assi manuali devono precedere gli assi liberi in ogni sharding delle dimensioni.
• Gli assi manuali non possono introdurre spaziatura interna. ovvero la dimensione della dimensione deve essere divisibile per la dimensione degli assi manuali corrispondenti.
• Le forme globali e locali degli argomenti/risultati delle regioni operative devono corrispondere.

Tratti: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfacce: ShardableDataFlowOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.mesh (sdy::MeshOp)

Mesh con nome

Sintassi:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

Definisce una nuova mesh denominata. Tutte le mesh di un modulo devono avere lo stesso numero di dispositivi (ad eccezione delle mesh con un singolo device_id). La mesh è un'operazione Symbol che viene visualizzata in SymbolTable del modulo e a cui è possibile fare riferimento tramite il relativo name.

Tratti: HasParent<ModuleOp>, SymbolName

Interfacce: Symbol

Attributi:

sdy.named_computation (sdy::NamedComputationOp)

Operazione di calcolo denominata

Sintassi:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

Raggruppa un calcolo, ovvero un blocco di operazioni, e gli assegna un nome. La propagazione avverrà all'interno e all'esterno della regione come se tutto fosse in linea.

Può essere utilizzato per gestire la propagazione delle istruzioni di chiamata ad altre funzioni. Gli utenti di Shardy devono scrivere una pass di importazione/esportazione che converte le operazioni di chiamata in operazioni sdy.named_computation, duplicando/copiando il corpo della funzione chiamata nel corpo di named_computation.

Il tipo di ogni argomento del blocco e dei valori restituiti nella regione deve essere lo stesso del tipo degli operandi e del tipo di risultato dell'operazione.

Esempio:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

Tratti: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfacce: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Funzionamento della barriera di propagazione

Sintassi:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

Questa operazione funziona come un'operazione di identità, restituendo lo stesso valore che ha ricevuto come input. Tuttavia, in termini di propagazione, consentirà il flusso solo in una determinata direzione.

Ciò impedisce la propagazione degli sharding tra gli utilizzi del risultato dell'operazione di barriera e del relativo operando.

• FORWARD significa che gli sharding possono fluire solo dall'operando al risultato.
• BACKWARD significa che gli sharding possono fluire solo dal risultato all'operando.
• NONE significa che lo sharding non può propagarsi tramite questa operazione.
• Impossibile specificare BOTH, in quanto questa operazione sarebbe ridondante.

Tratti: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfacce: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effetti: MemoryEffects::Effect{}

Attributi:

Operandi:

Risultati:

sdy.reduce_scatter (sdy::ReduceScatterOp)

Esegue una comunicazione di riduzione della dispersione lungo gli assi

Sintassi:

operation ::= `sdy.reduce_scatter` ($reduction_op^)? $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Riduce i blocchi di un tensore lungo gli assi specificati in reduce_scatter_axes e poi distribuisce il risultato lungo gli stessi assi. Questa operazione è essenzialmente una combinazione di un sdy.all_reduce seguito da un sdy.all_slice lungo lo stesso reduce_scatter_axes.

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• Gli elementi in reduce_scatter_axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• L'applicazione di reduce_scatter_axes allo sharding dell'operando restituisce out_sharding.

Tratti: SameOperandsAndResultType

Interfacce: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Sposta gli assi replicati in modo implicito o esplicito sugli assi non ridotti.

Sintassi:

operation ::= `sdy.replicated_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes deve essere replicato in modo implicito o esplicito nell'operando. Questa operazione li rende non ridotti nel risultato. Abbiamo il seguente rapporto:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

Esempio:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}, {}\], replicated={"c", "d"}, unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.replicated_to_unreduced {"a", "c", "f"} %1 out_sharding=<@mesh, [{"b"}, {}, {}\], replicated={"d"}, unreduced={"a", "c", "e", "f"}> : tensor<8x8x8xf32>

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• axes deve soddisfare i vincoli elencati in AxisRefListAttr.
• axes deve essere ordinato rispetto alla mesh.
• axes non sono vuoti.
• Lo sharding di input e output deve avere gli stessi shard di dimensione.
• axes deve essere replicato in modo implicito o esplicito nello sharding degli operandi.
• inUnreducedAxes + axes = outUnreducedAxes.

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.reshard (sdy::ReshardOp)

Esegue il resharding di un tensore in uno sharding diverso

Sintassi:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

Esegue il resharding del tensore di input con lo sharding specificato, che è diverso dallo sharding esistente del tensore di input.

Sia ShardingConstraintOp che ReshardOp collegano uno sharding a un tensore. La loro durata è:

1. Prima della propagazione dello sharding, ShardingConstraintOp viene aggiunto dagli utenti.
2. La propagazione dello sharding utilizza ShardingConstraintOp. Non è presente ShardingConstraintOp nei risultati della propagazione dello sharding. In alternativa, ReshardOp può essere aggiunto se necessario.
3. Un partizionatore converte un'operazione ReshardOp in un'operazione collettiva (o un'operazione identità). Non deve essere presente alcun ReshardOp nei risultati del partizionatore.

Tratti: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfacce: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), SymbolUserOpInterface

Effetti: MemoryEffects::Effect{}

Attributi:

Operandi:

Risultati:

sdy.return (sdy::ReturnOp)

L'operazione sdy.return termina le regioni associate alle operazioni basate sulla regione sdy e a qualsiasi altra operazione basata sulla regione Shardy. È variadica: accetta come argomenti un elenco di valori i cui tipi possono essere qualsiasi (ma dello stesso tipo, ad es. AnyTensor) e pertanto può essere riutilizzata a vari livelli dello stack Shardy IR.

Sintassi:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

Tratti: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

Interfacce: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effetti: MemoryEffects::Effect{}

Operandi:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

Sposta alcuni assi suddivisi dell'operando sugli assi non ridotti del risultato.

Sintassi:

operation ::= `sdy.sharded_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes deve essere utilizzato per partizionare l'operando. Questa operazione li rende non ridotti nel risultato. Abbiamo il seguente rapporto:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), dove all-gather, sharded-to-unreduced, all-reduce vengono applicati sugli stessi assi.

Esempio:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\], unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.sharded_to_unreduced [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\], unreduced={"b", "c", "d", "e"}> : tensor<8x8x8xf32>

Vincoli:

• Deve soddisfare i vincoli elencati in Sdy_CollectiveOpInterface.
• Gli elementi in axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• L'applicazione di axes allo sharding dell'operando restituisce out_sharding.

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Vincola un tensore allo sharding specificato

Sintassi:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

Collega uno sharding a un tensore intermedio (ad es. il risultato di una moltiplicazione di matrici) per indicare che questo è il modo in cui deve essere suddiviso in shard il tensore o un sottoinsieme dei suoi utilizzi.

Se lo sharding ha dimensioni aperte e assi senza vincoli, significa che il tensore può essere ulteriormente suddiviso lungo le dimensioni aperte.

Questa operazione può:

• Non hanno utilizzi (dangling), il che significa che lo sharding allegato indica come deve essere suddiviso il tensore di input.
• Hanno utilizzi, il che significa che lo sharding allegato è il modo in cui gli utilizzi dell'operazione di vincolo di sharding devono essere suddivisi, mentre altri utilizzi del tensore di input potrebbero avere uno sharding diverso (se il tensore di input non ha altri utilizzi, il comportamento è lo stesso del caso senza utilizzi).

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface, SymbolUserOpInterface

Attributi:

Operandi:

Risultati:

sdy.sharding_group (sdy::ShardingGroupOp)

Vincola i tensori nel gruppo ad avere lo stesso sharding.

Sintassi:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

Questa operazione fornisce un'interfaccia per assegnare i tensori ai gruppi di sharding ( gruppi di tensori a cui verrà imposto di avere sharding identici). Durante la propagazione, non appena viene partizionato un elemento del gruppo, tutti gli altri membri verranno partizionati esattamente nello stesso modo. Questa operazione accetta l'ID gruppo di argomenti e non restituisce alcun risultato, ma modifica la rappresentazione interna del gruppo di sharding per aggiungere il tensore di input al gruppo con l'ID specificato.

Interfacce: InferTypeOpInterface

Attributi:

Operandi:

Attributi

AllToAllParamAttr

Parametro all-to-all

Sintassi:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

Una tupla contenente gli assi e le dimensioni di origine/destinazione su cui eseguire l'operazione all-to-all.

Parametri:

AllToAllParamListAttr

Elenco di tutti i parametri all-to-all

Sintassi:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

Parametri:

AxisRefAttr

Riferimento a un asse completo o a un sottoasse diviso

Sintassi:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

Vincoli:

• name deve essere presente in MeshAttr associato.
• Se sub_axis_info è presente, deve soddisfare i vincoli di SubAxisInfoAttr.

Parametri:

AxisRefListAttr

Elenco dei riferimenti agli assi

Sintassi:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

Vincoli:

• Gli elementi in value devono soddisfare i vincoli di AxisRefAttr.
• Non sono presenti riferimenti agli assi o assi secondari duplicati che si sovrappongono tra loro.
• Due riferimenti agli assi adiacenti non sono sotto-assi consecutivi dello stesso asse completo, ovvero possono essere uniti in un unico sotto-asse o nell'asse completo.

Parametri:

AxisToPropagationDetailsAttr

Dettagli del flusso di propagazione per un'origine e un asse specifici.

Sintassi:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

Mappa un riferimento al valore di origine a un elenco di riferimenti al valore target lungo un asse specifico.

Parametri:

DimMappingAttr

Elenco degli indici dei fattori per una dimensione

Un elenco vuoto indica che si tratta di una mappatura nulla (analizzata/stampata con *), ovvero la dimensione non è mappata a nessun fattore.

Vincoli:

• Esiste almeno un indice dei fattori.
• Gli indici dei fattori devono essere compresi nell'intervallo [0, $factor_sizes).
• Se sono presenti più fattori, nessuno può avere dimensione 1.
• Nessun indice di fattori duplicato.

Parametri:

DimensionShardingAttr

Partizionamento orizzontale delle dimensioni

Elenco dei nomi degli assi su cui partizionare una dimensione del tensore da maggiore a minore, un valore booleano che indica se la dimensione può essere ulteriormente partizionata e un numero intero facoltativo che indica la priorità di questa partizione della dimensione, che verrà rispettata durante la propagazione della partizione. Le priorità derivano dalle annotazioni di sharding degli utenti e un valore più basso indica una priorità più elevata. La priorità più alta viene presupposta quando la priorità non è presente nell'annotazione.

Vincoli:

• Gli elementi in axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• Se uno sharding delle dimensioni ha una priorità:
  • La priorità è maggiore o uguale a 0.
  • Se la dimensione è chiusa, ha almeno un asse.

Parametri:

EdgeValueRefAttr

Riferimento a un particolare indice di un bordo di valore di tipo type.

Sintassi:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

Parametri:

ListOfAxisRefListsAttr

Elenco degli elenchi di riferimento degli assi

Sintassi:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

Parametri:

ManualAxesAttr

Un elenco di assi su cui un'operazione ManualComputationOp è manuale

Sintassi:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

Parametri:

MeshAttr

Mesh di assi e un elenco di dispositivi

Sintassi:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

Una mesh è un elenco di assi e un elenco facoltativo di ID dispositivo che specifica l'ordine dei dispositivi.

Se l'elenco degli assi è vuoto

• Se device_ids non viene fornito, si tratta di una mesh vuota.
• Se viene fornito device_ids, deve essere un singolo numero intero non negativo, che chiamiamo mesh di sharding massimale.

Se viene fornito l'elenco degli assi

• Se viene specificato un elenco di ID dispositivo, il prodotto delle dimensioni degli assi deve corrispondere al numero di dispositivi.
• Se non viene specificato un elenco di ID dispositivo, l'elenco implicito di ID dispositivo è iota(product(axes)). Per semplicità, non è consentito specificare un elenco di ID dispositivo uguale a iota(product(axes)); in questo caso, non deve essere specificato un elenco di ID dispositivo.
• Non è una mesh con sharding massimo anche se la dimensione totale degli assi è 1.

Ecco alcuni esempi di mesh:

• Una mesh vuota rappresenta una mesh segnaposto che può essere sostituita durante la propagazione: <[]>
• Una mesh senza elenco di assi e un singolo ID dispositivo non negativo, ovvero una mesh con sharding massimo: <[], device_ids=[3]>
• Una mesh con due assi e ID dispositivo impliciti iota(6): <["a"=2, "b"=3]>
• Una mesh con due assi e ID dispositivo espliciti che specificano l'ordine dei dispositivi: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Vincoli:

• Gli elementi in device_ids non devono essere negativi.
• Se axes è vuoto, la dimensione di device_ids può essere 0 (mesh vuota) o 1 (mesh di sharding massima).
• Se axes non è vuoto,
  • Gli elementi in axes non devono avere nomi duplicati.
  • Se device_ids è specificato, l'device_ids originale non è iota(product(axis_sizes)) e l'device_ids ordinato è iota(product(axis_sizes)).

Parametri:

MeshAxisAttr

Asse denominato in una mesh

Sintassi:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

Parametri:

OpShardingRuleAttr

Specifica come può essere partizionata un'operazione.

Sintassi:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  bool   # is_custom_rule
>

Una regola di sharding specifica come un'operazione può essere partizionata in base a varie proprietà dell'operazione, ad esempio attributi, forma degli operandi, forma dei risultati e così via. Ad esempio:

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>

%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

Tieni presente che consentiamo fattori di dimensione 1 anche se non possono essere suddivisi in shard. Ciò è dovuto principalmente al fatto che molte operazioni, come quelle puntuali, hanno dimensioni pari a 1 che corrispondono a operandi e risultati.

Tipi di fattori:

• reduction_factors contiene gli indici dei fattori che richiedono la riduzione, come le dimensioni di contrazione in un'operazione dot. Questi fattori possono essere negli operandi, ma non nei risultati.
• need_replication_factors contiene gli indici dei fattori che richiedono la replica completa, ad esempio la dimensione ordinata in un'operazione di ordinamento.
• permutation_factors contiene gli indici dei fattori che richiedono collective-permute se sono suddivisi in shard, ad esempio le dimensioni del padding in un'operazione di padding.
• Tutti gli altri fattori sono considerati fattori pass-through, ovvero fattori che non richiedono alcuna comunicazione se vengono suddivisi nello stesso modo in tutti i tensori a cui sono mappati.

blocked_propagation_factors contiene i fattori in base ai quali non è consentita la propagazione delle distribuzioni. È ortogonale ai tipi di fattori. ovvero, un fattore di propagazione bloccata può essere di qualsiasi tipo.

is_custom_rule indica se si tratta di una regola definita da un utente. Gli utenti possono definire regole di partizionamento per le chiamate personalizzate o sovrascrivere le regole di partizionamento predefinite per le operazioni standard. Una regola personalizzata viene sempre conservata/mai rimossa.

Vincoli:

• Il numero di mapping operandi/risultati deve corrispondere al numero di operandi/risultati dell'operazione.
• Esiste almeno una mappatura (non può esistere una regola per un'operazione senza operandi/risultati).
• Il rango di ogni TensorMappingAttr corrisponde al rango del tipo di tensore corrispondente.
• Per ogni gruppo di fattori (reduction_factors, need_replication_factors, permutation_factors):
  • Gli elementi devono essere compresi nell'intervallo [0, $factor_sizes].
  • Nessun indice di fattore duplicato all'interno di ogni gruppo e tra i gruppi.

Parametri:

PropagationEdgesAttr

Metadati edge di propagazione per tutti i passaggi di propagazione.

Sintassi:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

Un elenco di dettagli di propagazione per asse per un valore, raggruppati per indice di passaggio.

Parametri:

PropagationOneStepAttr

Metadati di propagazione per passaggio.

Sintassi:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

Dettagli della propagazione per tutti gli assi per un singolo passaggio di propagazione.

Parametri:

SubAxisInfoAttr

Informazioni su come questo asse secondario viene derivato dall'asse completo

Sintassi:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

Quando dividi un asse completo in n sotto-assi, l'asse viene rimodellato in [k_1,...,k_n] e l'i-esimo sotto-asse può essere espresso dal prodotto di tutte le dimensioni dell'asse alla sua sinistra m=prod(k_1,...,k_(i-1)) (ovvero la dimensione precedente) e dalla dimensione k_i. Pertanto, l'attributo sub-axis-info contiene questi due numeri ed è indicato come segue: (m)k per la pre-taglia m e la taglia k.

Vincoli:

• pre-size è almeno 1.
• size è maggiore di 1.
• pre-size deve dividere la dimensione dell'asse completo, ovvero sia pre-size sia size dividono la dimensione dell'asse completo e l'asse secondario non va oltre l'asse completo.
• La dimensione dell'asse secondario non è uguale a quella dell'asse completo corrispondente, nel qual caso è necessario utilizzare l'asse completo.

Parametri:

TensorMappingAttr

Mappature dei fattori per ogni dimensione di un tensore.

Sintassi:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

Vincoli:

• Gli elementi in dim_mappings devono soddisfare i vincoli in DimMappingAttr.
• Nessun indice di fattori duplicato nelle dimensioni.

Parametri:

TensorShardingAttr

Tensor sharding

Sintassi:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

Lo sharding dei tensori è associato a una mesh specifica e può fare riferimento solo ai nomi degli assi di quella mesh. Le suddivisioni delle dimensioni indicano per ogni dimensione del tensore lungo quali assi (o sotto-assi) viene suddiviso da maggiore a minore. Tutti gli altri assi che non suddividono una dimensione vengono replicati in modo implicito o esplicito (se vengono visualizzati nell'elenco degli assi replicati).

Tieni presente che nessun attributo di sharding su un tensore equivale a uno sharding del tensore completamente aperto.

La mesh a cui è associata questa suddivisione può essere specificata da un nome di simbolo, che fa riferimento a un simbolo MeshOp corrispondente, o da un MeshAttr incorporato.

Uno sharding può avere assi non ridotti (specificati da unreduced_axes), il che significa che il tensore non è ridotto lungo questi assi. Ad esempio, se la dimensione di contrazione di un'operazione matmul viene suddivisa lungo l'asse x sia nel lato sinistro che in quello destro, il risultato non viene ridotto lungo x. L'applicazione di un'operazione all-reduce sul tensore lungo gli assi non ridotti farà sì che il tensore venga replicato lungo questi assi. Tuttavia, un tensore con assi non ridotti non deve essere ridotto immediatamente, può rimanere non ridotto quando viene passato a operazioni lineari come stablehlo.add (purché sia lhs che rhs non siano ridotti) e ridotto in seguito. Supponiamo che il tipo di riduzione sia la somma, altre riduzioni potrebbero essere supportate in futuro.

Vincoli:

• Gli elementi in dim_shardings devono soddisfare i vincoli elencati in DimensionShardingAttr.
• Gli elementi in replicated_axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• Gli elementi in unreduced_axes devono soddisfare i vincoli elencati in AxisRefListAttr.
• Se il tipo di tensore corrispondente non è un ShapedType, lo sharding deve avere rango 0 e nessun asse replicato.
• Se si tratta di un ShapedType:
  • Il tensore deve avere un rango.
  • Il numero di shard di dimensioni è uguale al rango del tensore.
  • Le dimensioni di dimensione 0 non vengono suddivise.
• Non sono presenti riferimenti agli assi o assi secondari duplicati che si sovrappongono tra loro in dim_shardings, replicated_axes e unreduced_axes.
• Gli elementi in replicated_axes e unreduced_axes sono ordinati in base a mesh_or_ref (vedi AxisRefAttr::getMeshComparator).

Parametri:

TensorShardingPerValueAttr

Suddivisione dei tensori per operando/risultato di un'operazione

Sintassi:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

Un elenco di TensorShardingAttr, uno per ogni operando/risultato di un'operazione.

Vincoli:

• Gli elementi in shardings devono soddisfare i vincoli di TensorShardingAttr.

Parametri:

Enum

EdgeNodeType

Enumerazione del tipo di nodo edge

Custodie:

PropagationDirection

Enumerazione della direzione di propagazione

Custodie:

ReductionOp

Enum. riduzione

Custodie: