-sdy-add-data-flow-edges

为每个数据流边插入 DataFlowEdgeOp

为数据流边的所有所有者值(即 getDataFlowEdgeOwners 针对模块中的每个操作返回的所有值)插入 DataFlowEdgeOp

插入的 DataFlowEdgeOp 将采用所有者目标的现有分片(如果存在)。

-sdy-apply-sharding-constraints

应用用于决定其输入分片的约束条件。

如果 ShardingConstraintOp 满足以下所有条件,则将其分片复制到其输入:

  • 输入没有现有的分片。
  • ShardingConstraintOp 的分片已完全关闭。
  • 输入中没有其他类型为 ShardingConstraintOpManualComputationOp 且采用不同分片的用户。

这些条件表明 ShardingConstraintOp 决定了其输入的分片方式。

如果输入由 DataFlowEdgeOp 生成,则我们会将所有 input 用法替换为 ShardingConstraintOp,而不是设置操作的分片,以避免限制边缘的所有目标的分片。

请注意,无论是否进行此传递,ShardingConstraintOp 的分片都会在传播期间传播到其输入或用户,但由于维度的 closed 属性不会传播,因此请务必复制分片,以便在上述情况下完全遵循约束条件。

此外,如果张量被满足以下所有条件的 ShardingConstraintOp 链使用:

  • 该张量不是由 ShardingConstraintOp 生成的,并且没有任何其他类型为 ShardingConstraintOpManualComputationOp 的其他用户;
  • 链中的所有 ShardingConstraintOp(除了最后一个)都不会重复使用;
  • 链中的最后一个 ShardingConstraintOp 没有任何类型为 ShardingConstraintOpManualComputationOp 的用户(否则,它不是链中的最后一个);

然后,此传递会将链的输入的所有其他用法(在链中的最后一个 ShardingConstraintOp 之后 [并在同一代码块内] 定义)替换为链的结果,因为它应决定这些用法的分片。

-sdy-constant-splitter

拆分常量子计算,以便每个子计算都只使用一次。

拆分常量子计算,使其具有单个用户。

这可确保在常量子计算的不同用法之间不会传播分片,因为这被视为虚假依赖项(常量的用法不应仅因使用相同的常量而以相同的方式进行分片)。实际上,每种用法可以具有不同的分片,这些分片可以独立传播到常量子计算的各自副本。

常量子计算可以是:

  • 常量或 iota 运算(无运算数)
  • 广播、切片或纯元素级运算,其操作数都由常量子计算(递归)定义,以及定义其操作数的整个子计算。

请注意,在常量子计算中,某个值可以在该子计算中多次使用。

-sdy-lift-inlined-meshes

将分片中的内嵌 MeshAttr 提升为符号 MeshOp

TensorShardingAttr 中的任何内嵌 MeshAttr 替换为网格符号名称,引用模块中的现有或新 MeshOp,以便没有两个 MeshOp 具有相同的 MeshAttr(现有 MeshOp 也会被删除重复项)。

每个新 MeshOp 的名称将是:

  • maximal_mesh_{device-id}(对于最大网格,即空轴列表和单个设备 ID),或
  • [meshmesh_0mesh_1、…] 中第一个可用的名称。

-sdy-manual-axes-cleanup

清理了 ManualComputationOp 中手动轴的使用

  1. 对于未指定手动轴的任何入/出分片,请将该手动轴添加到其 replicated_axes。这是为了确保始终完全指定手动轴。
  2. 按网格轴声明顺序对手动轴进行排序。

-sdy-sharding-group-import

对分片组进行了规范化和验证。

在导入时对分片群组应用规范化和验证。具体包括:

  1. 分片群组统一

    使用组成员资格的传递性属性组合分片组。每当张量 T 同时位于分片组 G1 分片组 G2 中时,我们都可以推断出 G1 和 G2 中的所有成员都应采用相同的方式进行分片。因此,我们可以将 G1 和 G2 合并为一个组。合并后,一组最小组的规范组 ID 将为 0,1,...N-1。

  2. 分片群组验证

    验证分片组的格式是否正确,以及是否符合实现中的假设。目前,此断言表示,如果分片组包含在 ManualComputationOp 块中定义的 Value,则该组中的所有其他值都必须位于同一块中。