OpInterface 定義

CollectiveOpInterface (Sdy_CollectiveOpInterface)

所有集合運算的介面。封裝 outSharding 屬性的常見 get/set。

限制：

  • 運算元必須具有分割，否則 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)

此操作介面可讓 Shardy 透過擴充此介面的操作資料流邊緣，傳播區塊處理作業。

某個運算子 X 的資料流動邊緣會在一系列來源 (每個來源都是 X 的運算元或 X 的區塊終結運算元) 和一系列目標 (每個目標都是 X 的結果或 X 的區塊引數) 之間定義橋接點，以便所有來源和目標都以相同方式分割。一個運算可擁有多個彼此垂直的資料流邊緣。

擁有者是使用者指定的目標，用於 Shardy 的傳播所使用的資料流動邊緣。使用者可以任意選擇，但必須是靜態。

例如：

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

根據 transformType 轉換目標的 sharding

詳情請參閱 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 的情況下，取得資料流邊緣的非擁有者目標。

ShardingRuleOpInterface (Sdy_ShardingRuleOpInterface)

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

方法：

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

傳回作業的分割規則。