「sdy」方言

The Shardy (SDY) dialect

Shardy（SDY）言語は、軸ベースのテンソル シャーディング表現と、シャーディングをテンソルに接続するための追加の API コンポーネントを定義します。

バージョンログ: 0.0.1: TensorShardingAttr に削減されていない軸を追加。

運用

sdy.all_gather（sdy::AllGatherOp）

軸に沿って all-gather 通信を実行します

構文:

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

gathering_axes で指定された軸に沿ってテンソルのチャンクを収集します。

gathering_axes は、軸のリストのリストです。外側のリストがテンソルの次元を超えています。内側の各リストは、それぞれのディメンションで個別の収集を実行する軸を指定します。これは、オペランド（tensor）のシャーディングに適用され、結果（out_sharding）のシャーディングを取得します。

out_sharding は結果のシャーディングを決定するために使用されません。代わりに、結果のシャーディングはオペランドと gathering_axes のシャーディングによって決定され、out_sharding はこの推論されたシャーディングと一致する必要があります。

例:

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

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• gathering_axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• オペランド シャーディングに gathering_axes を適用すると、out_sharding が得られます。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、Sdy_CollectiveOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.all_reduce（sdy::AllReduceOp）

軸に沿って all-reduce 通信を実行する

構文:

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

reduction_axes で指定された軸に沿ってテンソルのチャンクを削減します。reduction_axes の順序は結果には影響しませんが、対応するレプリカ グループの順序に影響する可能性があります。

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• reduction_axes は AxisRefListAttr に記載されている制約を満たしている必要があります。
• reduction_axes はメッシュに関してソートする必要があります。
• オペランド シャーディングと out_sharding は同等のディメンション シャーディングを持つ必要があります。
• reduction_axes は、オペランドのディメンション シャーディングと複製された軸と重複してはなりません（縮小されていない軸と重複することは可能です）。
• reduction_axes は、out_sharding の削減されていない軸と重複してはなりません。つまり、out_sharding は reduction_axes に沿って複製される必要があります（暗黙的または明示的に）。

特性: SameOperandsAndResultType

インターフェース: CollectiveOpInterface、InferTypeOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.all_slice（sdy::AllSliceOp）

軸に沿って動的スライス オペレーションを実行します

構文:

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

slicing_axes で指定された軸に沿ってテンソルをスライスします。sdy.all_slice と sdy.all_gather の間には代数的な双対性があります。

slicing_axes は、軸のリストのリストです。外側のリストがテンソルの次元を超えています。内側の各リストは、対応するディメンションでスライスを実行する軸を指定します。これは、オペランド（tensor）のシャーディングに適用され、結果（out_sharding）のシャーディングを取得します。

out_sharding は結果のシャーディングを決定するために使用されません。代わりに、結果のシャーディングはオペランドと slicing_axes のシャーディングによって決定され、out_sharding はこの推論されたシャーディングと一致する必要があります。

例:

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

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• slicing_axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• オペランド シャーディングに slicing_axes を適用すると、out_sharding が得られます。

特性: SameOperandsAndResultType

インターフェース: CollectiveOpInterface、InferTypeOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.all_to_all（sdy::AllToAllOp）

軸に沿って全対全通信を実行します

構文:

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

パラメータ リスト内の各（axes、src_dim、tgt_dim）タプルについて、このオペレーションは、テンソルのチャンクをディメンション tgt_dim と axes で指定された軸に沿ってスライスし、これらのチャンクを軸に沿って分散させ、ディメンション src_dim に沿って連結します。

このオペレーションは、基本的に src_dim と axes に沿った all-gather と、tgt_dim と axes に沿った all-slice の組み合わせです。つまり、入力テンソルの軸シャーディング ディメンション src_dim の接尾辞が出力テンソルの軸シャーディング ディメンション tgt_dim に追加されます。

all-to-all は、オペランド（tensor）のシャーディングに適用され、結果（out_sharding）のシャーディングを取得します。

out_sharding は結果のシャーディングを決定するために使用されません。代わりに、結果のシャーディングはオペランド src_dim、tgt_dim、axes のシャーディングによって決まり、out_sharding はこの推論されたシャーディングと一致する必要があります。

例:

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

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• パラメータ リストを空にすることはできません。
• params の各パラメータについて:
  • axes の要素は AxisRefAttr の制約を満たす必要があります。
  • src_dim と tgt_dim は有効なディメンション（負の値ではなく、テンソルのランクより小さい）である必要があります。
  • src_dim または tgt_dim は、すべてのパラメータで一意である必要があります。
  • src_dim は、すべてのパラメータで昇順に並べ替える必要があります。
• オペランド シャーディングで axes を src_dim から tgt_dim に移動すると、out_sharding が取得されます。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、Sdy_CollectiveOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.collective_permute（sdy::CollectivePermuteOp）

軸を置き換える集団置換通信を実行します

構文:

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

各デバイスから別のデバイスに入力テンソルのチャンクを送信して、テンソルをシャーディングする軸を並べ替えたり置換したりします。

集合置換では、各ディメンションが以前と同じようにシャーディングされるように入力シャーディングを変換できます。つまり、サイズの積が以前にテンソルをシャーディングした軸のサイズの積と一致する軸に沿ってシャーディングする必要があります。

これは、単一のディメンションまたは複数のディメンションで軸の順序を変更したり、シャードされた軸を複製された軸と入れ替えたりする場合に便利です。

次の例では、シャード テンソルのサイズは tensor<1x4x2xf32> であり、これは集団置換によって保持されます。

例:

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>

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• 入力と出力のシャーディングのメッシュが異なる場合、これらのメッシュは軸がまったく同じで、デバイス ID の順序が異なる必要があります。
• 各ディメンションについて、out_sharding のシャーディング軸サイズの積は、対応するオペランド ディメンションのシャーディングの積と一致する必要があります。

特性: SameOperandsAndResultType

インターフェース: CollectiveOpInterface、InferTypeOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.constant（sdy::ConstantOp）

定数オペレーション

定数 value から output テンソルを生成します。

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

例:

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

特性: AlwaysSpeculatableImplTrait

インターフェース: ConditionallySpeculatable、InferTypeOpInterface、NoMemoryEffect (MemoryEffectOpInterface)

影響: MemoryEffects::Effect{}

属性:

結果:

sdy.data_flow_edge（sdy::DataFlowEdgeOp）

データフロー エッジ オペレーション。

構文:

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

ある op X のデータフロー エッジは、一連のソース（それぞれが X のオペランドまたは X のブロック終端子のオペランド）と一連のターゲット（それぞれが X の結果または X のブロック引数）の間のブリッジを定義します。すべてのソースとターゲットは同じ方法でシャーディングする必要があります。

Op には、互いに直交する複数のデータフロー エッジを設定できます。

次に例を示します。

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

この while オペレーションには n 個のデータフロー エッジがあります。i 番目のデータフロー エッジは、ソース x_i、return_value_i とターゲット y_i、pred_arg_i、body_arg_i の間にあります。

sdy.data_flow_edge は、エッジの所有者（ターゲットのいずれか。ブロック引数ではなく、演算結果が望ましい）を入力として受け取ります。この所有者は他の用途に使用しないでください。この op は、元々使用されていなかった入力を受け取ることができるため、純粋ではありません。

sdy.data_flow_edge には、エッジのすべてのターゲットのオプションのシャーディングも保持されます。伝播中に、ターゲットのシャーディング（アタッチ可能な場合）ではなく、そのシャーディングを更新する必要があります。これは、op に多くのエッジがある場合に便利です。

• 各エッジを個別に伝播します。
• すべてのターゲットを一度にではなく、各エッジのシャーディングを個別に更新します（例: オペレーションには結果シャーディング用の単一の不変の TensorShardingPerValueAttr があります）。
• ソースのシャーディングが変更されたときに、各エッジをワークリストに個別に追加します。

伝播は、sdy.data_flow_edge のすべてのソースとターゲット間でシャーディングを伝播します。これは、ソースをオペランド、ターゲットを結果、ID sdy.op_sharding_rule とする通常のオペレーションと同様です。つまり、順方向伝播はソースからターゲットへの伝播であり、逆方向伝播はターゲットからソースへの伝播です。

sdy.data_flow_edge の入力を SdyDialect op で定義することは許可されていないため、登録されていない sdy.sharding 属性を持つ op で定義されていると想定できます。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.func_data_flow_edge（sdy::FuncDataFlowEdgeOp）

Func 入出力データフロー エッジ オペレーション。

構文:

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

データフロー エッジ オペレーションですが、関数引数または呼び出し結果用です。オペランドが BlockArgument の場合、呼び出し元 callOp の引数から func 引数のユーザーへのブリッジになります。func 引数ごとに 1 つの func データフロー エッジがあります。オペランドが OpResult の場合、呼び出された funcOp の戻り値から呼び出し結果のユーザーへのブリッジになります。呼び出し結果ごとに 1 つの func データフロー エッジがあります。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、SymbolUserOpInterface

オペランド:

結果:

sdy.manual_computation（sdy::ManualComputationOp）

手動コレクティブを使用したマルチデバイス並列処理

構文:

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)

明示的なコレクティブを使用してデバイスごとのローカルコードで記述されたリージョンにジャンプします。ここでは、論理形状がデバイスごとのローカルの物理バッファ形状と一致し、コレクティブがクロスデバイス物理通信に正確に対応します。

本体は manual_axes に対してローカルです。伝播は、手動軸リストに含まれていない自由軸の本体を介して行われます。

ランク付けされていないテンソルは、ランク 0 のシャーディング（完全に複製されたシャーディング）を持つことが想定されています。

制約:

• in_shardings と out_shardings の要素は、TensorShardingAttr に記載されている制約を満たす必要があります。
• op リージョンのグローバル テンソル入力/出力の数は一致している必要があります。
• 手動軸は、各 dim シャーディングの自由軸の前に配置する必要があります。
• 手動軸ではパディングを導入できません。つまり、ディメンションのサイズは、対応する手動軸のサイズで割り切れる必要があります。
• op リージョン引数/結果のグローバル シェイプとローカル シェイプは一致している必要があります。

トレイト: IsolatedFromAbove、RecursiveMemoryEffects、SingleBlockImplicitTerminator<ReturnOp>、SingleBlock

インターフェース: ShardableDataFlowOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.mesh（sdy::MeshOp）

名前付きメッシュ

構文:

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

新しい名前付きメッシュを定義します。モジュール内のすべてのメッシュのデバイス数は同じである必要があります（単一の device_id を持つメッシュを除く）。メッシュは、モジュールの SymbolTable に表示され、その name で参照できる Symbol オペレーションです。

トレイト: HasParent<ModuleOp>、SymbolName

インターフェース: Symbol

属性:

sdy.named_computation（sdy::NamedComputationOp）

名前付き計算オペレーション

構文:

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)

計算（一連のオペレーション）をグループ化して名前を付けます。伝播は、すべてがインライン化されているかのようにリージョン内外に流れます。

これは、呼び出し命令を介して他の関数に伝播する処理に使用できます。Shardy のユーザーは、呼び出しオペレーションを sdy.named_computation オペレーションに変換し、呼び出された関数の本体を named_computation の本体に複製/コピーするインポート/エクスポート パスを記述する必要があります。

リージョン内の各ブロック引数と戻り値の型は、オペランドの型と op の結果の型と同じである必要があります。

例:

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

トレイト: IsolatedFromAbove、RecursiveMemoryEffects、RecursivelySpeculatableImplTrait、SingleBlockImplicitTerminator<ReturnOp>、SingleBlock

インターフェース: ConditionallySpeculatable、InferTypeOpInterface、ShardableDataFlowOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.propagation_barrier（sdy::PropagationBarrierOp）

伝播バリア オペレーション

構文:

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

この op は、入力として受け取ったのと同じ値を出力する ID op のように動作します。ただし、伝播の観点から見ると、この設定では特定の方向にのみ伝播が流れるようになります。

これにより、バリア オペレーションの結果とそのオペランドの使用間でシャーディングが伝播されるのを防ぎます。

• FORWARD は、シャーディングがオペランドから結果にのみフローできることを意味します。
• BACKWARD は、シャーディングが結果からオペランドにのみフローできることを意味します。
• NONE は、このオペレーションを介してシャーディングを伝播できないことを意味します。
• この op は冗長になるため、BOTH は指定できません。

トレイト: AlwaysSpeculatableImplTrait、SameOperandsAndResultType

インターフェース: ConditionallySpeculatable、InferTypeOpInterface、NoMemoryEffect (MemoryEffectOpInterface)

影響: MemoryEffects::Effect{}

属性:

オペランド:

結果:

sdy.reduce_scatter（sdy::ReduceScatterOp）

軸に沿って reduce-scatter 通信を実行します

構文:

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

reduce_scatter_axes で指定された軸に沿ってテンソルのチャンクを削減し、同じ軸に沿って結果を分散します。このオペレーションは、基本的に同じ reduce_scatter_axes に沿って sdy.all_reduce の後に sdy.all_slice が続く組み合わせです。

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• reduce_scatter_axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• オペランド シャーディングに reduce_scatter_axes を適用すると、out_sharding が取得されます。

特性: SameOperandsAndResultType

インターフェース: CollectiveOpInterface、InferTypeOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.replicated_to_unreduced（sdy::ReplicatedToUnreducedOp）

暗黙的または明示的に複製された軸を、縮小されていない軸に移動します。

構文:

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

axes は、オペランドで暗黙的または明示的に複製される必要があります。このオペレーションにより、結果で縮小されなくなります。次の関係があります。

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

例:

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

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• axes は AxisRefListAttr に記載されている制約を満たしている必要があります。
• axes はメッシュに関してソートする必要があります。
• axes は空ではありません。
• 入力と出力のシャーディングは、同じディメンション シャーディングを持つ必要があります。
• axes は、オペランド シャーディングで暗黙的または明示的に複製する必要があります。
• inUnreducedAxes + axes = outUnreducedAxes。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、Sdy_CollectiveOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.reshard（sdy::ReshardOp）

テンソルを別のシャーディングに再シャーディングします

構文:

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

入力テンソルの既存のシャーディングとは異なる、指定されたシャーディングで入力テンソルを再シャーディングします。

ShardingConstraintOp と ReshardOp はどちらも、テンソルにシャーディングをアタッチします。有効期間は次のとおりです。

1. シャーディングの伝播の前に、ユーザーが ShardingConstraintOp を追加します。
2. シャーディングの伝播は ShardingConstraintOp を消費します。シャーディング伝播の結果に ShardingConstraintOp がありません。代わりに、必要に応じて ReshardOp を追加できます。
3. パーティショナーは、ReshardOp を集合演算（または ID 演算）に変換します。パーティショナーの結果に ReshardOp が含まれていない必要があります。

トレイト: AlwaysSpeculatableImplTrait、SameOperandsAndResultType

インターフェース: ConditionallySpeculatable、InferTypeOpInterface、NoMemoryEffect (MemoryEffectOpInterface)、SymbolUserOpInterface

影響: MemoryEffects::Effect{}

属性:

オペランド:

結果:

sdy.return（sdy::ReturnOp）

sdy.return オペレーションは、sdy リージョンベースのオペレーションと、その他の Shardy リージョンベースのオペレーションにアタッチされたリージョンを終了します。可変長引数です。型が任意（ただし、同じ種類、たとえば AnyTensor）の値のリストを引数として受け取るため、Shardy IR スタックのさまざまなレベルで再利用できます。

構文:

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

トレイト: AlwaysSpeculatableImplTrait、ReturnLike、Terminator

インターフェース: ConditionallySpeculatable、NoMemoryEffect (MemoryEffectOpInterface)、RegionBranchTerminatorOpInterface

影響: MemoryEffects::Effect{}

オペランド:

sdy.sharded_to_unreduced（sdy::ShardedToUnreducedOp）

オペランドのシャード化された軸の一部を、結果の縮小されていない軸に移動します。

構文:

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

axes は、オペランドのシャーディングに使用する必要があります。このオペレーションにより、結果で縮小されなくなります。次の関係が成り立ちます。

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes)。ここで、all-gather、sharded-to-unreduced、all-reduce は同じ軸に適用されます。

例:

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

制約:

• Sdy_CollectiveOpInterface に記載されている制約を満たす必要があります。
• axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• オペランド シャーディングに axes を適用すると、out_sharding が得られます。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、Sdy_CollectiveOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.sharding_constraint（sdy::ShardingConstraintOp）

テンソルを指定されたシャーディングに制約します

構文:

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

中間テンソル（matmul の結果など）にシャーディングを適用して、そのテンソルまたはその使用のサブセットをシャーディングする方法を指定します。

シャーディングに開いたディメンションと制約のない軸がある場合、テンソルは開いたディメンションに沿ってさらにシャーディングできます。

このオペレーションは次のいずれかを行います。

• 使用されていない（ぶら下がっている） - つまり、アタッチされたシャーディングは入力テンソル自体がシャーディングされる方法です。
• 使用がある - つまり、アタッチされたシャーディングは、シャーディング制約 op の使用方法をシャーディングする方法です。一方、入力テンソルの他の使用方法では、異なるシャーディングが使用される可能性があります（入力テンソルに他の使用方法がない場合、動作は使用がない場合と同じです）。

特性: SameOperandsAndResultType

インターフェース: InferTypeOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

sdy.sharding_group（sdy::ShardingGroupOp）

グループ内のテンソルが同じシャーディングを持つように制約します。

構文:

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

この op は、テンソルをシャーディング グループ（同じシャーディングが適用されるテンソルのグループ）に割り当てるインターフェースを提供します。伝播中に、1 つのグループ要素がシャーディングされると、他のすべてのメンバーもまったく同じ方法でシャーディングされます。このオペレーションは引数グループ ID を受け取り、結果を返しません。代わりに、内部シャーディング グループ表現を変更して、指定された ID を持つグループに入力テンソルを追加します。

インターフェース: InferTypeOpInterface

属性:

オペランド:

属性

AllToAllParamAttr

All-to-all パラメータ

構文:

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

all-to-all を実行する軸とソース/ターゲット ディメンションを含むタプル。

パラメータ:

AllToAllParamListAttr

all-to-all パラメータのリスト

構文:

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

パラメータ:

AxisRefAttr

完全な軸または分割されたサブ軸のいずれかへの参照

構文:

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

制約:

• name は、境界 MeshAttr に存在している必要があります。
• sub_axis_info が存在する場合は、SubAxisInfoAttr の制約を満たしている必要があります。

パラメータ:

AxisRefListAttr

軸参照のリスト

構文:

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

制約:

• value の要素は AxisRefAttr の制約を満たす必要があります。
• 重複する軸参照や、互いに重なり合うサブ軸はありません。
• 隣接する 2 つの axis-ref が同じフル軸の連続するサブ軸であることはありません。つまり、1 つのサブ軸またはフル軸に統合できます。

パラメータ:

AxisToPropagationDetailsAttr

特定の軸とソースの伝播エッジフローの詳細。

構文:

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

特定の軸に沿って、ソース値の参照をターゲット値の参照のリストにマッピングします。

パラメータ:

DimMappingAttr

ディメンションのファクタ インデックスのリスト

空のリストは、これが null マッピング（* で解析/出力）であることを示します。つまり、ディメンションはどのファクタにもマッピングされていません。

制約:

• 少なくとも 1 つのファクタ インデックスがあります。
• 係数インデックスは [0, $factor_sizes) の範囲内である必要があります。
• 複数の要因がある場合、サイズ 1 の要因は存在できません。
• 重複するファクタ インデックスはありません。

パラメータ:

DimensionShardingAttr

ディメンション シャーディング

テンソル ディメンションをシャードする軸名のリスト（メジャーからマイナー）、ディメンションをさらにシャードできるかどうかを示すブール値、このディメンション シャーディングの優先度を示すオプションの整数。この優先度は、シャーディングの伝播時に考慮されます。優先度はユーザー シャーディング アノテーションから取得され、値が小さいほど優先度が高くなります。アノテーションに優先度が指定されていない場合は、最も高い優先度が想定されます。

制約:

• axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• ディメンション シャーディングに優先度がある場合:
  • 優先度が 0 以上である。
  • ディメンションが閉じている場合、少なくとも 1 つの軸があります。

パラメータ:

EdgeValueRefAttr

type 型の値エッジの特定のインデックスへの参照。

構文:

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

パラメータ:

ListOfAxisRefListsAttr

軸参照リストのリスト

構文:

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

パラメータ:

ManualAxesAttr

ManualComputationOp が手動で計算する軸のリスト

構文:

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

パラメータ:

MeshAttr

軸のメッシュとデバイスのリスト

構文:

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

メッシュは、軸のリストと、デバイスの順序を指定するデバイス ID のオプションのリストです。

軸のリストが空の場合

• device_ids が指定されていない場合、空のメッシュになります。
• device_ids が指定されている場合は、単一の非負の整数にする必要があります。これを最大シャーディング メッシュと呼びます。

軸のリストが指定されている場合

• デバイス ID リストが指定されている場合、軸サイズの積はデバイスの数と一致する必要があります。
• デバイス ID リストが指定されていない場合、暗黙的なデバイス ID リストは iota(product(axes)) になります。簡略化のため、iota(product(axes)) と同じデバイス ID リストの指定も許可しません。この場合、デバイス ID リストは指定されません。
• 軸の合計サイズが 1 であっても、最大シャーディング メッシュではありません。

メッシュの例を次に示します。

• 空のメッシュは、伝播中に置き換えることができるプレースホルダ メッシュを表します: <[]>
• 軸リストのないメッシュと 1 つの非負のデバイス ID（最大シャーディング メッシュ）: <[], device_ids=[3]>
• 2 つの軸と暗黙的なデバイス ID iota(6) を持つメッシュ: <["a"=2, "b"=3]>
• 2 つの軸とデバイスの順序を指定する明示的なデバイス ID を含むメッシュ: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

制約:

• device_ids の要素は負の値であってはなりません。
• axes が空の場合、device_ids のサイズは 0（空のメッシュ）または 1（最大シャーディング メッシュ）になります。
• axes が空でない場合、
  • axes 内の要素に重複する名前を指定することはできません。
  • device_ids が指定されている場合、元の device_ids は iota(product(axis_sizes)) ではなく、並べ替えられた device_ids は iota(product(axis_sizes)) です。

パラメータ:

MeshAxisAttr

メッシュ内の名前付き軸

構文:

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

パラメータ:

OpShardingRuleAttr

オペレーションをパーティショニングする方法を指定します。

構文:

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

シャーディング ルールでは、op のさまざまなプロパティ（属性、オペランドの形状、結果の形状など）に基づいてオペレーションをパーティショニングする方法を指定します。例:

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

シャードできないサイズ 1 の要素も許可されます。これは主に完全性を確保するためです。多くのポイントワイズ演算などの演算には、オペランドと結果に対応するサイズ 1 のディメンションがあります。

要素のタイプ:

• reduction_factors には、ドット演算の縮約ディメンションなど、削減が必要な因子のインデックスが含まれます。これらの要素はオペランドに存在しても、結果には存在しません。
• need_replication_factors には、並べ替えオペレーションの並べ替えられたディメンションなど、完全なレプリケーションを必要とする因子のインデックスが含まれます。
• permutation_factors には、パディング オペレーションのパディング ディメンションなど、シャード化されている場合に collective-permute が必要な因子のインデックスが含まれます。
• 他のすべての要因はパススルー要因と見なされます。つまり、それらにマッピングされたすべてのテンソルで同じようにシャーディングされた場合、通信を必要としない要因です。

blocked_propagation_factors には、シャーディングの伝播が許可されない要因が含まれています。これは要素タイプと直交しています。つまり、ブロックされた伝播係数は任意の係数タイプにできます。

is_custom_rule は、これがユーザー定義のルールかどうかを示します。ユーザーは、カスタム呼び出しのシャーディング ルールを定義したり、標準オペレーションの事前定義されたシャーディング ルールをオーバーライドしたりできます。カスタムルールは常に保持され、削除されることはありません。

制約:

• オペランド/結果のマッピングの数は、オペランド/結果の数と一致している必要があります。
• マッピングが 1 つ以上ある（オペランド/結果のない op のルールは存在できない）。
• 各 TensorMappingAttr のランクは、対応するテンソル型のランクと一致します。
• 各要因グループ（reduction_factors、need_replication_factors、permutation_factors）について:
  • 要素は [0, $factor_sizes] の範囲内である必要があります。
  • 各グループ内およびグループ間で重複するファクタ インデックスがない。

パラメータ:

PropagationEdgesAttr

すべての伝播ステップの伝播エッジ メタデータ。

構文:

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

値の軸ごとの伝播の詳細のリスト。ステップ インデックスでグループ化されます。

パラメータ:

PropagationOneStepAttr

ステップごとの伝播メタデータ。

構文:

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

単一の伝播ステップのすべての軸の伝播の詳細。

パラメータ:

SubAxisInfoAttr

このサブ軸が完全な軸からどのように導出されるかについての情報

構文:

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

軸全体を n 個のサブ軸に分割する場合、軸は [k_1,...,k_n] に再形成されます。i 番目のサブ軸は、左側のすべての軸サイズの積 m=prod(k_1,...,k_(i-1))（事前サイズ）とサイズ k_i の積で表すことができます。したがって、sub-axis-info 属性はこれら 2 つの数値を保持し、サイズ m とサイズ k の場合は (m)k と表記されます。

制約:

• pre-size は 1 以上です。
• size は 1 より大きい。
• pre-size は軸全体のサイズを分割する必要があります。つまり、pre-size と size の両方が軸全体のサイズを分割し、補助軸が軸全体を超えないようにします。
• サブ軸のサイズが対応するフル軸のサイズと等しくない場合、代わりにフル軸を使用する必要があります。

パラメータ:

TensorMappingAttr

テンソルの各ディメンションのファクタ マッピング。

構文:

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

制約:

• dim_mappings の要素は DimMappingAttr の制約を満たす必要があります。
• ディメンション間で重複するファクター インデックスはありません。

パラメータ:

TensorShardingAttr

Tensor シャーディング

構文:

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

テンソル シャーディングは特定のメッシュにバインドされ、そのメッシュの軸名のみを参照できます。ディメンション シャーディングは、テンソルの各ディメンションについて、どの軸（またはサブ軸）に沿ってメジャーからマイナーにシャーディングされるかを示します。ディメンションをシャーディングしない他のすべての軸は、暗黙的または明示的に（複製された軸のリストに表示される場合）複製されます。

テンソルにシャーディング属性がない場合、完全にオープンなテンソル シャーディングと同等になります。

このシャーディングがバインドされるメッシュは、対応する MeshOp シンボルを参照するシンボル名、またはインライン MeshAttr で指定できます。

シャーディングには、縮小されていない軸（unreduced_axes で指定）を含めることができます。つまり、テンソルはこれらの軸に沿って縮小されません。たとえば、matmul の縮約ディメンションが lhs と rhs の両方で軸 x に沿ってシャーディングされている場合、結果は x に沿って縮約されません。縮約されていない軸に沿ってテンソルに all-reduce を適用すると、テンソルがそれらの軸に沿って複製されます。ただし、縮約されていない軸を持つテンソルは、すぐに全縮約する必要はありません。stablehlo.add などの線形演算に渡されるときに縮約されていない状態を維持し（lhs と rhs の両方が縮約されていない場合）、後で全縮約できます。削減タイプは合計であると想定しています。他の削減は今後サポートされる可能性があります。

制約:

• dim_shardings の要素は、DimensionShardingAttr に記載されている制約を満たす必要があります。
• replicated_axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• unreduced_axes の要素は、AxisRefListAttr に記載されている制約を満たす必要があります。
• 対応するテンソル型が ShapedType でない場合、シャーディングのランクは 0 で、複製された軸があってはなりません。
• ShapedType の場合は、次の手順を行います。
  • テンソルにはランクが必要です。
  • ディメンション シャーディングの数は、テンソルのランクと同じです。
  • サイズ 0 のディメンションはシャーディングされません。
• dim_shardings、replicated_axes、unreduced_axes の間で重複する軸参照やサブ軸はありません。
• replicated_axes と unreduced_axes のアイテムは mesh_or_ref に従って順序付けられます（AxisRefAttr::getMeshComparator を参照）。

パラメータ:

TensorShardingPerValueAttr

オペレーションのオペランド/結果ごとのテンソル シャーディング

構文:

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

TensorShardingAttr のリスト。オペレーションのオペランド/結果ごとに 1 つ。

制約:

• shardings の要素は TensorShardingAttr の制約を満たす必要があります。

パラメータ:

列挙型

EdgeNodeType

エッジノードタイプの列挙型

Cases:

PropagationDirection

伝播方向の列挙型

Cases:

ReductionOp

削減演算子列挙型

Cases: