CollectiveOpInterface(Sdy_CollectiveOpInterface)
すべての集約オペレーションのインターフェース。outSharding 属性の一般的な get/set をカプセル化します。
制約:
- オペランドにシャーディングがないと、
allowMissingInputSharding()は true を返します。 out_shardingは、対応する型に対して有効です。allowDifferentMeshes()が false を返す場合、オペランドと結果のシャーディングは同じメッシュにする必要があります。- オペランドと結果のシャーディングに同じランク。
メソッド:
getOutSharding
::mlir::sdy::TensorShardingAttr getOutSharding();
集約オペレーションの出力テンソルのシャーディングを返します。
setOutShardingAttr
void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);
集約オペレーションの出力テンソルのシャーディングを設定します。
getTensor
::mlir::TypedValue<::mlir::TensorType> getTensor();
集約オペレーションのテンサー オペランドを取得します。
getType
::mlir::Type getType();
集約オペレーションの結果の型を取得します。
allowDifferentMeshes
bool allowDifferentMeshes();
集約オペレーションで入力シャーディングと出力シャーディングに異なるメッシュを許可するかどうかを示します。
allowMissingInputSharding
bool allowMissingInputSharding();
集約オペレーションで入力にシャーディングを適用しない(暗黙的に完全に複製する)かどうかを示します。
ShardableDataFlowOpInterface(Sdy_ShardableDataFlowOpInterface)
shardy がこのインターフェースを拡張するオペレーションのデータフロー エッジを介してシャーディングを伝播できるようにするオペレーション インターフェース。
ある演算 X のデータフロー エッジは、一連のソース(それぞれ X のオペランドまたは X のブロック終端子のオペランド)と一連のターゲット(それぞれ X の結果または X のブロック引数)間のブリッジを定義します。これにより、すべてのソースとターゲットが同じ方法でシャーディングされるようにします。1 つのオペレーションには、互いに直交する複数のデータフロー エッジを含めることができます。
オーナーは、シャーディーのプロパゲーションによって使用されるデータフロー エッジのユーザー指定のターゲットです。ユーザーが任意に選択できますが、静的である必要があります。
次に例を示します。
y_1, ..., y_n = custom_op (x_1, ..., x_n)
((body_arg_1,..., body_arg_n) {
...
return return_value_1, ..., return_value_n
})
この custom_op には、データフロー エッジに 2 つのタイプがあります。return_value_i(ソース)と y_i(ターゲット)の間に n 個のエッジ、x_i(ソース)と body_arg_i(ターゲット)の間に n 個のエッジです。この場合、エッジ オーナーはターゲットと同じです。
複数のターゲットを持つオペレーションの例を次に示します。
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 の間にあります。
メソッド:
getBlockArgumentEdgeOwnerShardings
mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();
すべてのブロック引数データフロー エッジ オーナーのシャーディングを返します。
setBlockArgumentEdgeOwnerShardings
void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
すべてのブロック引数のエッジ オーナーの shardings を設定します。
getOpResultEdgeOwnerShardings
mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();
すべてのオペレーション結果データフロー エッジ オーナーのシャーディングを返します。
setOpResultEdgeOwnerShardings
void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
すべてのオペレーション結果エッジ オーナーの shardings を設定します。
transformTargetSharding
mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);
transformType に応じてターゲットの sharding を変換します。
詳しくは、DataFlowShardingTransformType をご覧ください。
getBlockArgumentEdgeOwners
mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();
すべてのブロック引数のエッジ オーナーを取得します。
getOpResultEdgeOwners
mlir::ResultRange getOpResultEdgeOwners();
すべてのオペレーション結果エッジ オーナーを取得します。
getEdgeSources
mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);
エッジ owner を指定して、データフロー エッジソースを取得します。
getEdgeOwnerFromTarget
mlir::Value getEdgeOwnerFromTarget(mlir::Value target);
所有者である可能性のある target を指定して、データフロー エッジの所有者 target を取得します。
getEdgeOwnerFromSource
mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);
source を指定して、データフロー エッジのオーナー ターゲットを取得します。
getNonEdgeOwnerTargets
mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);
エッジ owner を指定して、データフロー エッジのオーナー以外のターゲットを取得します。
ShardingRuleOpInterface(Sdy_ShardingRuleOpInterface)
OP が独自のシャーディング ルールを定義できる OP インターフェース。シャーディング ルールでは、オペレーションのさまざまなプロパティ(属性、オペランドの形状、結果の形状など)に応じてオペレーションをパーティショニングする方法が指定されます。詳細については、OpShardingRuleAttr をご覧ください。
メソッド:
getShardingRule
mlir::sdy::OpShardingRuleAttr getShardingRule();
OP のシャーディング ルールを返します。