ShardableDataFlowOpInterface（ShardableDataFlowOpInterface）
このインターフェースを拡張するオペレーションのデータフロー エッジを介して、シャーディングを伝播できるようにするオペレーション インターフェース。
ある演算 X のデータフロー エッジは、一連のソース（それぞれ X のオペランドまたは X のブロック終端子のオペランド）と一連のターゲット（それぞれ X の結果または X のブロック引数）間のブリッジを定義します。これにより、すべてのソースとターゲットが同じ方法でシャーディングされるようになります。op は、互いに直交する複数のデータフロー エッジを持つことができます。
オーナーは、シャーディーのプロパゲーションによって使用されるデータフロー エッジのユーザー指定のターゲットです。ユーザーが任意に選択できますが、静的である必要があります。
例:
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::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();
すべてのブロック引数データフロー エッジ オーナーのシャーディングを返します。
setBlockArgumentEdgeOwnerSharding
void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);
指定された index で、ブロック引数のエッジオーナーの sharding を設定します。
setBlockArgumentEdgeOwnerShardings
void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
すべてのブロック引数のエッジ オーナーの shardings を設定します。
getOpResultEdgeOwnerShardings
mlir::ArrayRef<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();
すべてのオペレーション結果データフロー エッジ オーナーのシャーディングを返します。
setOpResultEdgeOwnerSharding
void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);
op 結果エッジ オーナーの sharding を指定された index に設定します。
setOpResultEdgeOwnerShardings
void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
すべてのオペレーション結果エッジ オーナーの shardings を設定します。
getBlockArgumentEdgeOwners
mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();
すべてのブロック引数のエッジ オーナーを取得します。
getOpResultEdgeOwners
mlir::ResultRange getOpResultEdgeOwners();
すべてのオペレーション結果エッジ オーナーを取得します。
getEdgeSources
mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);
target 値を指定して、データフロー エッジソースを取得します。
getEdgeOwnerFromTarget
mlir::Value getEdgeOwnerFromTarget(mlir::Value target);
所有者である可能性のある target を指定して、データフロー エッジの所有者 target を取得します。
getEdgeOwnerFromSource
mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);
指定された source で指定された Dataflow エッジのオーナー ターゲットを取得します。
getNonEdgeOwnerTargets
mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);
エッジ owner を指定して、データフロー エッジのオーナー以外のターゲットを取得します。
ShardingRuleOpInterface（ShardingRuleOpInterface）
演算で独自のシャーディング ルールを定義できる演算インターフェース。
シャーディング ルールでは、オペレーションのさまざまなプロパティ（属性、オペランドの形状、結果の形状など）に応じてオペレーションをパーティショニングする方法が指定されます。詳細については、OpShardingRuleAttr をご覧ください。
メソッド:
getShardingRule
mlir::sdy::OpShardingRuleAttr getShardingRule();
演算のシャーディング ルールを返します。