Dialekt „sdy“

Der Shardy-Dialekt (SDY) definiert eine achsebasierte Tensor-Sharding-Darstellung und zusätzliche API-Komponenten, um Shardings an Tensoren anzuhängen.

Vorgänge

sdy.all_gather (sdy::AllGatherOp)

Durchführt eine All-Gather-Kommunikation entlang von Achsen

Syntax:

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

Hiermit werden Teile eines Tensors entlang der in gathering_axes angegebenen Achsen zusammengefasst.

gathering_axes ist eine Liste von Listen von Achsen. Die äußere Liste ist größer als die Dimensionen des Tensors. In jeder inneren Liste werden die Achsen angegeben, entlang derer eine separate Datenerhebung für die jeweilige Dimension durchgeführt werden soll. Sie wird auf das Sharding des Operanden (tensor) angewendet, um das Sharding des Ergebnisses (out_sharding) zu erhalten.

out_sharding wird nicht zum Festlegen des Sharding der Ergebnisse verwendet. Stattdessen wird die Sharding-Ebene des Ergebnisses durch die Sharding-Ebene des Operanden und der gathering_axes bestimmt. out_sharding muss mit dieser abgeleiteten Sharding-Ebene ü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 die in Sdy_CollectiveOpInterface aufgeführten Einschränkungen erfüllen.
• Die Elemente in gathering_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Wenn gathering_axes auf das Operanden-Sharding angewendet wird, ergibt sich out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, Sdy_CollectiveOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.all_reduce (sdy::AllReduceOp)

All-Reduce-Kommunikation entlang von Achsen ausführen

Syntax:

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

Reduziert Blöcke 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 die in Sdy_CollectiveOpInterface aufgeführten Einschränkungen erfüllen.
• reduction_axes muss den in AxisRefListAttr aufgeführten Einschränkungen entsprechen.
• reduction_axes darf sich nicht mit den Sharding-Achsen des Operanden überschneiden.

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.all_slice (sdy::AllSliceOp)

Führt einen dynamischen Schnittvorgang entlang von Achsen aus.

Syntax:

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

Schneidet Teile eines Tensors entlang der in slicing_axes angegebenen Achsen. Zwischen sdy.all_slice und sdy.all_gather besteht eine algebraische Dualität.

slicing_axes ist eine Liste von Listen von Achsen. Die äußere Liste ist größer als die Dimensionen des Tensors. In jeder inneren Liste sind die Achsen angegeben, entlang derer eine Aufschlüsselung für die jeweilige Dimension erfolgen soll. Sie wird auf das Sharding des Operanden (tensor) angewendet, um das Sharding des Ergebnisses (out_sharding) zu erhalten.

out_sharding wird nicht zum Festlegen des Sharding der Ergebnisse verwendet. Stattdessen wird die Sharding-Ebene des Ergebnisses durch die Sharding-Ebene des Operanden und der slicing_axes bestimmt. out_sharding muss mit dieser abgeleiteten Sharding-Ebene ü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:

• Die Elemente in slicing_axes müssen die in AxisRefListAttr aufgeführten Einschränkungen erfüllen.
• Muss die in Sdy_CollectiveOpInterface aufgeführten Einschränkungen erfüllen.
• Wenn slicing_axes auf das Operanden-Sharding angewendet wird, ergibt sich out_sharding.

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.all_to_all (sdy::AllToAllOp)

Ermöglicht eine All-to-All-Kommunikation entlang von Achsen.

Syntax:

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

Bei dieser Operation werden für jedes Tupel (Achsen, Src_dim, tgt_dim) in der Parameterliste Teile eines Tensors entlang der Dimension tgt_dim und der in axes angegebenen Achsen zugeschnitten, entlang der Achsen verstreut und entlang der Dimension src_dim zusammengefügt.

Dieser Vorgang ist im Wesentlichen eine Kombination aus einem All-Gather entlang von src_dim und axes, gefolgt von einem All-Slice entlang von tgt_dim und axes. Das heißt, ein Suffix der Dimensionsachse für die Sharding-Achse src_dim des Eingabetensors wird an die Dimensionsachse für die Sharding-Achse tgt_dim des Ausgabetensors angehängt.

Die All-to-All-Speichersegmentierung wird auf die Segmentierung des Operanden (tensor) angewendet, um die Segmentierung des Ergebnisses (out_sharding) zu erhalten.

out_sharding wird nicht zum Festlegen des Sharding der Ergebnisse verwendet. Stattdessen wird die Sharding-Ebene des Ergebnisses durch die Sharding-Ebene des Operanden bestimmt. src_dim, tgt_dim und axes müssen mit dieser abgeleiteten Sharding-Ebene übereinstimmen.out_sharding

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

Merkmale: SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface, Sdy_CollectiveOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.collective_permute (sdy::CollectivePermuteOp)

Durchführt eine kollektive Permutation, um Achsen zu ersetzen

Syntax:

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

Es wird ein Teil des Eingabetensors von jedem Gerät an ein anderes gesendet, um die Achsen neu anzuordnen oder zu ersetzen, die den Tensor in Shards aufteilen.

Eine kollektive Permutation kann das Eingabe-Sharding so transformieren, dass jede Dimension so shardet werden muss wie zuvor, d.h., sie muss entlang von Achsen shardet werden, deren Produkt der Größe der Achsen entspricht, die den Tensor zuvor shardet haben.

Das ist nützlich, um Achsen in einer einzelnen Dimension oder in verschiedenen Dimensionen neu anzuordnen und shardierte Achsen durch replizierte Achsen zu ersetzen.

Im folgenden Beispiel beträgt die Größe des ge shardeten 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 die in Sdy_CollectiveOpInterface aufgeführten Einschränkungen erfüllen.
• Wenn das Eingabe- und Ausgabe-Sharding unterschiedliche Netze hat, müssen diese Netze genau dieselben Achsen und eine andere Reihenfolge der Geräte-IDs haben.
• Für jede Dimension muss das Produkt der Größen der Sharding-Achsen in out_sharding mit dem der Sharding-Achsen der entsprechenden Operandendimension übereinstimmen.

Merkmale: SameOperandsAndResultType

Schnittstellen: CollectiveOpInterface, InferTypeOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.constant (sdy::ConstantOp)

Kontinuierlicher Betrieb

Erstellt einen output-Tensor aus einer Konstante value.

Siehe: 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

Oberflächen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Ergebnisse:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Dataflow-Edge-Operation

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 Gruppe von Quellen (jede ist entweder ein Operand von X oder ein Operand des Blockterminators von X) und einer Gruppe von Zielen (jedes ist entweder ein Ergebnis von X oder ein Blockargument von X), sodass alle Quellen und Ziele auf dieselbe Weise geSharded 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
                  })

Diese While-Operation 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.

Ein sdy.data_flow_edge nimmt als Eingabe den Inhaber einer Kante an (kann eines der Ziele sein, aber vorzugsweise ein OP-Ergebnis anstatt eines Blockarguments), der keine anderen Verwendungen haben sollte. Dieser Operator ist nicht rein, da er eine Eingabe annehmen kann, die ursprünglich keine Verwendung hatte.

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

• werden die Änderungen separat über jede Kante übertragen.
• das Sharding für jede Kante separat aktualisieren, anstatt alle Ziele gleichzeitig (z.B. hat eine Operation eine einzelne unveränderliche TensorShardingPerValueAttr für das Ergebnis-Sharding).
• Fügen Sie jede Kante der Arbeitsliste separat hinzu, wenn sich das Sharding einer Quelle geändert hat.

Bei der Propagation werden Shardings zwischen allen Quellen und Zielen einer sdy.data_flow_edge so weitergegeben, als wäre es eine reguläre Operation mit den Quellen als Operanden und den Zielen als Ergebnissen und einer Identität sdy.op_sharding_rule. Das bedeutet, dass die Vorwärtsweitergabe von den Quellen zu den Zielen und die Rückwärtsweitergabe von den Zielen zu den Quellen erfolgt.

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

Merkmale: SameOperandsAndResultType

Oberflächen: InferTypeOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.manual_computation (sdy::ManualComputationOp)

Parallele Ausführung auf mehreren Geräten 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)

Springen Sie zu einer Region, die in Form von lokalem Code pro Gerät mit expliziten Kollektiven geschrieben ist, bei der logische Formen mit lokalen physischen Bufferformen pro Gerät übereinstimmen und Kollektive genau der physischen geräteübergreifenden Kommunikation entsprechen.

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

Einschränkungen:

• Die Elemente in in_shardings und out_shardings müssen den in TensorShardingAttr aufgeführten Einschränkungen entsprechen.
• Die Anzahl der globalen und lokalen Tensoreingaben/-ausgaben der Operatorregion muss übereinstimmen.
• Die manuellen Achsen müssen in jeder Dimensionsaufteilung vor allen kostenlosen Achsen stehen.
• Die manuellen Achsen können kein Padding einfügen. Die Dimensionsgröße muss durch die entsprechende Größe der manuellen Achsen teilbar sein.
• Die globalen und lokalen Formen der Argumente/Ergebnisse der Operatorregionen müssen übereinstimmen.
• Manuelle Achsen werden nicht aufgeteilt.

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

Oberflächen: ShardableDataFlowOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.mesh (sdy::MeshOp)

Benanntes Netzwerk

Syntax:

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

Definiert ein neues benanntes Mesh. Alle Meshes in einem Modul müssen dieselbe Anzahl von Geräten haben (außer Meshes mit einer einzelnen device_id). Das Mesh ist ein Symbol-Vorgang, der in der SymbolTable des Moduls angezeigt wird und auf den über seine name verwiesen werden kann.

Merkmale: HasParent<ModuleOp>

Oberflächen: Symbol

Attribute:

sdy.named_computation (sdy::NamedComputationOp)

Benannter Berechnungsvorgang

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)

Hiermit werden Berechnungen, d.h. Befehlsblöcke, gruppiert und benannt. Die Weiterleitung erfolgt in die Region und aus der Region, als wäre alles inline.

So können Sie die Weiterleitung von Aufrufanweisungen an andere Funktionen steuern. Alle Nutzer von Shardy sollten einen Import-/Export-Pass schreiben, der ihre Aufrufvorgänge in sdy.named_computation-Vorgänge umwandelt, indem der Body der aufgerufenen Funktion in den Body der named_computation dupliziert oder kopiert wird.

Der Typ der einzelnen Blockargumente und zurückgegebenen Werte in der Region muss mit dem Typ der Operanden und dem Ergebnistyp des Operators ü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

Oberflächen: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Vorgang für die Propagation Barrier

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. Bei der Ausbreitung ist dies jedoch nur in eine bestimmte Richtung möglich.

Dadurch wird verhindert, dass Shardings zwischen den Verwendungen des Ergebnisses der Barriere-Operation und ihres 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 keine Sharding-Informationen über diese Operation weitergegeben werden können.
• BOTH kann nicht angegeben werden, da diese Operation redundant wäre.

Merkmale: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Oberflächen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Operanden:

Ergebnisse:

sdy.reshard (sdy::ReshardOp)

Er teilt einen Tensor in einem anderen Sharding neu auf.

Syntax:

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

Der Eingabetensor wird mit der angegebenen Sharding-Methode neu gesplittet. Diese unterscheidet sich von der vorhandenen Sharding-Methode des Eingabetensors.

Sowohl ShardingConstraintOp als auch ReshardOp hängen einem Tensor eine Sharding-Einheit an. Die Lebensdauer beträgt:

1. Vor der Sharding-Weitergabe wird ShardingConstraintOp von Nutzern hinzugefügt.
2. Für die Sharding-Weitergabe wird ShardingConstraintOp verwendet. In den Ergebnissen der Sharding-Weitergabe ist keine ShardingConstraintOp enthalten. Stattdessen kann bei Bedarf ReshardOp hinzugefügt werden.
3. Ein Partitionierer wandelt eine ReshardOp in eine kollektive Operation (oder eine Identitätsoperation) um. In den Ergebnissen des Partitionierungstools sollte keine ReshardOp vorhanden sein.

// TODO(b/331680067). Fügen Sie ein Canonicalisierungsmuster hinzu, um redundante // reshard-Vorgänge zu entfernen.

Merkmale: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Oberflächen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Operanden:

Ergebnisse:

sdy.return (sdy::ReturnOp)

Durch den Vorgang sdy.return werden die Regionen beendet, die mit sdy-basierten Vorgängen und anderen Shardy-basierten Vorgängen verknüpft sind. Es ist variabel: Es nimmt als Argumente eine Liste von Werten an, deren Typen beliebig sein können (aber derselben Art, z.B. AnyTensor), und kann daher auf verschiedenen Ebenen des Shardy-IR-Stacks wiederverwendet werden.

Syntax:

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

Merkmale: AlwaysSpeculatableImplTrait, Terminator

Schnittstellen: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Operanden:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Begrenzt einen Tensor auf die angegebene Sharding-Methode.

Syntax:

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

Hiermit wird einem Zwischentensor (z.B. dem Ergebnis einer Matrixmultiplikation) ein Sharding hinzugefügt, um anzugeben, dass dieser Tensor oder ein Teil seiner Verwendungen so geSharded werden soll.

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

Diese Operation kann entweder:

• Sie werden nicht verwendet (dangling). Das bedeutet, dass der angehängte Sharding-Wert dem Sharding des Eingabetensors selbst entspricht.
• Es gibt Anwendungsfälle. Das bedeutet, dass die Sharding-Einschränkungs-Operationen mit dem angehängten Sharding gesplittet werden sollten, während andere Verwendungen des Eingabetensors möglicherweise eine andere Sharding-Methode haben. Wenn der Eingabetensor keine anderen Verwendungen hat, entspricht das Verhalten dem Fall „Keine Anwendungsfälle“.

Merkmale: SameOperandsAndResultType

Oberflächen: InferTypeOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.sharding_group (sdy::ShardingGroupOp)

Schränkt Tensoren in der Gruppe auf dasselbe Sharding ein.

Syntax:

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

Dieser Operator bietet eine Schnittstelle, um Tensoren Chunking-Gruppen zuzuweisen (Gruppen von Tensoren, die dieselbe Chunking-Methode haben müssen). Sobald bei der Replikation ein Gruppenelement geSharded wird, werden alle anderen Mitglieder auf genau dieselbe Weise geSharded. Bei diesem Vorgang wird die Gruppen-ID des Arguments verwendet und es wird kein Ergebnis zurückgegeben. Stattdessen wird die interne Darstellung der Sharding-Gruppe geändert, um der Gruppe mit der angegebenen ID den Eingabetensor hinzuzufügen.

Oberflächen: InferTypeOpInterface

Attribute:

Operanden:

Attribute

AllToAllParamAttr

Parameter für die All-to-All-Kommunikation

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 die All-to-All-Analyse durchgeführt werden soll.

Parameter:

AlltoAllParamListAttr

Liste der Parameter für die All-to-All-Kommunikation

Syntax:

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

Parameter:

AxisRefAttr

Verweis auf eine vollständige Achse oder eine untergeordnete Achse mit Split

Syntax:

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

Einschränkungen:

• name muss in der Grenze MeshAttr vorhanden sein.
• Wenn sub_axis_info vorhanden ist, muss es die Einschränkungen von SubAxisInfoAttr erfüllen.

Parameter:

AxisRefListAttr

Liste der Achsenreferenzen

Syntax:

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

Einschränkungen:

• Die Elemente in value müssen den Einschränkungen von AxisRefAttr entsprechen.
• Es gibt keine doppelten Achsenreferenzen oder untergeordneten Achsen, die sich überschneiden.
• Keine zwei benachbarten Achsenreferenzen sind aufeinanderfolgende Teilachsen derselben Vollachse, d.h., sie können zu einer Teilachse oder zur Vollachse zusammengeführt werden.

Parameter:

DimMappingAttr

Liste der Faktorindizes für eine Dimension

Eine leere Liste weist auf eine Nullzuordnung hin (wird mit * geparst/gedruckt). Das bedeutet, dass der Dimension keine Faktoren zugeordnet sind.

Einschränkungen:

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

Parameter:

DimensionShardingAttr

Dimensionsfragmentierung

Liste der Achsennamen, nach denen eine Tensordimension von Haupt- zu Nebenachse gesplittet werden soll, ein boolescher Wert, der angibt, ob die Dimension weiter gesplittet werden kann, und eine optionale Ganzzahl, die die Priorität dieser Dimensionsshardierung angibt, die bei der Sharding-Ausbreitung berücksichtigt wird. Prioritäten stammen aus Anmerkungen zur Nutzeraufteilung. Je niedriger der Wert, desto höher die Priorität. Wenn die Priorität in der Anmerkung fehlt, wird die höchste Priorität angenommen.

Einschränkungen:

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

Parameter:

ListOfAxisRefListsAttr

Liste der Achsenreferenzlisten

Syntax:

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

Parameter:

ManualAxesAttr

Liste der Achsen, für die eine manuelle Berechnung aktiviert ist

Syntax:

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

Parameter:

MeshAttr

Achsennetz und Liste der Geräte

Syntax:

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

Ein Mesh besteht aus einer Liste von Achsen und einer optionalen Liste von Geräte-IDs, die die Gerätereihenfolge angeben.

Wenn die Liste der Achsen leer ist, hat das Mesh eine implizite unbenannte Achse der Größe 1. Wenn in diesem Fall keine Geräte-ID-Liste angegeben wird, ist die implizite Geräte-ID-Liste [0]. Wenn eine Geräte-ID-Liste angegeben wird, muss sie eine einzelne Ganzzahl mit einem beliebigen positiven Wert enthalten. Wir nennen das den Fall mit maximaler Sharding-Anzahl.

Wenn für alle Fälle, in denen kein maximales Sharding verwendet wird, eine Geräte-ID-Liste angegeben wird, muss das Produkt der Achsengrößen der Anzahl der Geräte entsprechen. 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.

Hier einige Beispiele für Netze:

• Ein leeres Mesh ist ein Platzhalter-Mesh, das während der Übertragung ersetzt werden kann: <[]>
• Ein Mesh mit einer unbenannten Achse und einer expliziten Geräte-ID, die in der Regel für die maximale Sharding verwendet wird: <[], 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 Gerätereihenfolge angeben: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Einschränkungen:

• Elemente in axes dürfen keine doppelten Namen haben.
• Wenn device_ids angegeben ist:
  • Das Produkt der Achsengrößen muss mit der Anzahl der Geräte übereinstimmen.
  • Alle Elemente dürfen nicht negativ sein.
  • device_ids darf nicht mit iota(product(axis_sizes)) übereinstimmen.
  • „Sortiert nach device_ids“ muss iota(product(axis_sizes)) lauten.

Parameter:

MeshAxisAttr

Benannte Achse in einem Netz

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 nach verschiedenen Eigenschaften des Vorgangs partitioniert werden kann, z. B. nach Attributen, der Form der Operanden oder der Form der Ergebnisse. Beispiele:

%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 geShardet werden können. Dies dient hauptsächlich der Vollständigkeit, da viele Vorgänge wie punktweise Vorgänge Dimensionen mit der Größe 1 haben, die über Operanden und Ergebnisse hinweg übereinstimmen.

Faktortypen:

• reduction_factors enthält die Indizes der Faktoren, die reduziert werden müssen, z. B. die kontrahierenden Dimensionen in einem Punktoperator.
• need_replication_factors enthält die Indizes der Faktoren, die vollständig repliziert werden müssen, z. B. die sortierte Dimension bei einem Sortiervorgang.
• permutation_factors enthält die Indizes der Faktoren, die eine gemeinsame Permutation erfordern, wenn sie geSharded sind, z. B. die Padding-Dimensionen bei einem Padding-Vorgang.
• Alle anderen Faktoren gelten als Durchlauffaktoren, d.h. Faktoren, für die keine Kommunikation erforderlich ist, wenn sie auf dieselbe Weise über alle zugewiesenen Tensoren verteilt werden.

blocked_propagation_factors enthält die Faktoren, entlang derer Shardings nicht weitergegeben werden dürfen. Sie ist orthogonal zu den Faktortypen. Ein Faktor für die blockierte Ausbreitung kann dabei jeden der Faktortypen haben.

Mit is_custom_rule wird angegeben, 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 Operanden/Ergebniszuordnungen muss mit der Anzahl der Operanden/Ergebnisse der Operation übereinstimmen.
• Es gibt mindestens eine Zuordnung. Es kann keine Regel für eine Operation ohne Operanden/Ergebnisse geben.
• Der Rang jedes TensorMappingAttr stimmt mit dem Rang des entsprechenden Tensortyps überein.
• Für jede Gruppe von Faktoren (reduction_factors, need_replication_factors, permutation_factors):
  • Die Elemente müssen im Bereich [0, $factor_sizes] liegen.
  • Keine doppelten Faktorindizes innerhalb und zwischen den Gruppen.

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 volle Achse in n Unterachsen aufgeteilt wird, wird die Achse in [k_1,…,k_n] umgewandelt. Die i. Unterachse kann durch das Produkt aller Achsengrößen links davon m=prod(k_1,...,k_(i-1)) (auch Vorabgröße genannt) und der Größe k_i ausgedrückt werden. Das Attribut „sub-axis-info“ enthält daher diese beiden Zahlen und wird folgendermaßen dargestellt: (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 gesamten Achse teilen, d.h. sowohl pre-size als auch size müssen die Größe der gesamten Achse teilen und die untergeordnete Achse darf nicht über die gesamte Achse hinausgehen.
• Die Größe der untergeordneten Achse entspricht nicht der Größe der entsprechenden Vollachse. In diesem Fall sollte stattdessen die Vollachse verwendet werden.

Parameter:

TensorMappingAttr

Faktorzuordnungen für jede Dimension eines Tensors.

Syntax:

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

Einschränkungen:

• Die Elemente in dim_mappings müssen den Einschränkungen in DimMappingAttr entsprechen.
• Keine doppelten Faktorindizes in Dimensionen.

Parameter:

TensorShardingAttr

Tensor-Sharding

Syntax:

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

Ein Tensor-Sharding ist an ein bestimmtes Mesh gebunden und kann nur auf Achsennamen aus diesem Mesh verweisen. Die Dimensionsshardings geben für jede Dimension des Tensors an, entlang welcher Achsen (oder Unterachsen) er von Haupt- zu Nebenachsen gesplittet wird. Alle anderen Achsen, die keine Dimensionen sharden, werden entweder implizit oder explizit (wenn sie in der Liste der replizierten Achsen aufgeführt sind) repliziert.

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 MeshAttr-Element, das eingefügt wurde.

Einschränkungen:

• Die Elemente in dim_shardings müssen die in DimensionShardingAttr aufgeführten Einschränkungen erfüllen.
• Die Elemente in replicated_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 haben und keine replizierten Achsen.
• Der Tensor sollte einen Rang haben.
• Die Anzahl der Dimensionsshards entspricht dem Rang des Tensors.
• Dimensionen der Größe 0 werden nicht geSharded.
• Die Elemente in replicated_axes werden nach mesh_or_ref sortiert (siehe AxisRefAttr::getMeshComparator).

Parameter:

TensorShardingPerValueAttr

Tensor-Sharding nach Operand/Ergebnis einer Operation

Syntax:

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

Eine Liste von TensorShardingAttr, eine für jeden Operanden/das Ergebnis einer Operation.

Einschränkungen:

• Die Elemente in shardings müssen den Einschränkungen von TensorShardingAttr entsprechen.

Parameter:

Enums

PropagationDirection

Propagation direction enum

Schutzhüllen: