'mpmd' 言語

運用

mpmd.assign（mpmd::AssignOp）

割り当てオペレーション

構文:

operation ::= `mpmd.assign` attr-dict $tensor `:` functional-type(operands, results)

ローカル テンソルをメッシュ内の完全なレプリケートとしてメッシュに割り当てます。

これは、jax ops を削減するときに導入される一時的な op で、ローカル型からメッシュ型に移行するために使用されます。これらのオペレーションは、func オペレーションの入力と結果がメッシュ テンソルになると、インポート時に削除されます。

結果タイプのメッシュ名はトポロジ内のメッシュに対応し、そのグローバル タイプはオペランド タイプと同じである必要があります。

割り当てオペレーションの送信元は、メッシュの送信元です（named_computation、メッシュ推論など）。

トレイト: AlwaysSpeculatableImplTrait、HasParent<::mlir::func::FuncOp, ForOp>

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

影響: MemoryEffects::Effect{}

属性:

オペランド:

結果:

mpmd.broadcast（mpmd::BroadcastOp）

ブロードキャスト オペレーション

構文:

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

テンソルが使用される任意のメッシュで、テンソルを転送（または複製）できます。転送されるたびに、転送元はオペランドの現在のロケーションになります。

トレイト: AlwaysSpeculatableImplTrait、SameOperandsAndResultType

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

影響: MemoryEffects::Effect{}

オペランド:

結果:

mpmd.call（mpmd::CallOp）

MPMD 固有の呼び出し関数

構文:

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

関数呼び出しオペレーション。たとえば、ループの本体を関数宣言でラップしてコードサイズを小さくする場合に便利です。

インターフェース: ArgAndResultAttrsOpInterface、CallOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

mpmd.for（mpmd::ForOp）

オペレーター向け

固定回数の反復で本体関数を実行した結果を返します。本体では反復インデックスを使用できます。

オプションのアンロール ファクタ（反復回数を分割する必要があります）を指定して、そのファクタで op の本体をアンロールできます。つまり、アンロール ファクタ N の場合、本体が複製されて N 個のコピーが作成され、反復回数が 1/N に減少します。最初のコピーを除く各コピーでは、ブロック引数ではなく前のコピーの結果が使用されます。また、反復インデックスはアンロール係数で乗算され、コピーごとにインクリメントされます。

for 演算子は任意の型を受け入れて返すことができますが、これらの TypeID は同じである必要があります（すべてのテンソル型、すべての MPMD メッシュ型など）。これにより、さまざまなレベルで op を使用し、実装と変換を共有できます。

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

インターフェース: ConditionallySpeculatable、LoopLikeOpInterface、OpAsmOpInterface、ShardableDataFlowOpInterface

属性:

オペランド:

結果:

mpmd.fragment（mpmd::FragmentOp）

フラグメント オペレーション

MPMD トポロジの特定のメッシュに計算（オペレーションのブロック）を割り当てます。これは、個々の SPMD プログラム フラグメントとして実行されることを目的としています。

フラグメントは、フラグメントと同じメッシュに割り当てられたメッシュ テンソルのみを受け取って返します。

フラグメントのメッシュ名は、トポロジのメッシュに対応している必要があります。

フラグメントには、オリジン（このフラグメントを形成した元の名前付き計算に関する情報を含むメタデータ）のリストと、ユーザー定義フラグメント（オリジンの空でないリストを含む）の場合に定義される staged_id が含まれます。オプションの in_shardings は、オペランドに対応するフラグメントのブロック引数のシャーディングを指定します。省略可能な out_shardings は、結果のシャーディングを指定します。

フラグメントのリージョンに自由変数があってはなりません。リージョン内の各ブロック引数と戻り値の型は、対応するメッシュ テンソルのグローバル テンソル型です。

トレイト: HasParent<::mlir::func::FuncOp, ForOp>、IsolatedFromAbove、RecursiveMemoryEffects、RecursivelySpeculatableImplTrait、SingleBlockImplicitTerminator<ReturnOp>、SingleBlock

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

属性:

オペランド:

結果:

mpmd.fragment_call（mpmd::FragmentCallOp）

フラグメント呼び出しオペレーション

MPMD フラグメント本体（MPMD トポロジの特定のメッシュに割り当てられた計算）を保持する関数呼び出しを表します。これは、個々の SPMD プログラム フラグメントとして実行されることを目的としています。

フラグメントのメッシュ名は、囲み関数のトポロジのメッシュに対応している必要があります。また、そのメッシュの形状は呼び出し先のメッシュの形状と一致している必要があります。

オリジンは、このフラグメント呼び出しに貢献したユーザー指定の計算（マージなど）を指定します。

呼び出し先の関数の入力型と結果型は、このオペレーションのオペランドと結果の対応するメッシュ テンソルのローカル テンソル型である必要があります。

例:

%2 = mpmd.fragment_call<mesh="m1",> origin=[] @my_fragment(%0, %1) :
  (m<esh>_tensor..., m<esh>_te>nsor...) - m<esh>_tensor...

トレイト: HasParent<::mlir::func::FuncOp>、MemRefsNormalizable

インターフェース: ArgAndResultAttrsOpInterface、CallOpInterface、SymbolUserOpInterface

属性:

オペランド:

結果:

mpmd.named_computation（mpmd::NamedComputationOp）

名前付きスコープ オペレーション

計算（一連のオペレーション）をグループ化し、UserOrigin 属性を介して名前と転置カウントを付与します。この NamedComputation は、MPMD で計算にメッシュを割り当てるため、または最適化のために使用できます。

転置カウント（デフォルト=0）は、名前付き計算が特定の数の JAX AD 転置変換によって生成されたかどうかを示します。

op のリージョンに自由変数があってはなりません。また、リージョン内の各ブロック引数と戻り値の型は、op の入力の型と戻り値の型と同じである必要があります。

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

インターフェース: ConditionallySpeculatable

属性:

オペランド:

結果:

mpmd.named_tensor（mpmd::NamedTensorOp）

テンソルをメッシュに割り当てる

構文:

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

テンソルの結果を特定の名前に関連付ける ID オペレーション。この NamedTensor は、MPMD でテンソルにメッシュを割り当てるために使用できます。

トレイト: AlwaysSpeculatableImplTrait、SameOperandsAndResultType

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

影響: MemoryEffects::Effect{}

属性:

オペランド:

結果:

mpmd.reduce（mpmd::ReduceOp）

クロス メッシュ削減オペレーション

構文:

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

テンソルをさまざまなメッシュで縮小し、使用する必要がある場所にブロードキャストできます。

トレイト: AlwaysSpeculatableImplTrait、SameOperandsAndResultType

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

影響: MemoryEffects::Effect{}

属性:

オペランド:

結果:

mpmd.return（mpmd::ReturnOp）

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

構文:

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

トレイト: AlwaysSpeculatableImplTrait、Terminator

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

影響: MemoryEffects::Effect{}

オペランド:

mpmd.transfer（mpmd::TransferOp）

転送オペレーション

構文:

operation ::= `mpmd.transfer` attr-dict $tensor `:` functional-type(operands, results)

分散テンソルをあるメッシュから別のメッシュに転送します。

オペランドと結果の型のメッシュ名は、トポロジ内のメッシュに対応している必要があり、グローバル型は同じである必要があります。

トレイト: AlwaysSpeculatableImplTrait、HasParent<::mlir::func::FuncOp>

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

影響: MemoryEffects::Effect{}

オペランド:

結果:

mpmd.unassign（mpmd::UnassignOp）

割り当て解除オペレーション

構文:

operation ::= `mpmd.unassign` attr-dict $tensor `:` functional-type(operands, results)

完全に複製されたテンソルをメッシュから割り当て解除します。

これは、jax ops を削減するときに導入される一時的な op で、ローカル型からメッシュ型に移行するために使用されます。これらのオペレーションは、func オペレーションの入力と結果がメッシュ テンソルになると、インポート時に削除されます。

オペランド タイプのメッシュ名はトポロジのメッシュに対応し、そのグローバル タイプは結果タイプと同一である必要があります。

トレイト: AlwaysSpeculatableImplTrait、HasParent<::mlir::func::FuncOp, ForOp>、InferTensorType

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

影響: MemoryEffects::Effect{}

属性:

オペランド:

結果:

属性

MeshWithOriginsAttr

メッシュの起源と一致します。

構文:

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

パラメータ:

MeshesWithOriginsAttr

オリジンを含むメッシュのリスト。

構文:

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

パラメータ:

NamedMeshAttr

名前と Mesh を含むペア。

構文:

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

パラメータ:

OriginAttr

メッシュ割り当てのオリジン。

構文:

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

メッシュ割り当てのオリジン。

origin_label は、人が読める形式のオリジンのラベルです。デバッグ目的で使用することを想定しています。

パラメータ:

ReductionAttr

削減を示します。

構文:

#mpmd.reduction<
  `none` | `add` | `max` | `min` | `mul` | `or` | `and`   # reduction_type
>

パラメータ:

TopologyAttr

名前付きメッシュのトポロジ。

構文:

#mpmd.topology<
  ::llvm::ArrayRef<NamedMeshAttr>   # meshes
>

パラメータ:

UserOriginAttr

ユーザー指定の計算の送信元。

構文:

#mpmd.user_origin<
  ::mlir::StringAttr,   # userName
  int64_t   # transposeCount
>

パラメータ:

型

MeshTensorType

メッシュ テンソル型

プログラムの MPMD トポロジのメッシュ内の特定の SPMD メッシュに RankedTensorType を割り当てます。この型には、SPMD メッシュに対するテンソルのシャーディング方法を指定するオプションのシャーディングが含まれます。シャーディングが存在しない場合、テンソルは完全に複製されます。

パラメータ:

列挙型

ReductionType

reduction type [割引タイプ] 属性

Cases:

EdgeNodeType

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

Cases:

PropagationDirection

伝播方向の列挙型

Cases:

ReductionOp

削減演算子列挙型

Cases: