ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)
Interfejs operacji, który umożliwia cząstkom propagowanie podziału za pomocą krawędzi przepływu danych w operacjach rozszerzających ten interfejs.
Krawędzia przepływu danych w operacji X definiuje most między zbiorem źródeł (każde z nich jest albo operandem X, albo operandem terminatora bloku X) a zbiorem docelów (każdy z nich jest albo wynikiem X, albo argumentem bloku X), tak aby wszystkie źródła i docelowe powinny być dzielone w taki sam sposób. Operacja może mieć wiele krawędzi przepływu danych, które są do siebie prostopadłe.
Właściciel to określony przez użytkownika cel krawędzi przepływu danych używanej przez propagowanie fragmentów. Użytkownik może wybrać dowolny kolor, ale musi być on statyczny.
Na przykład:
y_1, ..., y_n = custom_op (x_1, ..., x_n)
((body_arg_1,..., body_arg_n) {
...
return return_value_1, ..., return_value_n
})
Ta operacja niestandardowa ma 2 typy krawędzi przepływu danych: n krawędzi między return_value_i (źródła) a y_i (docelowe) oraz n krawędzi między x_i(źródła) a body_arg_i(docelowe). W tym przypadku właściciele krawędzi są tacy sami jak cele.
Oto przykład operacji z większą liczbą celów:
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
})
Ta operacja while ma n krawędzi przepływu danych, a i-ta krawędź przepływu danych znajduje się między źródłami x_i i return_value_i oraz celami y_i, pred_arg_i i body_arg_i.
Metody:
getBlockArgumentEdgeOwnerShardings
mlir::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();
Zwraca podział wszystkich właścicieli krawędzi przepływu danych argumentu bloku.
setBlockArgumentEdgeOwnerSharding
void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);
Ustawia wartość sharding argumentu bloku edge owner na podany index.
setBlockArgumentEdgeOwnerShardings
void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
Ustawia wartość shardings wszystkich właścicieli krawędzi argumentu bloku.
getOpResultEdgeOwnerShardings
mlir::ArrayRef<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();
Zwraca fragmenty kodu ze wszystkich właścicieli brzegów przepływu danych wyniku operacji.
setOpResultEdgeOwnerSharding
void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);
Ustawia sharding właściciela krawędzi wyniku op za pomocą podanego index.
setOpResultEdgeOwnerShardings
void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);
Ustawia shardings wszystkich właścicieli krawędzi wyników op.
getBlockArgumentEdgeOwners
mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();
Pobiera wszystkich właścicieli krawędzi argumentu bloku.
getOpResultEdgeOwners
mlir::ResultRange getOpResultEdgeOwners();
Pobiera wszystkich właścicieli krawędzi wyniku operacji.
getEdgeSources
mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);
Pobiera źródła brzegowe przepływu danych o danej wartości target.
getEdgeOwnerFromTarget
mlir::Value getEdgeOwnerFromTarget(mlir::Value target);
Pobiera właściciela target krawędzi przepływu danych podany jako target, który może, ale nie musi być właścicielem.
getEdgeOwnerFromSource
mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);
Pobiera właściciela celu krawędzi przepływu danych na podstawie source.
getNonEdgeOwnerTargets
mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);
Pobiera cele niebędące właścicielami krawędzi przepływu danych, biorąc pod uwagę krawędź owner.
ShardingRuleOpInterface (ShardingRuleOpInterface)
Interfejs operatora, który umożliwia mu zdefiniowanie własnej reguły dzielenia.
Reguła podziału określa, jak operacja może być dzielona na podstawie różnych właściwości operacji, np. atrybutów, kształtu operandów, kształtu wyników itp. Więcej informacji znajdziesz w artykule OpShardingRuleAttr.
Metody:
getShardingRule
mlir::sdy::OpShardingRuleAttr getShardingRule();
Zwraca regułę podziału operacji.