-mpmd-copy-topology-from-main

将拓扑从主函数复制到 mpmd.call 引用的函数。

将拓扑属性从主函数复制到 mpmd.call 引用的任何函数。这还会将 mpmd.call 被调用方设置为私有可见性,以避免被误认为是入口点函数。

-mpmd-enforce-input-output-equisharding

对 MPMD 函数强制执行等分片约束。

通过在必要时引入 TransferOps,对 MPMD 函数强制执行输入-输出等分片约束。

选项

-constraints : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.

-mpmd-generate-sdy-meshes-from-topology

根据 MPMD 拓扑生成分片网格。

此传递会生成移除所有现有的分片网格操作,并将其替换为基于 MPMD 拓扑的操作。它还会更新张量分片,以引用新的网格操作。

-mpmd-infer-mesh-assign-mesh-func-leaves

_使用 use_set 和 srcset 分析,为每个未使用的计算、函数输出和函数输入分配网格。

此传递会使用 use_set 和 src_set 信息,通过创建 AssignOps 或更改类型,为函数正文叶(即未使用计算的结果、未使用的函数实参和函数输出)分配网格。

为了进行分析,我们还将某些中间值视为叶。具体而言:mpmd.reduce 和 mpmd.broadcast 的操作数被视为叶,并且系统会在其上创建分配-取消分配对。

此分配将清除所有非叶操作的 use_set,因为之前带注释的用途将过时,原因是推断 reduce 操作会更改某些值的 use_set:初始 use_set 传播不知道 reduce 操作,但现在我们已经推断出 reduce 操作,传播将有所不同。

如果未正确填充叶操作的 use 集和 src 集,此传递将失败并发出错误。

前提条件:每个操作都有一个非空 src 集,或者我们推断出转移。

选项

-infer-transfers : Whether to create transfers when needed, instead of erroring.
-error-limit     : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-infer-mesh-assign-using-input-output-constraints

根据输入-输出分配约束,为输入和输出分配网格。

此传递使用输入-输出等分配约束,将输入和输出分配给同一网格。

请注意,无论输入随后是否转移到其他网格,这都能保证输入位于同一网格上。但这表示我们不应在此传递后运行 populate-src-set

要求

  • 对于可能是等分配约束一部分的入口点函数的任何输入 ii 具有 MeshTensorType,或者具有明确定义的 use 集。
  • 对于可能是等分配约束一部分的入口点函数的任何输出 oo 具有 MeshTensorType 类型,或者具有明确定义的 src 集和 use 集。

其中,值的明确定义的 use 集包括通过 mpmd.assign 操作(而非其他网格)将值(传递)分配到的所有网格。明确定义的 src 集包括允许张量存在的所有网格,而不包括其他网格。

虽然这仅在入口点函数上运行,但我们将其作为模块操作传递,因为这需要先完成对现有函数的所有传递,然后才能运行。例如,如果我们将其作为 EntryPointFunctionPass,则传递管理器可能会在非入口点函数上的验证完成之前运行此传递。

选项

-verbose-logging : Whether to enable verbose logging
-constraints     : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.

-mpmd-infer-mesh-convert-reduce-ops

将带注释的 reduce 操作转换为 mpmd.reduce 操作,并展平 reduce 操作链。

将带注释的 reduce 操作转换为 mpmd.reduce 操作,并展平这些 reduce 操作链。

符号如下:

x = add(w0, w1) {mpmd.reduce = #mpmd.reduce} y = add(x, w2) {mpmd.reduce = #mpmd.reduce} ~~> r = mpmd.reduce(w0,w1,w2)

选项

-infer-cross-mesh-reductions : Whether to infer cross-mesh reductions. Will be enabled by default once stable.

-mpmd-infer-mesh-finalize

在模式网格推断后应用最终清理。

选项

-infer-transfers : Whether to create transfers when needed, instead of erroring.

-mpmd-infer-mesh-populate-src-set

_初始化 UnassignOps 和函数实参的 src_set,并传播 srcset。

此传递会初始化 src_set 并传播它,从而使用 src_set 信息填充图。

前提条件:如需让函数实参具有 src_set,必须填充 use_set。

初始化:UnassignOp 的 src_set 设置为其分配到的网格。函数实参的 src_set 设置为其 use_set。

传播:src_set 从操作数向前传播到操作本身,并取操作数的交集。如需了解详情,请参阅 PropagateSrcSet

-mpmd-infer-mesh-populate-use-set

_初始化 AssignOps 的 use_set,并向后传播 useset。

此传递会初始化 use_set 并向后传播它,从而使用 use_set 信息填充图。

初始化:AssignOp 的 use_set 设置为其分配到的网格。

传播:use_set 从用户向后传播到操作本身,并取用户的并集。根据定义,操作的 use_set 是其用户的 use_set 的并集,因为 use_set 是传递用途的集合。

-mpmd-infer-mesh-rewrite-using-analysis

_根据 useset. 重写操作。_

此传递使用 use_set 和 src_set 分析,通过将无网格操作封装在 fragment 中来分配这些操作。

它还会移除 use_set 和 src_set 属性作为清理的一部分,因为在此之后不再需要这些分析。

前提条件:每个操作都有一个 use_set,即分析已完成。 前提条件:非入口点函数的每个实参至少被一个非终止符操作使用。

待办事项:jupvfranco - 考虑到此传递不再那么依赖分析,请考虑重命名此传递。

选项

-max-clones : How many copies of a meshless operation we allow. Setting it to 1 means we never clone the op.

-mpmd-infer-mesh-validate-no-additional-transfers-needed

验证网格分配不需要额外的转移。

此传递验证是否可以为所有无网格操作分配网格,而无需引入任何额外的转移。

对于不是函数操作的无网格操作,如果出现以下情况,则会出错:

  1. use_set 不包含在给定操作的 src_set 中,即需要转移。

对于函数操作,只需检查函数实参的上述条件,因为函数返回无网格操作或块实参。

选项

-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-infer-mesh-validate-src-set-not-empty

_验证每个无网格操作都有一个非空 srcset。

此传递会验证所有无网格操作,检查是否可以将操作分配到某个位置。也就是说,对于不是函数操作的无网格操作,如果操作的 src_set 为空,或者如果推断出它是跨网格缩减但未转换,则会出错。这是函数叶分配的前提条件。

对于函数操作,只需检查函数实参的上述条件,因为函数返回无网格操作或块实参。

这需要在模块级进行传递,因为如果被调用方出现错误,我们希望打印调用方。

前提条件:在运行此传递之前,应将跨网格缩减转换为 reduce 操作。

选项

-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-inline-nested-user-exposed-ops

内联嵌套在命名计算中的任何用户公开的 mpmd 操作。

内联嵌套在命名计算中的任何命名计算、命名张量、广播和 reduce 操作,检查其网格分配(如果已定义)是否与其父项的网格分配匹配。

选项

-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.

-mpmd-insert-nameless-clone-of-negligible-ops

在命名计算之外克隆可忽略的操作。

在命名计算之外克隆可忽略的操作,即单个结果、零操作数操作,只要它们被计算的返回操作使用,就用克隆替换命名计算的结果。这是必需的,因为如果此类结果被分配给不同网格的命名计算使用,则可能会导致网格推断冲突。通过应用此传递,我们允许网格推断克隆这些可忽略的操作。

此传递不会更改命名计算。

-mpmd-introduce-transfers

根据用户网格分配创建数据转移。

创建一个传递,根据用户网状网分配引入转移操作。其中包括:

  1. 如果 UnassignOp 的结果稍后在被调用方中分配,则将 UnassignOp 推送到 mpmd 调用中。
  2. 将 UnassignOp 的 AssignOp 替换为 TransferOp。
  3. 如果 fragment 之间存在无网格加法,则将加法分配给使用网格并引入转移。

-mpmd-map-input-output-to-mesh

为函数输入和输出分配网格。

创建一个传递,根据用户定义的网格分配将函数输入/输出映射到网格。

对于输入实参,此传递:

  1. 将应放置在网格上的输入张量转换为网格张量。
  2. 更新函数签名。
  3. 在张量使用之前添加 mpmd.unassign。

对于输出实参,此传递会在返回张量之前添加 mpmd.assign 并更新函数签名。

要求:每个输入/输出索引都是有效的,并且每个映射的网格都是拓扑中的有效网格。

选项

-input-assignment  : Mapping between function input indices and assigned mesh names.E.g., '0@m0,1@m1' defines that input with index 0 will be assigned to mesh m0 and input with index 1 will be assigned to mesh m1.
-output-assignment : Mapping between function output indices and assigned mesh names.E.g., '0@m0,1@m1' defines that output with index 0 will be assigned to mesh m0 and output with index 1 will be assigned to mesh m1.

-mpmd-map-named-ops-to-mpmd-ops

为用户定义的操作分配网格。

创建一个传递,用于选择性地将 mpmd.named_tensor 分配给 Assign(Unassign(%v))(取决于 assignment 中是否有条目),并使用用户定义的命名计算与网格名称之间的映射将每个命名计算映射到网格。这意味着将每个命名计算替换为 fragment,并为这些 fragment 的操作数创建 AssignOps,为这些 fragment 的结果创建 UnassignOps。现在引入的模式 Assign(Unassign(%v)) 将重写为 Transfer(%v)。 在此传递之后,将不存在任何命名计算/命名张量操作。

要求:所有命名计算和命名张量都位于函数的顶层。

选项

-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.

-mpmd-simplify-named-computations

简化命名计算操作的输入和输出。

独立简化每个命名计算。具体而言,它:

  • 对结果及其对应的返回值进行重复数据删除;
  • 对操作数及其对应的块实参进行重复数据删除;
  • 移除其对应的返回操作数是操作的块实参的结果;
  • 移除其对应的块实参不再使用(或一开始就没有使用)的操作数;以及
  • 移除未使用的结果。
  • 将命名计算中的模式 arg -> stablehlo.optimization_barrier -> return 替换为模式 arg -> return,从而允许 进一步简化。

-mpmd-validate-named-ops-in-mpmd-func

验证命名操作是否仅嵌套在 mpmd 函数中。

验证 NamedComputationOp 和 NamedTensorOp 是否仅嵌套在 mpmd 函数中,即具有拓扑属性的函数。