OpInterface 定義

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

此操作介面可讓 Shardy 透過擴充此介面的操作資料流邊緣，傳播分割作業。

某個運算 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 有兩種資料流邊緣類型，分別是 return_value_i (來源) 和 y_i (目標) 之間的 n 個邊，以及 x_i(來源) 和 body_arg_i(目標) 之間的 n 個邊。在這種情況下，邊緣擁有者與目標相同

以下是含有多個目標的作業範例：

  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
                  })

這個 while 運算子有 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);

使用指定的 index 設定區塊引數邊緣擁有者的 sharding。

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);

使用指定的 index 設定運算結果邊緣擁有者的 sharding。

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)

可讓運算作業定義自身資料分割規則的 O 介面。資料分割規則會指定如何根據運算上的各種屬性 (任何屬性、運算元形狀、結果形狀等) 為作業分區。詳情請參閱 OpShardingRuleAttr。

方法：

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

傳回作業的區塊處理規則。