Dialekt „sdy“

Der Shardy-Dialekt (SDY)

Der Shardy-Dialekt (SDY) definiert eine achsenbasierte Tensor-Sharding-Darstellung und zusätzliche API-Komponenten zum Anhängen von Shardings an Tensoren.

Versionslog: 0.0.1: Unreduced axes to TensorShardingAttr hinzufügen.

Vorgänge

sdy.all_gather (sdy::AllGatherOp)

Führt eine All-Gather-Kommunikation entlang von Achsen durch

Syntax:

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

Sammelt Chunks eines Tensors entlang der in gathering_axes angegebenen Achsen.

gathering_axes ist eine Liste von Achsenlisten. Die äußere Liste überschreitet die Dimensionen des Tensors. Jede innere Liste gibt die Achsen an, entlang derer ein separater Gather für die jeweilige Dimension ausgeführt werden soll. Sie wird auf die Shardierung des Operanden (tensor) angewendet, um die Shardierung des Ergebnisses (out_sharding) zu erhalten.

out_sharding wird nicht verwendet, um das Sharding des Ergebnisses zu bestimmen. Stattdessen wird die Aufteilung des Ergebnisses durch die Aufteilung des Operanden und gathering_axes bestimmt. out_sharding muss mit dieser abgeleiteten Aufteilung übereinstimmen.

Beispiel:

%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>

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• Elemente in gathering_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn Sie gathering_axes auf das Sharding des Operanden anwenden, erhalten Sie out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.all_reduce (sdy::AllReduceOp)

All-Reduce-Kommunikation entlang von Achsen durchführen

Syntax:

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

Reduziert Chunks eines Tensors entlang der in reduction_axes angegebenen Achsen. Die Reihenfolge von reduction_axes ist für das Ergebnis nicht wichtig, kann aber die Reihenfolge der entsprechenden Replikagruppen beeinflussen.

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• reduction_axes muss den in AxisRefListAttr aufgeführten Einschränkungen entsprechen.
• reduction_axes muss in Bezug auf das Mesh sortiert werden.
• Das Sharding des Operanden und out_sharding müssen ein entsprechendes Dimensions-Sharding haben.
• reduction_axes darf sich nicht mit dem Sharding der Operanden-Dimension und den replizierten Achsen überschneiden, kann sich aber mit nicht reduzierten Achsen überschneiden.
• reduction_axes darf sich nicht mit den nicht reduzierten Achsen von out_sharding überschneiden. Mit anderen Worten: out_sharding muss entlang von reduction_axes repliziert werden (implizit oder explizit).

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.all_slice (sdy::AllSliceOp)

Führt einen dynamischen Slice-Vorgang entlang von Achsen aus

Syntax:

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

Teilt einen Tensor entlang der in slicing_axes angegebenen Achsen in Abschnitte. Es besteht eine algebraische Dualität zwischen sdy.all_slice und sdy.all_gather.

slicing_axes ist eine Liste von Achsenlisten. Die äußere Liste überschreitet die Dimensionen des Tensors. Jede innere Liste gibt die Achsen an, entlang derer ein Slice für die jeweilige Dimension ausgeführt werden soll. Sie wird auf die Aufteilung des Operanden (tensor) angewendet, um die Aufteilung des Ergebnisses (out_sharding) zu erhalten.

out_sharding wird nicht verwendet, um das Sharding des Ergebnisses zu bestimmen. Stattdessen wird die Aufteilung des Ergebnisses durch die Aufteilung des Operanden und slicing_axes bestimmt. out_sharding muss mit dieser abgeleiteten Aufteilung übereinstimmen.

Beispiel:

%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>

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• Elemente in slicing_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn Sie slicing_axes auf das Sharding des Operanden anwenden, erhalten Sie out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.all_to_all (sdy::AllToAllOp)

Führt eine All-to-All-Kommunikation entlang von Achsen durch

Syntax:

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

Für jedes Tupel (axes, src_dim, tgt_dim) in der Parameterliste werden mit dieser Operation Teile eines Tensors entlang der Dimension tgt_dim und der in axes angegebenen Achsen aufgeteilt, entlang der Achsen verteilt und entlang der Dimension src_dim verkettet.

Dieser Vorgang ist im Wesentlichen eine Kombination aus einem „all-gather“ entlang src_dim und axes, gefolgt von einem „all-slice“ entlang tgt_dim und axes. Das bedeutet, dass ein Suffix der Achsen-Sharding-Dimension src_dim im Eingabetensor an die Achsen-Sharding-Dimension tgt_dim im Ausgabetensor angehängt wird.

Der All-to-All-Vorgang wird auf die Shardierung des Operanden (tensor) angewendet, um die Shardierung des Ergebnisses (out_sharding) zu erhalten.

out_sharding wird nicht verwendet, um das Sharding des Ergebnisses zu bestimmen. Stattdessen wird die Aufteilung des Ergebnisses durch die Aufteilung des Operanden src_dim, tgt_dim und axes bestimmt und out_sharding muss mit dieser abgeleiteten Aufteilung übereinstimmen.

Beispiel:

%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>

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• Die Parameterliste darf nicht leer sein.
• Für jeden Parameter in params:
  • Elemente in axes müssen die Einschränkungen von AxisRefAttr erfüllen.
  • src_dim und tgt_dim müssen gültige Dimensionen sein (nicht negativ und kleiner als der Rang des Tensors).
  • Alle src_dim oder tgt_dim müssen für alle Parameter eindeutig sein.
  • src_dim muss für alle Parameter in aufsteigender Reihenfolge sortiert sein.
• Wenn Sie axes von src_dim nach tgt_dim im Operanden-Sharding verschieben, erhalten Sie out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.collective_permute (sdy::CollectivePermuteOp)

Führt eine kollektive Permutationskommunikation durch, um Achsen zu ersetzen

Syntax:

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

Sendet einen Teil des Eingabetensors von jedem Gerät an ein anderes, um die Achsen neu anzuordnen/zu ersetzen, die den Tensor aufteilen.

Bei einer kollektiven Permutation kann das Eingabe-Sharding so transformiert werden, dass jede Dimension so stark wie zuvor geshardet werden muss. Das heißt, sie muss entlang von Achsen geshardet werden, deren Produkt der Größen dem der Achsen entspricht, die den Tensor zuvor geshardet haben.

Das ist nützlich, um Achsen in einer einzelnen Dimension oder über verschiedene Dimensionen hinweg neu anzuordnen und um Shard-Achsen durch replizierte Achsen zu ersetzen.

Im folgenden Beispiel ist die Größe des aufgeteilten Tensors tensor<1x4x2xf32>. Diese Größe wird durch die kollektive Permutation beibehalten.

Beispiel:

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>

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• Wenn sich die Meshs für die Eingabe- und Ausgabeverteilung unterscheiden, müssen sie genau dieselben Achsen und eine unterschiedliche Reihenfolge der Geräte-IDs haben.
• Für jede Dimension muss das Produkt der Größen der Sharding-Achsen in out_sharding mit dem des entsprechenden Operanden-Dimensionsshardings übereinstimmen.

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.constant (sdy::ConstantOp)

Konstanter Vorgang

Erstellt einen output-Tensor aus einer konstanten value.

Weitere Informationen finden Sie unter https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant.

Beispiel:

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

Merkmale: AlwaysSpeculatableImplTrait

Schnittstellen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Ergebnisse:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Datenfluss-Edge-Vorgang

Syntax:

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

Eine Datenflusskante eines bestimmten Vorgangs X definiert eine Brücke zwischen einer Reihe von Quellen (jede ist entweder ein Operand von X oder ein Operand des Blockterminators von X) und einer Reihe von Zielen (jedes ist entweder ein Ergebnis von X oder ein Blockargument von X), sodass alle Quellen und Ziele auf dieselbe Weise partitioniert werden sollten.

Ein Vorgang kann mehrere Datenflusskanten haben, die orthogonal zueinander sind.

Beispiel:

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

Dieser While-Vorgang hat „n“ Datenflusskanten. Die i-te Datenflusskante verläuft zwischen den Quellen x_i, return_value_i und den Zielen y_i, pred_arg_i, body_arg_i.

Eine sdy.data_flow_edge nimmt den Inhaber einer Kante als Eingabe entgegen. Das kann eines der Ziele sein, vorzugsweise aber ein Vorgangsergebnis und kein Blockargument. Die Eingabe sollte keine anderen Verwendungen haben. Dieser Vorgang ist nicht rein, da er eine Eingabe annehmen kann, die ursprünglich keine Verwendungen hatte.

Die sdy.data_flow_edge enthält auch ein optionales Sharding für alle Ziele des Edge. Dieses Sharding sollte während der Weitergabe anstelle des Shardings der Ziele aktualisiert werden (sofern es angehängt werden kann). Das ist nützlich, wenn ein Vorgang viele Kanten hat, da es viel effizienter ist,

• separat über jede Kante weitergegeben.
• Die Shardierung jeder Kante wird separat aktualisiert, anstatt alle Ziele gleichzeitig zu aktualisieren. Beispiel: Ein Vorgang hat eine einzelne unveränderliche TensorShardingPerValueAttr für die Shardierung von Ergebnissen.
• Fügen Sie jede Kante der Arbeitsliste separat hinzu, wenn sich das Sharding einer Quelle geändert hat.

Bei der Weitergabe werden Shardings zwischen allen Quellen und Zielen eines sdy.data_flow_edge so weitergegeben, als wäre es ein regulärer Vorgang mit den Quellen als Operanden und den Zielen als Ergebnissen sowie einer Identitäts-sdy.op_sharding_rule. Das bedeutet, dass die Vorwärtsweitergabe von Quellen zu Zielen und die Rückwärtsweitergabe von Zielen zu Quellen erfolgt.

Die Eingabe von sdy.data_flow_edge darf nicht durch einen SdyDialect-Vorgang definiert werden. Wir können also davon ausgehen, dass sie durch einen Vorgang mit einem nicht registrierten sdy.sharding-Attribut definiert wird.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

Func input/output data flow edge op.

Syntax:

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

Ein Datenfluss-Edge-Vorgang, aber für Funktionsargumente oder Aufrufergebnisse. Wenn der Operand ein BlockArgument ist, wird das Argument des callOp des Aufrufers an die Nutzer des func-Arguments übergeben. Für jedes „func“-Argument gibt es eine „func“-Datenflusskante. Wenn der Operand ein OpResult ist, wird der Rückgabewert des aufgerufenen funcOp an die Nutzer des Aufrufergebnisses weitergeleitet. Für jedes Aufrufergebnis gibt es eine Kante für den Funktionsdatenfluss.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, SymbolUserOpInterface

Operanden:

Ergebnisse:

sdy.manual_computation (sdy::ManualComputationOp)

Parallelbetrieb für verschiedene Geräte mit manuellen Kollektiven

Syntax:

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)

Wechseln Sie zu einem Bereich, der in Bezug auf den lokalen Code pro Gerät mit expliziten Kollektiven geschrieben ist, wobei logische Formen mit lokalen physischen Pufferformen pro Gerät übereinstimmen und Kollektive genau der physischen geräteübergreifenden Kommunikation entsprechen.

Der Körper ist lokal in Bezug auf die manual_axes. Die Ausbreitung erfolgt über den Body auf allen kostenlosen Achsen, die nicht in der Liste „manual_axes“ enthalten sind.

Nicht gerankte Tensoren müssen eine Shardierung mit Rang 0 haben, d.h. vollständig repliziert sein.

Einschränkungen:

• Elemente in in_shardings und out_shardings müssen die in TensorShardingAttr aufgeführten Einschränkungen erfüllen.
• Die Anzahl der globalen und lokalen Tensoreingaben/-ausgaben der Op-Region muss übereinstimmen.
• Die manuellen Achsen müssen in jedem Dimensions-Sharding vor allen kostenlosen Achsen stehen.
• Bei manuellen Achsen kann es nicht zu Auffüllungen kommen. Die Größe der Dimension muss durch die Größe der entsprechenden manuellen Achse teilbar sein.
• Die globalen und lokalen Formen der Argumente/Ergebnisse der Vorgangsregionen müssen übereinstimmen.

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

Schnittstellen: ShardableDataFlowOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.mesh (sdy::MeshOp)

Benanntes Mesh

Syntax:

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

Definiert ein neues benanntes Mesh. Alle Meshs in einem Modul müssen dieselbe Anzahl von Geräten haben (mit Ausnahme von Meshs mit einer einzelnen device_id). Das Mesh ist ein Symbol-Vorgang, der im SymbolTable des Moduls angezeigt wird und über seinen name referenziert werden kann.

Merkmale: HasParent<ModuleOp>, SymbolName

Schnittstellen: Symbol

Attribute:

sdy.named_computation (sdy::NamedComputationOp)

Benannter Rechenvorgang

Syntax:

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)

Gruppiert eine Berechnung, d.h. einen Block von Vorgängen, und gibt ihr einen Namen. Die Weitergabe erfolgt in die Region hinein und aus der Region heraus, als ob alles inline wäre.

Damit kann die Weitergabe von Aufrufanweisungen an andere Funktionen gesteuert werden. Alle Shardy-Nutzer sollten einen Import-/Export-Pass schreiben, der ihre Call-Vorgänge in sdy.named_computation-Vorgänge umwandelt und den Hauptteil der aufgerufenen Funktion in den Hauptteil von named_computation dupliziert bzw. kopiert.

Der Typ der einzelnen Blockargumente und zurückgegebenen Werte in der Region muss mit dem Typ der Operanden und dem Ergebnistyp des Vorgangs übereinstimmen.

Beispiel:

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

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

Schnittstellen: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Betrieb der Propagationsbarriere

Syntax:

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

Diese Operation funktioniert wie eine Identitätsoperation und gibt denselben Wert aus, der als Eingabe verwendet wurde. In Bezug auf die Ausbreitung wird die Ausbreitung jedoch nur in einer bestimmten Richtung ermöglicht.

Dadurch wird verhindert, dass Shardings zwischen den Verwendungen des Ergebnisses des Barrier-Vorgangs und seines Operanden weitergegeben werden.

• FORWARD bedeutet, dass Shardings nur vom Operanden zum Ergebnis fließen können.
• BACKWARD bedeutet, dass Shardings nur vom Ergebnis zum Operanden fließen können.
• NONE bedeutet, dass kein Sharding durch diesen Vorgang weitergegeben werden kann.
• BOTH kann nicht angegeben werden, da dieser Vorgang redundant wäre.

Merkmale: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Schnittstellen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Operanden:

Ergebnisse:

sdy.reduce_scatter (sdy::ReduceScatterOp)

Führt eine Reduce-Scatter-Kommunikation entlang von Achsen durch

Syntax:

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

Reduziert Chunks eines Tensors entlang der in reduce_scatter_axes angegebenen Achsen und verteilt das Ergebnis dann entlang derselben Achsen. Dieser Vorgang ist im Grunde eine Kombination aus einem sdy.all_reduce und einem sdy.all_slice entlang derselben reduce_scatter_axes.

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• Elemente in reduce_scatter_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn Sie reduce_scatter_axes auf das Operand-Sharding anwenden, erhalten Sie out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Implizit oder explizit replizierte Achsen in nicht reduzierte Achsen verschieben:

Syntax:

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

axes sollte implizit oder explizit im Operanden repliziert werden. Durch diesen Vorgang werden sie im Ergebnis nicht reduziert. Wir haben die folgende Beziehung:

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

Beispiel:

%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>

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• axes muss den in AxisRefListAttr aufgeführten Einschränkungen entsprechen.
• axes muss in Bezug auf das Mesh sortiert werden.
• axes sind nicht leer.
• Die Eingabe- und Ausgabeshards müssen dieselben Dimensionshards haben.
• axes muss implizit oder explizit im Operand-Sharding repliziert werden.
• inUnreducedAxes + axes = outUnreducedAxes.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.reshard (sdy::ReshardOp)

Teilt einen Tensor in ein anderes Sharding auf.

Syntax:

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

Teilt den Eingabetensor mit dem angegebenen Sharding neu auf. Dieses Sharding unterscheidet sich vom vorhandenen Sharding des Eingabetensors.

Sowohl ShardingConstraintOp als auch ReshardOp weisen einem Tensor eine Shardierung zu. Die Lebensdauer ist:

1. Vor der Shard-Weitergabe wird ShardingConstraintOp von Nutzern hinzugefügt.
2. Die Weitergabe der Fragmentierung verwendet ShardingConstraintOp. In den Ergebnissen der Sharding-Weitergabe ist kein ShardingConstraintOp vorhanden. Stattdessen kann bei Bedarf ReshardOp hinzugefügt werden.
3. Ein Partitioner wandelt einen ReshardOp in einen kollektiven Vorgang (oder einen Identitätsvorgang) um. In den Ergebnissen des Partitionierers darf es keine ReshardOp geben.

Merkmale: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Operanden:

Ergebnisse:

sdy.return (sdy::ReturnOp)

Durch den Vorgang sdy.return werden die Regionen beendet, die an sdy-region-based ops und alle anderen Shardy-region-based ops angehängt sind. Sie ist variadisch: Sie akzeptiert als Argumente eine Liste von Werten, deren Typen beliebig sein können (aber vom selben Typ, z. B. AnyTensor). Daher kann sie auf verschiedenen Ebenen des Shardy IR-Stacks wiederverwendet werden.

Syntax:

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

Merkmale: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

Schnittstellen: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Auswirkungen: MemoryEffects::Effect{}

Operanden:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

Verschieben einiger aufgeteilter Achsen des Operanden zu nicht reduzierten Achsen des Ergebnisses

Syntax:

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

axes sollte verwendet werden, um den Operanden zu sharden. Durch diesen Vorgang werden sie im Ergebnis nicht reduziert. Wir haben die folgende Beziehung:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), wobei all-gather, sharded-to-unreduced und all-reduce auf dieselben Achsen angewendet werden.

Beispiel:

%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>

Einschränkungen:

• Muss den in Sdy_CollectiveOpInterface aufgeführten Einschränkungen entsprechen.
• Elemente in axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn Sie axes auf das Sharding des Operanden anwenden, erhalten Sie out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Beschränkt einen Tensor auf das angegebene Sharding

Syntax:

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

Weist einem Zwischen-Tensor (z.B. dem Ergebnis einer Matmul-Operation) eine Shardierung zu, um anzugeben, wie dieser Tensor oder eine Teilmenge seiner Verwendungen shardiert werden soll.

Wenn das Sharding offene Dimensionen und nicht eingeschränkte Achsen hat, kann der Tensor entlang der offenen Dimensionen weiter aufgeteilt werden.

Dieser Vorgang kann entweder:

• Keine Verwendungen (dangling): Das bedeutet, dass das angehängte Sharding angibt, wie der Eingabetensor selbst partitioniert werden soll.
• Verwendungen haben. Das bedeutet, dass die angehängte Partitionierung angibt, wie die Verwendungen des Sharding-Constraint-Vorgangs partitioniert werden sollen. Andere Verwendungen des Eingabetensors können eine andere Partitionierung haben, wenn der Eingabetensor keine anderen Verwendungen hat. In diesem Fall ist das Verhalten dasselbe wie im Fall ohne Verwendungen.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, SymbolUserOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.sharding_group (sdy::ShardingGroupOp)

Schränkt Tensoren in der Gruppe so ein, dass sie dasselbe Sharding haben.

Syntax:

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

Dieser Vorgang bietet eine Schnittstelle zum Zuweisen von Tensoren zu Sharding-Gruppen (Gruppen von Tensoren, für die identische Shardings erzwungen werden). Sobald ein Gruppenelement während der Weitergabe fragmentiert wird, werden alle anderen Elemente auf genau dieselbe Weise fragmentiert. Bei diesem Vorgang wird die Argumentgruppen-ID verwendet und es wird kein Ergebnis zurückgegeben. Stattdessen wird die interne Sharding-Gruppendarstellung geändert, um den Eingabetensor der Gruppe mit der angegebenen ID hinzuzufügen.

Schnittstellen: InferTypeOpInterface

Attribute:

Operanden:

Attribute

AllToAllParamAttr

All-to-all-Parameter

Syntax:

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

Ein Tupel mit den Achsen und Quell-/Zieldimensionen für die All-to-All-Operation.

Parameter:

AllToAllParamListAttr

Liste aller All-to-All-Parameter

Syntax:

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

Parameter:

AxisRefAttr

Verweis auf eine vollständige Achse oder eine untergeordnete Achse

Syntax:

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

Einschränkungen:

• name muss im gebundenen MeshAttr vorhanden sein.
• Wenn sub_axis_info vorhanden ist, muss es den Einschränkungen von SubAxisInfoAttr entsprechen.

Parameter:

AxisRefListAttr

Liste der Achsenreferenzen

Syntax:

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

Einschränkungen:

• Elemente in value müssen die Einschränkungen von AxisRefAttr erfüllen.
• Es gibt keine doppelten Achsenreferenzen oder untergeordneten Achsen, die sich überschneiden.
• Keine zwei benachbarten Achsenreferenzen sind aufeinanderfolgende untergeordnete Achsen derselben vollständigen Achse. Sie können also in einer untergeordneten Achse oder der vollständigen Achse zusammengeführt werden.

Parameter:

AxisToPropagationDetailsAttr

Details zum Weiterleitungs-Edge-Flow für eine bestimmte Achse und Quelle.

Syntax:

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

Ordnet eine Quellwertreferenz einer Liste von Zielwertreferenzen entlang einer bestimmten Achse zu.

Parameter:

DimMappingAttr

Liste der Faktorindexe für eine Dimension

Eine leere Liste gibt an, dass es sich um eine Nullzuordnung handelt (diese wird mit * geparst/ausgegeben), d.h., die Dimension ist keinen Faktoren zugeordnet.

Einschränkungen:

• Es gibt mindestens einen Faktorindex.
• Faktorindexe müssen im Bereich [0, $factor_sizes) liegen.
• Wenn es mehrere Faktoren gibt, darf keiner die Größe 1 haben.
• Keine doppelten Faktorindexe.

Parameter:

DimensionShardingAttr

Dimensionen aufteilen

Liste der Achsennamen, nach denen eine Tensordimension von primär nach sekundär aufgeteilt werden soll, ein boolescher Wert, der angibt, ob die Dimension weiter aufgeteilt werden kann, und eine optionale Ganzzahl, die die Priorität dieser Dimensionsaufteilung angibt, die bei der Weitergabe der Aufteilung berücksichtigt wird. Prioritäten stammen aus Anmerkungen zum Sharding von Nutzern. Ein niedrigerer Wert bedeutet eine höhere Priorität. Wenn die Priorität in der Anmerkung fehlt, wird die höchste Priorität angenommen.

Einschränkungen:

• Elemente in axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn ein Dimension-Sharding eine Priorität hat:
  • Die Priorität ist größer oder gleich 0.
  • Die Dimension hat mindestens eine Achse, wenn sie geschlossen ist.

Parameter:

EdgeValueRefAttr

Verweis auf einen bestimmten Index einer Wertkante vom Typ type.

Syntax:

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

Parameter:

ListOfAxisRefListsAttr

Liste der Achsenreferenzlisten

Syntax:

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

Parameter:

ManualAxesAttr

Eine Liste der Achsen, für die ein ManualComputationOp manuell ist

Syntax:

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

Parameter:

MeshAttr

Achsennetz und Geräteliste

Syntax:

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

Ein Mesh ist eine Liste von Achsen und eine optionale Liste von Geräte-IDs, die die Geräteanordnung angeben.

Wenn die Liste der Achsen leer ist

• Wenn device_ids nicht angegeben ist, ist es ein leeres Mesh.
• Wenn device_ids angegeben wird, muss es sich um eine einzelne nicht negative Ganzzahl handeln. Wir nennen sie maximales Sharding-Mesh.

Wenn die Liste der Achsen angegeben wird

• Wenn eine Geräte‑ID-Liste angegeben wird, muss das Produkt der Achsengrößen mit der Anzahl der Geräte übereinstimmen.
• Wenn keine Geräte-ID-Liste angegeben ist, ist die implizite Geräte-ID-Liste iota(product(axes)). Aus Gründen der Einfachheit ist es auch nicht zulässig, eine Geräte-ID-Liste anzugeben, die mit iota(product(axes)) identisch ist. In diesem Fall sollte keine Geräte-ID-Liste angegeben werden.
• Es handelt sich nicht um ein Mesh mit maximalem Sharding, auch wenn die Gesamtgröße der Achsen 1 beträgt.

Hier einige Beispiele für Meshes:

• Ein leeres Mesh ist ein Platzhalter-Mesh, das während der Weitergabe ersetzt werden kann: <[]>
• Ein Mesh ohne Achsenliste und mit einer einzelnen nicht negativen Geräte-ID, die ein maximales Sharding-Mesh ist: <[], device_ids=[3]>
• Ein Mesh mit zwei Achsen und impliziten Geräte-IDs iota(6): <["a"=2, "b"=3]>
• Ein Mesh mit zwei Achsen und expliziten Geräte-IDs, die die Reihenfolge der Geräte angeben: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Einschränkungen:

• Die Elemente in device_ids dürfen nicht negativ sein.
• Wenn axes leer ist, kann die Größe von device_ids 0 (leeres Mesh) oder 1 (maximales Sharding-Mesh) sein.
• Wenn axes nicht leer ist,
  • Elemente in axes dürfen keine doppelten Namen haben.
  • Wenn device_ids angegeben ist, ist das ursprüngliche device_ids nicht iota(product(axis_sizes)) und das sortierte device_ids ist iota(product(axis_sizes)).

Parameter:

MeshAxisAttr

Benannte Achse in einem Mesh

Syntax:

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

Parameter:

OpShardingRuleAttr

Gibt an, wie ein Vorgang partitioniert werden kann.

Syntax:

#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
>

Eine Sharding-Regel gibt an, wie ein Vorgang anhand verschiedener Eigenschaften des Vorgangs partitioniert werden kann, z. B. Attribute, Form von Operanden, Form von Ergebnissen usw. Beispiel:

%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>

Faktoren mit der Größe 1 sind zulässig, auch wenn sie nicht aufgeteilt werden können. Das ist hauptsächlich der Vollständigkeit geschuldet, da viele Operationen wie punktweise Operationen Dimensionen der Größe 1 haben, die Operanden und Ergebnissen entsprechen.

Faktortypen:

• reduction_factors enthält die Indexe der Faktoren, die reduziert werden müssen, z. B. die kontrahierenden Dimensionen in einer Punktoperation. Diese Faktoren können in Operanden, aber nicht in Ergebnissen enthalten sein.
• need_replication_factors enthält die Indexe der Faktoren, die eine vollständige Replikation erfordern, z. B. die sortierte Dimension in einem Sortiervorgang.
• permutation_factors enthält die Indexe der Faktoren, für die „collective-permute“ erforderlich ist, wenn sie aufgeteilt sind, z. B. die Padding-Dimensionen in einem Pad-Vorgang.
• Alle anderen Faktoren werden als Pass-through-Faktoren betrachtet, d.h. Faktoren, für die keine Kommunikation erforderlich ist, wenn sie auf dieselbe Weise auf alle Tensoren verteilt werden, die ihnen zugeordnet sind.

blocked_propagation_factors enthält die Faktoren, für die keine Shardings weitergegeben werden dürfen. Sie ist orthogonal zu den Faktortypen. Ein Faktor für blockierte Weitergabe kann ein beliebiger Faktortyp sein.

is_custom_rule gibt an, ob es sich um eine vom Nutzer definierte Regel handelt. Nutzer können Sharding-Regeln für ihre benutzerdefinierten Aufrufe definieren oder die vordefinierten Sharding-Regeln für die Standardvorgänge überschreiben. Eine benutzerdefinierte Regel wird immer beibehalten und nie entfernt.

Einschränkungen:

• Die Anzahl der Operand-/Ergebniszuordnungen muss mit der Anzahl der Operanden/Ergebnisse des Vorgangs übereinstimmen.
• Es gibt mindestens eine Zuordnung (es kann keine Regel für einen Vorgang ohne Operanden/Ergebnisse geben).
• Der Rang jedes TensorMappingAttr entspricht dem Rang des entsprechenden Tensortyps.
• Für jede Gruppe von Faktoren (reduction_factors, need_replication_factors, permutation_factors):
  • Die Elemente müssen im Bereich [0, $factor_sizes] liegen.
  • Es gibt keine doppelten Faktorindexe innerhalb der einzelnen Gruppen und gruppenübergreifend.

Parameter:

PropagationEdgesAttr

Metadaten für die Weiterleitungskanten für alle Weiterleitungsschritte.

Syntax:

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

Eine Liste mit Details zur Weitergabe pro Achse für einen Wert, gruppiert nach Schrittindex.

Parameter:

PropagationOneStepAttr

Metadaten für die Weitergabe pro Schritt

Syntax:

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

Details zur Weitergabe für alle Achsen für einen einzelnen Weitergabeschritt.

Parameter:

SubAxisInfoAttr

Informationen dazu, wie diese untergeordnete Achse aus der vollständigen Achse abgeleitet wird

Syntax:

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

Wenn eine vollständige Achse in n Unterachsen aufgeteilt wird, wird die Achse in [k_1,...,k_n] umgeformt. Die i-te Unterachse kann durch das Produkt aller Achsengrößen links davon m=prod(k_1,...,k_(i-1)) (auch als „Pre-Size“ bezeichnet) und die Größe k_i ausgedrückt werden. Daher enthält das Attribut „sub-axis-info“ diese beiden Zahlen und wird so angegeben: (m)k für die Vorabgröße m und die Größe k.

Einschränkungen:

• pre-size ist mindestens 1.
• size ist größer als 1.
• pre-size muss die Größe der vollständigen Achse teilen, d.h. sowohl pre-size als auch size teilen die Größe der vollständigen Achse und die untergeordnete Achse geht nicht über die vollständige Achse hinaus.
• Die Größe der untergeordneten Achse entspricht nicht der Größe der entsprechenden vollständigen Achse. In diesem Fall sollte stattdessen die vollständige Achse verwendet werden.

Parameter:

TensorMappingAttr

Faktormappings für jede Dimension eines Tensors.

Syntax:

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

Einschränkungen:

• Elemente in dim_mappings müssen die Einschränkungen in DimMappingAttr erfüllen.
• Keine doppelten Faktorindexe für Dimensionen.

Parameter:

TensorShardingAttr

Tensor-Sharding

Syntax:

#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
>

Ein Tensor-Sharding ist an ein bestimmtes Mesh gebunden und kann nur auf Achsennamen aus diesem Mesh verweisen. Die Dimensions-Shards geben für jede Dimension des Tensors an, entlang welcher Achsen (oder untergeordneten Achsen) sie von der Haupt- zur Nebenachse aufgeteilt wird. Alle anderen Achsen, die keine Dimension aufteilen, werden entweder implizit oder explizit (wenn sie in der Liste der replizierten Achsen enthalten sind) repliziert.

Wenn für einen Tensor kein Sharding-Attribut angegeben ist, entspricht das einem vollständig offenen Tensor-Sharding.

Das Mesh, an das dieses Sharding gebunden ist, kann entweder durch einen Symbolnamen angegeben werden, der auf ein entsprechendes MeshOp-Symbol verweist, oder durch ein Inline-MeshAttr.

Ein Sharding kann nicht reduzierte Achsen haben (angegeben durch unreduced_axes), was bedeutet, dass der Tensor entlang dieser Achsen nicht reduziert wird. Wenn die kontrahierende Dimension einer Matmul beispielsweise sowohl in der linken als auch in der rechten Seite entlang der Achse x partitioniert wird, ist das Ergebnis entlang x nicht reduziert. Wenn Sie „all-reduce“ auf den Tensor entlang der nicht reduzierten Achsen anwenden, wird der Tensor entlang dieser Achsen repliziert. Ein Tensor mit nicht reduzierten Achsen muss jedoch nicht sofort vollständig reduziert werden. Er kann nicht reduziert bleiben, wenn er an lineare Operationen wie stablehlo.add übergeben wird (solange sowohl „lhs“ als auch „rhs“ nicht reduziert sind), und danach vollständig reduziert werden. Wir gehen davon aus, dass der Reduzierungstyp „Summe“ ist. Andere Reduzierungen werden möglicherweise in Zukunft unterstützt.

Einschränkungen:

• Elemente in dim_shardings müssen die in DimensionShardingAttr aufgeführten Einschränkungen erfüllen.
• Elemente in replicated_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Elemente in unreduced_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn der entsprechende Tensortyp kein ShapedType ist, muss das Sharding den Rang 0 und keine replizierten Achsen haben.
• Wenn es sich um ein ShapedType handelt, gilt Folgendes:
  • Der Tensor muss einen Rang haben.
  • Die Anzahl der Dimensions-Shards entspricht dem Rang des Tensors.
  • Dimensionen der Größe 0 werden nicht partitioniert.
• Es gibt keine doppelten Achsenreferenzen oder untergeordneten Achsen, die sich in dim_shardings, replicated_axes und unreduced_axes überschneiden.
• Die Elemente in replicated_axes und unreduced_axes werden in Bezug auf mesh_or_ref sortiert (siehe AxisRefAttr::getMeshComparator).

Parameter:

TensorShardingPerValueAttr

Tensor-Sharding pro Operanden/Ergebnis eines Vorgangs

Syntax:

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

Eine Liste von TensorShardingAttrs, eine für jeden Operanden/jedes Ergebnis eines Vorgangs.

Einschränkungen:

• Elemente in shardings müssen die Einschränkungen von TensorShardingAttr erfüllen.

Parameter:

Enums

EdgeNodeType

Edge-Knotentyp-Enum

Schutzhüllen:

PropagationDirection

Enumeration für die Weitergaberichtung

Schutzhüllen:

ReductionOp

Enum für die Reduzierung

Schutzhüllen: