Definicje interfejsu OpInterface

Zadbaj o dobrą organizację dzięki kolekcji Zapisuj i kategoryzuj treści zgodnie ze swoimi preferencjami.

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.