Definicje interfejsu OpInterface

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Interfejs do wszystkich operacji zbiorczych. Zawiera wspólne operacje get/set dla atrybutu outSharding.

Ograniczenia:

  • Operand musi mieć podział, w przeciwnym razie funkcja allowMissingInputSharding() zwróci wartość true.
  • out_sharding jest prawidłowy w powiązaniu z odpowiednim typem.
  • Jeśli funkcja allowDifferentMeshes() zwraca wartość false, operacja i podzielenie wyników muszą mieć tę samą siatkę.
  • Ten sam ranking dla operandu i podzielenia wyników.

Metody:

getOutSharding

::mlir::sdy::TensorShardingAttr getOutSharding();

Zwraca podział tensora wyjściowego zbiorczego działania.

setOutShardingAttr

void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);

Określa podział tensora wyjściowego zbiorczego działania.

getTensor

::mlir::TypedValue<::mlir::TensorType> getTensor();

Pobierz operand tensora operacji zbiorczej.

getType

::mlir::Type getType();

Pobiera typ zbiorczego wyniku operacji.

allowDifferentMeshes

bool allowDifferentMeshes();

Wskazanie, czy operacja zbiorcza zezwala na podzielenie danych wejściowych i wyjściowych na różne siatki.

allowMissingInputSharding

bool allowMissingInputSharding();

Wskazanie, czy operacja zbiorcza pozwala na brak podziału danych wejściowych, czyli czy są one domyślnie w pełni replikowane.

ShardableDataFlowOpInterface (Sdy_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) i zbiorem docelów (każdy z nich jest albo wynikiem X, albo argumentem bloku X), tak aby wszystkie źródła i docelowe miały taką samą liczbę fragmentów. 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ą identyczni z docelowymi.

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_ireturn_value_i oraz celami y_i, pred_arg_ibody_arg_i.

Metody:

getBlockArgumentEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Zwraca podział wszystkich właścicieli krawędzi przepływu danych argumentu bloku.

setBlockArgumentEdgeOwnerShardings

void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Ustawia wartość shardings wszystkich właścicieli krawędzi argumentu blokowania.

getOpResultEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

Zwraca podział wszystkich właścicieli krawędzi przepływu danych z wynikiem operacji.

setOpResultEdgeOwnerShardings

void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Ustawia shardings wszystkich właścicieli krawędzi wyniku op.

transformTargetSharding

mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);

Przekształca sharding celu w zależności od transformType

Aby dowiedzieć się więcej, przeczytaj DataFlowShardingTransformType.

getBlockArgumentEdgeOwners

mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();

Pobiera wszystkich właścicieli krawędzi argumentu blokady.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Pobiera wszystkich właścicieli krawędzi wyniku operacji.

getEdgeSources

mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);

Pobiera źródła krawędzi przepływu danych, podając krawędzie owner.

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 (Sdy_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, takich jak atrybuty, kształt operandów, kształt wyników itp. Więcej informacji znajdziesz w artykule OpShardingRuleAttr.

Metody:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Zwraca regułę podziału na części op.