Определения OpInterface

CollectiveOpInterface ( Sdy_CollectiveOpInterface )

Интерфейс для всех коллективных операций. Инкапсулирует общий метод get/set для атрибута outSharding.

Ограничения:

  • Операнд должен иметь сегментирование, иначе 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 )

Интерфейс операции, который позволяет шарди распространять сегменты через границы потока данных операций, которые расширяют этот интерфейс.

Граница потока данных некоторой операции 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::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);

Преобразует sharding цели в зависимости от transformType

Дополнительные сведения см. в разделе 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 ребра.

shouldKeepEdgeOwnerShardingsDivisible

bool shouldKeepEdgeOwnerShardingsDivisible();

Возвращает true, если сегментирование владельцев ребер делит размер измерения, чтобы избежать необходимости заполнения.

ShardingRuleOpInterface ( Sdy_ShardingRuleOpInterface )

Интерфейс оператора, который позволяет оператору определять собственное правило сегментирования. Правило сегментирования определяет, как операция может быть разделена в соответствии с различными свойствами операции — любыми атрибутами, формой операндов, формой результатов и т. д. Дополнительные сведения см. в OpShardingRuleAttr .

Методы:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Возвращает правило сегментирования операции.

shouldKeepOutputShardingsDivisible

bool shouldKeepOutputShardingsDivisible();

Возвращает true, если выходное сегментирование делит размер измерения, чтобы избежать необходимости заполнения.