Dialetto 'sdy'

Il dialetto Shardy (SDY) definisce una rappresentazione di sharding dei tensori basata sull'asse e componenti API aggiuntivi per collegare gli sharding ai tensori.

Operazioni

sdy.constant (sdy::ConstantOp)

Operazione 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 blocco di X) e un insieme di destinazioni (ognuna è un risultato di X o un argomento del blocco di X), in modo che tutte le origini e le destinazioni debbano essere suddivise in modo uguale.

Un'operazione può avere più spigoli di 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
                  })

Questa operazione while ha n bordi di flusso di dati, il bordo di flusso di dati i-esimo è tra le origini x_i, return_value_i e le destinazioni y_i, pred_arg_i, body_arg_i.

Un elemento sdy.data_flow_edge prende come input la destinazione principale di un perimetro (può essere una qualsiasi delle destinazioni, ma preferibilmente un risultato dell'operazione piuttosto che un argomento blocco), che non dovrebbe avere altri usi. Questa operazione non è pura perché può accettare un input che inizialmente non aveva alcun utilizzo.

sdy.data_flow_edge include anche uno sharding facoltativo per tutte le destinazioni del perimetro e che durante la propagazione dovrebbe essere aggiornato lo sharding anziché lo sharding delle destinazioni (se possibile). Questo è utile quando un'operazione ha molti bordi, in quanto è molto più efficiente:

• si propagano separatamente in ogni bordo.
• aggiornare lo sharding di ogni perimetro separatamente invece che di tutte le destinazioni contemporaneamente (ad es. un'operazione ha un singolo TensorShardingPerValueAttr immutabile per la suddivisione dei risultati).
• Aggiungi ogni bordo alla lista di lavoro separatamente quando lo sharding di un'origine è cambiato.

La propagazione propagherà gli shard 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 avviene dalle origini alle destinazioni e che la propagazione all'indietro avviene dalle destinazioni alle origini.

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

Tratti: SameOperandsAndResultType

Interfacce: InferTypeOpInterface

Attributi:

Operandi:

Risultati:

sdy.manual_computation (sdy::ManualComputationOp)

Operazione di parallelismo multi-dispositivo con i 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)

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

Il corpo è locale rispetto agli assi manual_axes. La propagazione avverrà tramite il corpo su eventuali assi liberi, ovvero quelli non presenti nell'elenco manual_axes.

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

Attributi:

Operandi:

Risultati:

sdy.mesh (sdy::MeshOp)

Mesh con nome

Sintassi:

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

Definisce un nuovo mesh denominato. Tutti i mesh in un modulo devono avere lo stesso numero di dispositivi (tranne i mesh con un singolo device_id). La mesh è un'operazione Symbol visualizzata nel SymbolTable del modulo e a cui è possibile fare riferimento tramite il relativo name.

Tratti: HasParent<ModuleOp>

Interfacce: Symbol

Attributi:

sdy.named_computation (sdy::NamedComputationOp)

Operazione di calcolo con nome

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 scorrerà dentro e fuori la regione come se tutto fosse allineato.

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

Il tipo di ogni argomento del blocco e dei valori restituiti nella regione deve essere uguale al tipo degli operandi e al tipo di risultati 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, ShardableDataFlowOpInterface

Attributi:

Operandi:

Risultati:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Operazione di 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 ricevuto come input. Tuttavia, in termini di propagazione, consentirà la propagazione solo in una determinata direzione.

In questo modo, gli sharding non vengono propagati tra gli utilizzi del risultato dell'operazione di barriera e del relativo operando.

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

Tratti: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Interfacce: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effetti: MemoryEffects::Effect{}

Attributi:

Operandi:

Risultati:

sdy.reshard (sdy::ReshardOp)

Esegui il ridimensionamento di un tensore su uno sharding diverso

Sintassi:

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

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

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

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

// TODO(b/331680067). Aggiungi un pattern di canonizzazione per rimuovere le operazioni // reshard ridondanti.

Tratti: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Interfacce: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Effetti: MemoryEffects::Effect{}

Attributi:

Operandi:

Risultati:

sdy.return (sdy::ReturnOp)

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

Sintassi:

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

Caratteristiche: AlwaysSpeculatableImplTrait, Terminator

Interfacce: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effetti: MemoryEffects::Effect{}

Operandi:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Limita 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 matriciale) per indicare in che modo deve essere suddiviso il tensore o un sottoinsieme dei relativi utilizzi.

Se lo sharding ha dimensioni aperte e assi non vincolati, significa che il tensore può essere ulteriormente suddiviso in base alle dimensioni aperte.

Questa operazione può:

• Non hanno utilizzi (non collegati), il che significa che lo sharding allegato è il modo in cui deve essere suddiviso il tensore di input stesso.
• Hanno utilizzi, il che significa che lo sharding associato è il modo in cui gli utilizzi dell'operazione del vincolo di partizionamento devono essere sottoposti a sharding, 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 uso).

Caratteristiche: Elementwise, SameOperandsAndResultType

Interfacce: InferTypeOpInterface

Attributi:

Operandi:

Risultati:

sdy.sharding_group (sdy::ShardingGroupOp)

Operazione di gruppo per lo sharding

Sintassi:

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

Questa operazione fornisce un'interfaccia per assegnare tensori ai gruppi di partizionamento orizzontale (gruppi di tensori che verranno imposti per avere uno sharding identico). Durante la propagazione, non appena un elemento di gruppo viene suddiviso in parti, tutti gli altri membri vengono suddivisi nello stesso modo. Questa operazione prende l'ID gruppo dell'argomento e non restituisce alcun risultato, ma modifica la rappresentazione interna del gruppo di suddivisione per aggiungere il tensore di input al gruppo con l'ID specificato.

Attributi:

Operandi:

Attributi

AxisRefAttr

Riferimento a un asse completo o a un asse secondario suddiviso

Sintassi:

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

Parametri:

DimMappingAttr

Elenco di indici di fattori per una dimensione

Tutti gli indici dei fattori devono rientrare nell'intervallo [0, num_factors) e un elenco vuoto indica che si tratta di una mappatura nulla (viene analizzata/stampata con *), ovvero la dimensione non è mappata a nessun fattore.

Parametri:

DimensionShardingAttr

Sharding delle dimensioni

Elenco di nomi degli assi in base ai quali eseguire lo sharding di una dimensione del tensore dal maggiore al minore, un valore booleano che indica se la dimensione può essere ulteriormente suddivisa e un valore intero facoltativo che indica la priorità di questo sharding della dimensione, che verrà rispettata durante la propagazione dello sharding. Le priorità provengono dalle annotazioni di sharding degli utenti e un valore più basso indica una priorità più elevata. Se la priorità non è presente nell'annotazione, viene assunta la priorità più alta.

Parametri:

ManualAxesAttr

Sintassi:

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

Parametri:

MeshAttr

Mesh di assi e 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 specificano l'ordine dei dispositivi.

Se l'elenco degli assi è vuoto, la mesh ha un asse implicito senza nome di dimensione 1. In questo caso, se non viene fornito un elenco di ID dispositivo, l'elenco di ID dispositivo implicito è [0]; se viene fornito un elenco di ID dispositivo, deve contenere un singolo numero intero di qualsiasi valore non negativo. Questo è il caso di sharding massimo.

Per tutti i casi di suddivisione non massima, se viene specificato un elenco di ID dispositivo, il prodotto delle dimensioni dell'asse deve corrispondere al numero di dispositivi. Se non viene specificato un elenco di ID dispositivo, l'elenco di ID dispositivo implicito è iota(product(axes)). Per semplicità, non è consentita nemmeno la specifica di un elenco di ID dispositivo uguale a iota(product(axes)); in questo caso, non deve essere specificato un elenco di ID dispositivo.

Ecco alcuni esempi di maglie:

• Un mesh vuoto rappresenta un mesh segnaposto che può essere sostituito durante la propagazione: <[]>
• Una mesh con un asse senza nome e un ID dispositivo esplicito, che in genere viene utilizzato per rappresentare lo 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]>

Parametri:

MeshAxisAttr

Asse denominato in una mesh

Sintassi:

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

Parametri:

OpShardingRuleAttr

Specifica come è possibile partizionare un'operazione.

Sintassi:

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

Una regola di suddivisione in parti specifica in che modo un'operazione può essere suddivisa 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 parti, principalmente per completezza, in quanto molte operazioni, come le operazioni punto per punto, hanno dimensioni di dimensione 1 che corrispondono a operandi e risultati.

is_custom_rule descrive se si tratta di una regola definita da un utente per un'operazione stablehlo.custom_call. Il partizionatore non sa come suddividere queste operazioni, quindi è necessario che un utente gli indichi come. Quando si tratta di una regola personalizzata, la regola viene sempre conservata/mai rimossa. is_custom_rule può essere true solo per le operazioni stablehlo.custom_call.

Parametri:

SubAxisInfoAttr

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

Sintassi:

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

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

Parametri:

TensorMappingAttr

Mappature dei fattori per ogni dimensione di un tensore.

Sintassi:

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

Parametri:

TensorShardingAttr

Sharding dei tensori

Sintassi:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

La suddivisione in blocchi di un tensore è associata a una mesh specifica e può fare riferimento solo ai nomi degli assi di quella mesh. I partizionamenti delle dimensioni ci dicono per ogni dimensione del tensore lungo quali assi (o assi secondari) è suddiviso da principale a secondario. Tutti gli altri assi che non eseguono lo shard di una dimensione vengono replicati implicitamente o esplicitamente (se compaiono nell'elenco degli assi replicati).

La mesh a cui è associato questo suddivisione può essere specificata tramite un nome simbolo, facendo riferimento a un simbolo MeshOp corrispondente o a un MeshAttr in linea.

Parametri:

TensorShardingPerValueAttr

Sharding dei tensori per operando/risultato di un'operazione

Sintassi:

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

Parametri:

Enum

PropagationDirection

enum direzione di propagazione

Custodie: