“mpmd”方言

操作

mpmd.assign (mpmd::AssignOp)

分配操作

语法：

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

将局部张量分配给网格，使其在该网格内完全复制。

这是在降低 JAX 操作时引入的临时操作，用于从本地类型移动到网格类型。当 func op 的输入和结果变为网格张量时，这些操作将在导入期间被消除。

结果类型的网格名称应与拓扑中的网格相对应，并且其全局类型应与操作数类型相同。

分配操作的来源是网格的来源，例如命名计算、网格推理等。

特征：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)

fragment 操作

将计算（即操作块）分配给 MPMD 拓扑中的特定网格，该网格旨在作为单独的 SPMD 程序片段执行。

该 fragment 仅接受和返回分配给与该 fragment 相同网格的网格张量。

片段的网格名称应与拓扑中的网格相对应。

该 fragment 包含一个来源列表，即包含有关构成此 fragment 的原始 named_computations 的信息的元数据，以及一个 staged_id（仅当它是用户定义的 fragment 时才定义）。也就是说，它具有非空的来源列表。可选的 in_shardings 指定了 fragment 的块实参（对应于操作数）的分片。 可选的 out_shardings 用于指定结果的分片。

fragment 的区域不应有任何自由变量，并且该区域中每个区块实参和返回值的类型都是相应网状网张量的全局张量类型。

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

接口：ConditionallySpeculatable、ShardableDataFlowOpInterface

属性：

实参：

结果：

mpmd.fragment_call (mpmd::FragmentCallOp)

fragment 调用操作

表示对包含 MPMD fragment body 的函数的调用，即分配给 MPMD 拓扑中特定网格的计算，旨在作为单独的 SPMD 程序 fragment 执行。

片段的网格名称应与封装函数的拓扑中的网格相对应，并且该网格形状应与被调用者的网格形状相匹配。

来源指定了促成此 fragment 调用的用户命名计算（例如通过合并）。

被调用方的函数输入和结果类型必须分别是相应操作数和结果的相应网状网张量的本地张量类型。

示例：

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

跨网格 reduce 操作

语法：

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 的输入和结果变为网格张量时，这些操作将在导入期间被消除。

操作数类型的网格名称应与拓扑中的网格相对应，并且其全局类型应与结果类型相同。

特征：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

网格张量类型

在程序的网格 MPMD 拓扑中，为特定 SPMD 网格分配 RankedTensorType。 该类型包含一个可选的分片，用于指定如何根据 SPMD 网格对张量进行分片。 如果不存在分片，则张量会完全复制。

参数：

枚举

ReductionType

减免类型属性

支持请求：

EdgeNodeType

边缘节点类型枚举

支持请求：

PropagationDirection

传播方向枚举

支持请求：

ReductionOp

缩减操作枚举

支持请求：