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.constant (sdy::ConstantOp)

Kontinuierlicher Betrieb

Erzeugt 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

Schnittstellen: 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 verlaufen.

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

Der Vorgang hat zwar n Datenflusskanten, die i-ten Datenflusskanten liegen zwischen den Quellen x_i, return_value_i und den Zielen y_i, pred_arg_i und body_arg_i.

Ein sdy.data_flow_edge nimmt als Eingabe das Stammziel einer Kante an (kann jedes der Ziele sein, aber vorzugsweise ein Operatorergebnis statt eines Blockarguments), das 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 über jeden Edge separat übertragen.
• das Sharding jeder 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, 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. Daher können wir davon ausgehen, dass sie durch einen Operator mit einem nicht registrierten sdy.sharding-Attribut definiert ist.

Merkmale: SameOperandsAndResultType

Schnittstellen: 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.

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

Attribute:

Operanden:

Ergebnisse:

sdy.mesh (sdy::MeshOp)

Benanntes Mesh

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)

Gruppiert eine Berechnung, d.h. einen Block von Operationen, und gibt ihr einen Namen. 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>

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

Schnittstellen: ConditionallySpeculatable, ShardableDataFlowOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Vorgang für die Ausbreitungsbarriere

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, Elementwise, SameOperandsAndResultType

Schnittstellen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Operanden:

Ergebnisse:

sdy.reshard (sdy::ReshardOp)

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 binden ein Sharding an einen Tensor. 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, Elementwise, SameOperandsAndResultType

Schnittstellen: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Auswirkungen: MemoryEffects::Effect{}

Attribute:

Operanden:

Ergebnisse:

sdy.return (sdy::ReturnOp)

Mit dem 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 das angegebene Sharding.

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 der Sharding-Wert des Eingabetensors selbst sein sollte.
• Es gibt Anwendungsfälle. Das bedeutet, dass die Sharding-Beschränkungsoperationen mit der angehängten Sharding-Funktion gesplittet werden sollten, während andere Verwendungen des Eingabetensors eine andere Sharding-Funktion haben können. Wenn der Eingabetensor keine anderen Verwendungen hat, entspricht das Verhalten dem Fall ohne Verwendungen.

Merkmale: Elementwise, SameOperandsAndResultType

Schnittstellen: InferTypeOpInterface

Attribute:

Operanden:

Ergebnisse:

sdy.sharding_group (sdy::ShardingGroupOp)

Sharding-Gruppenvorgang

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, die so erzwungen werden, dass sie identische Fragmente haben). Sobald bei der Replikation ein Gruppenelement geSharded wird, werden alle anderen Mitglieder auf genau dieselbe Weise geSharded. Dieser Vorgang verwendet die Argumentgruppen-ID und gibt kein Ergebnis zurück. Stattdessen wird die Darstellung der internen Fragmentierungsgruppe geändert, um den Eingabetensor der Gruppe mit der angegebenen ID hinzuzufügen.

Attribute:

Operanden:

Attribute

AxisRefAttr

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

Syntax:

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

Parameter:

DimMappingAttr

Liste der Faktorindizes für eine Dimension

Alle Faktorindizes müssen im Bereich [0, num_factors] liegen. Eine leere Liste weist auf eine Nullzuordnung hin (wird mit * geparst/gedruckt), d. h., die Dimension ist keinen Faktoren zugeordnet.

Parameter:

DimensionShardingAttr

Dimensionssharding

Liste der Achsennamen, auf die eine Tensordimension von Dur zu Moll fragmentiert werden soll, ein boolescher Wert, der angibt, ob die Dimension weiter fragmentiert werden kann, und eine optionale Ganzzahl, die die Priorität dieser Dimensionsfragmentierung angibt, die bei der Weitergabe der Fragmentierung 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.

Parameter:

ManualAxesAttr

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 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 maximalen Sharding-Fall.

Wenn für alle Fälle, in denen kein maximales Sharding verwendet wird, eine Geräte-ID-Liste angegeben wird, sollte 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 Mesh-Netzwerke:

• 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-Netzwerk 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]>

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

Mit einer Sharding-Regel wird angegeben, wie ein Vorgang nach verschiedenen Eigenschaften des Vorgangs partitioniert werden kann, z. B. nach Attributen, der Form von 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.

is_custom_rule gibt an, ob dies eine vom Nutzer für eine stablehlo.custom_call-Operation definierte Regel ist. Der Partitionierer weiß nicht, wie diese Vorgänge partitioniert werden sollen, daher muss ein Nutzer dies angeben. Bei einer benutzerdefinierten Regel wird die Regel immer beibehalten und nie entfernt. is_custom_rule kann nur für stablehlo.custom_call-Vorgänge „true“ sein.

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 und die i-te Unterachse kann durch das Produkt aller Achsengrößen zu ihrer linken m=prod(k_1,...,k_(i-1)) (auch Vorabgröße genannt) und 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 Größe k.

Parameter:

TensorMappingAttr

Zuordnungen für jede Dimension eines Tensors faktorisieren.

Syntax:

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

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 partitionieren, werden entweder implizit oder explizit (wenn sie in der Liste der replizierten Achsen aufgeführt sind) repliziert.

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

Parameter:

TensorShardingPerValueAttr

Tensor-Sharding nach Operand/Ergebnis einer Operation

Syntax:

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

Parameter:

Enums

PropagationDirection

propagation direction enum

Schutzhüllen: