'mpmd' Dialect

作業

mpmd.assign (mpmd::AssignOp)

指派作業

語法：

operation ::= `mpmd.assign` attr-dict $tensor `:` functional-type(operands, results)

將本機張量指派給網格，在該網格中完全複製。

這是降低 jax 作業時導入的暫時性作業，可從本機型別移至網格型別。匯入時，當 func op 的輸入和結果變成網格張量，這些 op 就會遭到淘汰。

結果類型的網格名稱應對應至拓撲中的網格，且其全域類型應與運算元類型相同。

指派作業的來源是網格的來源，例如 named_computation、網格推論等。

特徵：AlwaysSpeculatableImplTrait、HasParent<::mlir::func::FuncOp, ForOp>

介面：ConditionallySpeculatable、NoMemoryEffect (MemoryEffectOpInterface)

影響：MemoryEffects::Effect{}

屬性：

運算元：

成果：

mpmd.broadcast (mpmd::BroadcastOp)

廣播作業

語法：

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

允許在任何使用張量的網格中轉移 (或複製) 張量。無論何時轉移，轉移的來源都是運算元的目前位置。

特徵：AlwaysSpeculatableImplTrait、SameOperandsAndResultType

介面：ConditionallySpeculatable、InferTypeOpInterface、NoMemoryEffect (MemoryEffectOpInterface)

影響：MemoryEffects::Effect{}

運算元：

成果：

mpmd.call (mpmd::CallOp)

MPMD 專屬的呼叫函式

語法：

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

函式呼叫作業。例如，將迴圈主體包裝在函式宣告中，可縮減程式碼大小。

介面：ArgAndResultAttrsOpInterface、CallOpInterface、SymbolUserOpInterface

屬性：

運算元：

成果：

mpmd.for (mpmd::ForOp)

適用於運算子

傳回執行主體函式固定疊代次數的結果，主體中提供疊代索引。

您可以指定選用的展開係數 (必須能整除疊代次數)，依該係數展開運算元的主體，也就是說，如果展開係數為 N，主體會複製 N 份，疊代次數則會減少 1/N 倍。除了第一個副本外，每個副本都會使用前一個副本的結果，而非區塊引數，且疊代索引會乘以展開係數，並在每個副本後遞增。

for 運算子可以接受及傳回任何型別，但這些型別的 TypeID 必須相同，例如所有張量型別或所有 MPMD 網格型別等。這讓我們可以在各種層級使用運算子，共用實作項目和轉換。

特徵：HLO_PairwiseSameOperandAndResultType、RecursiveMemoryEffects、RecursivelySpeculatableImplTrait、SingleBlockImplicitTerminator<ReturnOp>、SingleBlock

介面：ConditionallySpeculatable、LoopLikeOpInterface、OpAsmOpInterface、ShardableDataFlowOpInterface

屬性：

運算元：

成果：

mpmd.fragment (mpmd::FragmentOp)

片段作業

將運算 (即一連串的作業) 指派給 MPDM 拓撲中的特定網格，目的是要將運算做為個別的 SPMD 程式片段執行。

片段只會接收和傳回指派給與片段相同網格的網格張量。

片段的網格名稱應對應拓撲中的網格。

片段包含來源清單，也就是含有資訊的中繼資料，這些資訊與構成這個片段的原始 named_computations 相關，以及定義 iff 的 staged_id (如果是使用者定義的片段，也就是含有非空白的來源清單)。選用的 in_shardings 會指定片段的區塊引數分片，對應至運算元。選用的 out_shardings 會指定結果的分片。

片段的區域不應有任何自由變數，且區域中每個區塊引數和傳回值的型別，都是對應網格張量的全域張量型別。

特徵：HasParent<::mlir::func::FuncOp, ForOp>、IsolatedFromAbove、RecursiveMemoryEffects、RecursivelySpeculatableImplTrait、SingleBlockImplicitTerminator<ReturnOp>、SingleBlock

介面：ConditionallySpeculatable、ShardableDataFlowOpInterface

屬性：

運算元：

成果：

mpmd.fragment_call (mpmd::FragmentCallOp)

片段呼叫作業

代表對函式的呼叫，該函式會保留 MPMD 片段主體，也就是指派給 MPMD 拓撲中特定網格的計算，目的是要以個別 SPMD 程式片段的形式執行。

片段的網格名稱應對應至封閉函式拓撲中的網格，且該網格形狀應與呼叫端相符。

來源會指定促成這個片段呼叫的使用者命名計算，例如透過合併。

被呼叫端的函式輸入和結果型別，必須分別是這個運算元和結果的對應網格張量的本機張量型別。

範例：

%2 = mpmd.fragment_call<mesh="m1",> origin=[] @my_fragment(%0, %1) :
  (m<esh>_tensor..., m<esh>_te>nsor...) - m<esh>_tensor...

特徵：HasParent<::mlir::func::FuncOp>、MemRefsNormalizable

介面：ArgAndResultAttrsOpInterface、CallOpInterface、SymbolUserOpInterface

屬性：

運算元：

成果：

mpmd.named_computation (mpmd::NamedComputationOp)

具名範圍作業

將計算 (即一連串的作業) 分組，並透過 UserOrigin 屬性為其命名及提供轉置計數。這個 NamedComputation 可用於在 MPMD 中將網格指派給運算，或用於最佳化。

轉置計數 (預設值為 0) 表示具名計算是否由特定數量的 JAX AD 轉置轉換產生。

運算元的區域不應有任何自由變數，且區域中每個區塊引數和傳回值的型別，必須與運算元的輸入型別和傳回類型相同。

特徵：IsolatedFromAbove、RecursiveMemoryEffects、RecursivelySpeculatableImplTrait、SingleBlockImplicitTerminator<ReturnOp>、SingleBlock

介面：ConditionallySpeculatable

屬性：

運算元：

成果：

mpmd.named_tensor (mpmd::NamedTensorOp)

將張量指派給網格

語法：

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

將張量結果與指定名稱建立關聯的身分運算。這個 NamedTensor 可用於在 MPMD 中將網格指派給張量。

特徵：AlwaysSpeculatableImplTrait、SameOperandsAndResultType

介面：ConditionallySpeculatable、InferTypeOpInterface、NoMemoryEffect (MemoryEffectOpInterface)

影響：MemoryEffects::Effect{}

屬性：

運算元：

成果：

mpmd.reduce (mpmd::ReduceOp)

跨網格減少作業

語法：

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

允許張量在不同網格間縮減，然後廣播至需要使用的位置。

特徵：AlwaysSpeculatableImplTrait、SameOperandsAndResultType

介面：ConditionallySpeculatable、NoMemoryEffect (MemoryEffectOpInterface)

影響：MemoryEffects::Effect{}

屬性：

運算元：

成果：

mpmd.return (mpmd::ReturnOp)

mpmd.return 作業會終止附加至 mpmd 區域的區域作業。這是可變長度引數：它會將值清單做為引數 (型別可以是任何型別，但必須是相同種類，例如 AnyTensor)，因此可在 MPMD IR 堆疊的各個層級重複使用。

語法：

operation ::= `mpmd.return` attr-dict $results (`:` type($results)^)?

特徵：AlwaysSpeculatableImplTrait、Terminator

介面：ConditionallySpeculatable、NoMemoryEffect (MemoryEffectOpInterface)

影響：MemoryEffects::Effect{}

運算元：

mpmd.transfer (mpmd::TransferOp)

轉移作業

語法：

operation ::= `mpmd.transfer` attr-dict $tensor `:` functional-type(operands, results)

將分散式張量從一個網格轉移至另一個網格。

運算元和結果類型的網格名稱應對應至拓撲中的網格，且兩者的全域類型應相同。

特徵：AlwaysSpeculatableImplTrait、HasParent<::mlir::func::FuncOp>

介面：ConditionallySpeculatable、NoMemoryEffect (MemoryEffectOpInterface)、ShardingRuleOpInterface

影響：MemoryEffects::Effect{}

運算元：

成果：

mpmd.unassign (mpmd::UnassignOp)

取消指派作業

語法：

operation ::= `mpmd.unassign` attr-dict $tensor `:` functional-type(operands, results)

從網格取消指派完全複製的張量。

這是降低 jax 作業時導入的暫時性作業，可從本機型別移至網格型別。匯入時，當 func op 的輸入和結果變成網格張量，這些 op 就會遭到淘汰。

運算元類型的網格名稱應對應至拓撲中的網格，且其全域類型應與結果類型相同。

特徵：AlwaysSpeculatableImplTrait、HasParent<::mlir::func::FuncOp, ForOp>、InferTensorType

介面：ConditionallySpeculatable、InferShapedTypeOpInterface、InferTypeOpInterface、NoMemoryEffect (MemoryEffectOpInterface)

影響：MemoryEffects::Effect{}

屬性：

運算元：

成果：

屬性

MeshWithOriginsAttr

與來源合併。

語法：

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

參數：

MeshesWithOriginsAttr

網格清單及其來源。

語法：

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

參數：

NamedMeshAttr

一對名稱和網格。

語法：

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

參數：

OriginAttr

網格指派的來源。

語法：

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

網格指派作業的來源。

origin_label 是指來源的人類可讀標籤。此方法用於偵錯。

參數：

ReductionAttr

表示減少。

語法：

#mpmd.reduction<
  `none` | `add` | `max` | `min` | `mul` | `or` | `and`   # reduction_type
>

參數：

TopologyAttr

具名網格的拓撲。

語法：

#mpmd.topology<
  ::llvm::ArrayRef<NamedMeshAttr>   # meshes
>

參數：

UserOriginAttr

使用者指定運算的來源。

語法：

#mpmd.user_origin<
  ::mlir::StringAttr,   # userName
  int64_t   # transposeCount
>

參數：

類型

MeshTensorType

網格張量類型

在程式的 MPMP 網格拓撲中，將 RankedTensorType 指派給特定 SPMD 網格。這個型別會保留選用的分片，指定張量相對於 SPMD 網格的分片方式。如果沒有分片，張量會完全複製。

參數：

列舉

ReductionType

減免類型屬性

案件：

EdgeNodeType

邊緣節點類型列舉

案件：

PropagationDirection

Propagation direction enum

案件：

ReductionOp

減少 op 列舉

案件：