ShardableDataFlowOpInterface ( ShardableDataFlowOpInterface )
Интерфейс операции, который позволяет шарди распространять сегменты через границы потока данных операций, которые расширяют этот интерфейс.
Граница потока данных некоторой операции X определяет мост между набором источников (каждый является либо операндом X, либо операндом терминатора блока X) и набором целей (каждый является либо результатом X, либо аргументом блока X). X), так что все источники и цели должны быть сегментированы одинаково. Операция может иметь несколько ребер потока данных, ортогональных друг другу.
Владелец — это указанная пользователем цель границы потока данных, используемая при распространении шарди. Пользователь может выбрать его произвольно, но он должен быть статичным.
Например:
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 имеет два типа ребер потока данных: по n ребер между return_value_i (источниками) и y_i (целями) и n ребер между x_i (источниками) и body_arg_i (целями). В этом случае владельцы ребер такие же, как и цели.
Вот пример операции с несколькими целями:
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
})
При этом, хотя op имеет 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);
Устанавливает sharding владельца края аргумента блока с заданным index .
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);
Устанавливает 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 .
getNonEdgeOwnerTargets
mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);
Получает целевые объекты ребра потока данных, не являющиеся владельцами, с учетом owner ребра.
ShardingRuleOpInterface ( ShardingRuleOpInterface )
Интерфейс оператора, который позволяет оператору определять собственное правило сегментирования. Правило сегментирования определяет, как операция может быть разделена в соответствии с различными свойствами операции — любыми атрибутами, формой операндов, формой результатов и т. д. Дополнительные сведения см. в OpShardingRuleAttr .
Методы:
getShardingRule
mlir::sdy::OpShardingRuleAttr getShardingRule();
Возвращает правило сегментирования операции.