-sdy-add-data-flow-edges

为每个数据流边插入 DataFlowEdgeOp

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

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

TODO(b/330339693):移除 getDataFlowEdgeOwners 后更新此文档。

-sdy-apply-sharding-constraints

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

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

  • 输入没有现有的分片。
  • 输入不是由 DataFlowEdgeOp 生成的,它存放着边缘的所有目标的分片。
  • ShardingConstraintOp 的分片已完全关闭。
  • 输入没有任何其他具有其他分片的 ShardingConstraintOpManualComputationOp 用户。

这表示 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) 对于任何未指定手动轴的输入/输出分片,请将该手动轴添加到其 duplicate_axes。这是为了确保始终完全指定手动轴。

2) 按网格轴声明顺序对手动轴进行排序。

-sdy-sharding-group-import

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

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

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

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