OneFlow源码阅读8：eager模式下的SBP Signature推导

oneflow 的 global tensor 有两个必要属性：

placement：决定了 tensor 数据分布在哪些设备上。
sbp：决定了 tensor 数据在这些设备上的分布方式。例如：
- split：将切分后的不同部分放到不同设备；同时指定切分的 axis。
- broadcast：将数据复制到各个设备。

如果参与运算的 tensor 的 sbp 不一样，结果 tensor 的 sbp 是什么呢？例如下面的代码：

# export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=0 LOCAL_RANK=0
# export MASTER_ADDR=127.0.0.1 MASTER_PORT=17789 WORLD_SIZE=2 RANK=1 LOCAL_RANK=1
import oneflow as flow

P0 = flow.placement("cpu", ranks=[0, 1])

t1 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.split(0))
# t1 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.broadcast)
t2 = flow.Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], placement=P0, sbp=flow.sbp.split(1))
t3 = t1 + t2
# oneflow.placement(type="cpu", ranks=[0, 1])
print(t3.placement)
# (oneflow.sbp.split(dim=0),)
print(t3.sbp)

t1和t2是分布在相同设备上的两个 tensor。t1.sbp是S(0)，在行上切分；t2.sbp是S(1)，在列上切分。
计算结果t3的 sbp 不需要用户手动指定，系统可以自动推导出t3.sbp是S(0)。这个过程中的一个核心步骤，就是 SBP Signature 的推导。

1 SBP Signature

本节以下文字摘自SBP Signature的官方文档：

对于一个孤立的 Tensor，我们可以随意设置它的 SBP 属性。但是，对于一个有输入、输出数据的算子，我们却不可以随意设置它的输入、输出的 SBP 属性。这是因为随意设置一个算子输入输出的 SBP 属性，可能不符合全局视角下算子的运算法则。
对于某个算子，其输入输出的一个特定的、合法的 SBP 组合，称为这个算子的一个 SBP Signature。
算子作者根据算子的运算法则，在开发算子时，就已经罗列并预设好该算子所有可能的 SBP Signature。
某一层算子只要有输入的 SBP，OneFlow 就可以根据 SBP Signature 推导出该层算子输出的 SBP。
所谓的 SBP Signature 自动推导，指的是：在给定所有算子的所有合法的 SBP Signature 的前提下，OneFlow 有一套算法，会基于传输代价为每种合法的 SBP Signature 进行打分，并选择传输代价最小的那个 SBP Signature。这样使得系统的吞吐效率最高。
如果 OneFlow 自动选择的 SBP Signature，上一层算子的输出与下一层算子的输入的 SBP 属性不匹配时，那怎么办呢？OneFlow 会检测到这种不一致，并且在上游的输出和下游的输入间插入一个算子，做相关的转换工作。这类自动加入做转换的算子，就称为 Boxing 算子。

总结一下，SBP Signature 的要点如下

它是针对算子的。
它包括算子的全部输入、输出的sbp。缺少（部分）输入，或（部分）输出，不能构成签名。
- 所以SbpSignature.bn_in_op2sbp_parallel是一个map结构，key就是各个input和output的标识。
输入与输出的sbp组合，在算子的运算法则下必须是合法的。
算子的作者需要列出合法 SBP Signature 的候选集。
如果推导出的SBP Signature，sbp与inputs不一致，会通过GetBoxingOutput转换为sbp一致的tensor。

1.1 NdSbp 及 NdSbpSignature

以下文字主要参考官方文档2D SBP。

我们可以通过ranks=[0, 1, 2, 3]指定tensor的数据分布在这4个设备上。这4个设备组成了一个一维的设备向量。对应的 SBP 如split(1)，是单个值，即 1D SBP。

Tensor 数据的分布也可以指定为ranks=[[0, 1], [2, 3]]。四个计算设备被划分为2x2的设备阵列。这时，SBP 也必须与之对应，是一个长度为 2 的数组。对应的NdSbp.sbp_parallel的类型就是数组。

例如sbp = (broadcast, split(0))。这个 2D SBP 的含义是：

在 ranks 的第一维度执行广播，将数据分别拷贝到[0, 1]和[2, 3]。
在 ranks 的第二维度分别执行split(0)。
- 例如，对于[0, 1]这个sub group，将上一步中分配给它的数据按行拆分给0和1。

如果 Tensor 的数据分布形式是多维的，如[[0, 1], [2, 3]]，算子对应的 SBP Signature 也是多维的，所以NdSbpSignature中，每个 input/output 对应的 sbp_parallel 都是数组。

2 placement.hierarchy

placement 对应的 C++ 类型是ParallelDesc。
构造 placement 的 ranks 可以是多维数组，表示设备的多维阵列分布。placement.hierarchy存储 ranks 在各个维度的 size 信息。

hierarchy 数组的长度是 ranks 的维数。
hierarchy 数组的元素值，是 ranks 对应维度的 size。
构造 hierarchy 的 C++ 代码可参考GetRanksShape。

运行下面的代码可以观察 hierarchy 的值。

import oneflow as flow

placements = [
    flow.placement("cpu", ranks=[ 0, 1, 2,   3, 4, 5]),
    flow.placement("cpu", ranks=[[0, 1, 2], [3, 4, 5]]),
]
for p in placements:
    print(p.hierarchy)
# outputs:
# [6]
# [2, 3]

3 tensor add 是哪个算子？

为了提高性能，从v0.8.0开始，Tensor 的接口基本都通过 C API 提供给Python。

PyTensorObject_methods中定义了很多 Tensor 方法。不过，add 方法是通过 Python C API 的 number protocol 实现的，指定 PyTensorObject_nb_add 实现加法操作，实际由functional::add实现。

functional::add的定义在functional_api.yaml.pybind.cpp中，这是一个在构建期自动生成的文件。
顺着这个找，容易发现示例代码对应的是AddFunctor。
Op的名字是"add_n"，
自动生成的文件op_generated.cpp中定义了add_n对应的Op是AddNOp。add_n_op.cpp中定义的几个函数，会在 SBP Signature 推导过程中用到。

4 一维 SBP 的推导过程

SBP Signature 推导相关的类关系如下：

示例代码中的 tensor add 操作（t1 + t2），执行到 Interpret 中调用GetOrInfer时，会进行 SBP Signature 的推导。
在ConsistentTensorInferCache::GetOrInfer中，会把推导结果存起来，不需要每次都进行推导。

ConsistentTensorMetaInferArgs的 hash 函数主要依赖输入 tensor 的如下信息：

shape
dtype
nd_sbp
placement
consumer_nd_sbp_constraint
不同的 tensor 对象，只要这些元信息相同，就可以复用同一个推导结果。

UserOpExpr通过ConsistentTensorInferCache持有所有推导过的结果。

4.1 ConsistentTensorInferCache 中的推导准备

实际的推导在ConsistentTensorInferCache::Infer中进行。

4.1.1 推导 output 的 shape 和 dtype

user_op_expr.InferLogicalTensorDesc的作用主要是推导 output 的 shape 和 data_type，结果保存到 output_mut_metas。

这里涉及到 UserOpExpr 和 Op 两个模块之间的交互关系。后面会总结一下几个模块之间的契约协议。
user_op_expr.InferLogicalTensorDesc中用到的两个函数对象，是在构造时从OpRegistryResult中拷贝的。OpRegistryResult 的函数对象来自 Op 注册。示例代码中 tensor add 对应的 Op 是 AddNOp。

AddNOp 场景的实际调用顺序示例如下：

user_op_expr.InferLogicalTensorDesc
- tensor_desc_infer_fn_ -> AddNOp::InferLogicalTensorDesc
  - [out.shape = in[0].shape](https://github.com/Oneflow-In...)
- dtype_infer_fn_ -> AddNOp::InferDataType
  - [out.data_type = in[0].data_type](https://github.com/Oneflow-In...)

4.1.2 构造 UserOp

MakeOp(user_op_expr...)返回一个Operator，具体类型是UserOp（参考之前静态图的讨论）。这个对象负责执行具体的推导。

CheckInputParallelDescIdentical要求所有 inputs 的 placement 是一致的。因为这里是针对UserOp做的推导，例如 tensor add、matmul 等操作，操作数都在相同的设备时，这些操作才能直接计算，否则，就需要通过系统 Op 将数据搬运到一起，再进行计算。

既然所有 inputs 的 placement 都是一样的，那就用第一个作为代表，并赋值给 UserOp 保存。

op->InferParallelSignatureIf()的作用是将 placement 填充到op.bn2parallel_desc_。
对于 AddNOp 来说，key是in_0, in_1, out_0，value 是 inputs[0].placement。

infer_args.MakeInputBlobDescs操作用伪码表示如下：

# for each input index i
blob_descs[i].shape = inputs[i].shape
blob_descs[i].stride = inputs[i].stride
blob_descs[i].data_type = inputs[i].data_type

infer_args.MakeNdSbpInferHints操作用伪码表示如下：

# for each input index i
hints[i].parallel_desc = inputs[i].parallel_desc
hints[i].blob_desc = blob_descs[i]
hints[i].nd_sbp = inputs[i].nd_sbp

blob_descs的作用是为了构造pd_infer_hints，pd_infer_hints是为了构造NdSbpInferHint4Ibn，将相关信息封装到这个函数对象中。这个函数对象被传递给UserOp进行推导。在UserOp中，通过这个函数对象，根据input/output的标识bn（blob name），获取NdSbpInferHint，从而可以得到上述元信息。

UserOp推导完毕后，ConsistentTensorInferCache会将 inputs/outputs 的元信息，连同推导得到的 NdSbp ，一起保存到ConsistentTensorInferResult。

4.2 Operator 中的推导准备

Operator::InferNdSbpSignatureIf中，调用InferNdSbpSignature进行实际的推导，然后调用FillNdSbpSignature保存推导结果。

InferNdSbpSignature是一个虚函数。UserOp会先检查Op有没有定义自己的 SBP Signature 推导函数，AddNOp 没有这方面的函数，就调用 Operator::InferNdSbpSignature。

InferNdSbpSignature 中会根据 parallel_desc.hierarchy() 判断是 1D SBP，还是 ND SBP。
先只看 1D SBP 的情况。调用传入的 NdSbpInferHint4Ibn 函数对象，查到 ConsistentTensorInferCache 中创建的 NdSbpInferHint，转为 NdSbpInferHint 并存到 map 中。因为是一维的，所以只需要取 sbp_parallel 的第一个元素。然后调用 InferSbpSignature（名字中少了 Nd），将推导结果写到 SbpSignature。
无论是一维还是多维，结果的类型都是 NdSbpSignature。所以要将 SbpSignature 转为 NdSbpSignature。

Operator::InferSbpSignature的作用主要是构造两个函数对象，SbpInferHint4Ibn 和 CalcOrderValue4SbpSig，然后调用子类 override 的、同名重载的虚函数 InferSbpSignature。
SbpInferHint4Ibn 是将传入的 map 数据封装到函数对象中，用于查询输入输出的元信息。
CalcOrderValue4SbpSig 给每个 SbpSignature 计算一个序值，用于对签名进行排序。

InferSbpSignature 也是一个虚函数。因为 AddNOp 没有定义签名推导函数，会调用 Operator::InferSbpSignature。

4.3 SbpSignature 的推导

之前都是做各种准备，[Operator::InferSbpSignature]()里才进行真正的推导。简单讲就3步：

获取候选集
过滤不合适的签名
排序

4.3.1 SbpSignature 的候选集

调用 GetSbpSignaturesIf会获取 SbpSignature 的候选集。

在这个函数中，先调用 GetSbpSignatures 获取初步的候选集，再补充 broadcast 的候选集。候选集都保存到sbp_sig_list。

GetSbpSignatures是一个虚函数，UserOp 实现了自己的版本。这个函数中最核心的操作就是val_->get_sbp_fn，实际调用AddNOp::GetSbp。UserOpSbpContext是 UserOp 与 AddNOp 等类之间的协议接口的一部分。

如前所述，提供 SBP Signature 的候选集，是算子的责任。AddNOp这个算子比较简单，只给出两类签名：

对输入 tensor 的 shape 的每个 axis i，所有的 input/output 都创建一个 split(i)。
- 对于 tensor add 来说，input/output 的 shape 一样才能直接计算，所以 split 的 axis 也都一样。
所有的 input/output 都创建一个 partialsum。
- broadcast 在 Operator 中会补充。

候选集数据示例如下：

 {"sbp_signature":[{"bn_in_op2sbp_parallel":{"in_0":{"split_parallel":{"axis":"0"}},"in_1":{"split_parallel":{"axis":"0"}},"out_0":{"split_parallel":{"axis":"0"}}}},{"bn_in_op2sbp_parallel":{"in_0":{"split_parallel":{"axis":"1"}},"in_1":{"split_parallel":{"axis":"1"}},"out_0":{"split_parallel":{"axis":"1"}}}},{"bn_in_op2sbp_parallel":{"in_0":{"partial_sum_parallel":{}},"in_1":{"partial_sum_parallel":{}},"out_0":{"partial_sum_parallel":{}}}},{"bn_in_op2sbp_parallel":{"in_0":{"broadcast_parallel":{}},"in_1":{"broadcast_parallel":{}},"out_0":{"broadcast_parallel":{}}}}]}

4.3.2 过滤不合适的签名

分两步过滤不合适的签名

FilterAndCheckValidSbpSignatureListByLogicalShape中，对于每个输入tensor ibn，签名中 ibn 的 split axis，必须小于 tensor ibn 的 shape axes 数量。换句话说，如果 tensor 是二维的，就无法接受split(2)，只能是split(0)或split(1)。
FilterSbpSignatureList的作用是检验sbp_sig_conf约束，也就是从ConsistentTensorInferCache一路传过来的参数nd_sbp_constraints。这个过滤规则要求，符合条件的签名，其内容必须包含sbp_sig_conf。

4.3.3 签名排序

SortSbpSignatureListByCopyCost对候选签名进行排序。

优先按 OrderValue 比较
OrderValue 相等时，按 CopyCost 比较
二者都是较小的值优先。

OrderValue4SbpSig是对CalcOrderValue4SbpSig的封装，预先计算所有签名的 OrderValue 存到 map 中，便于 sort 函数查找。IbnCopyCost4SbpSig也是同理。

回过头来看CalcOrderValue4SbpSig的定义。因为AddNOp是有输入的，对于每个输入 tensor ibn 会加上一个权重，当 ibn 的 sbp 与签名中对应的 sbp 相同时，权重值为-10，即增加了选中的机会，因为 sbp 一致通常就不需要数据搬运。而parallel_num的条件在UserOp下应该是都成立的。

当 sbp_sig_conf 不空时，CalcOrderValue4SbpSig 直接返回0。因为如果签名不包含 sbp_sig_conf，即使 sbp 都一致，签名也不符合要求，所以直接返回0。

签名成本由ComputeIbnCopyCost4SbpSig计算。主要是根据输入和签名的 sbp 计算 cost：

如果 sbp 一致，cost 为0
partial_sum 和 broadcast 的 cost 都是一个超大的数字。
否则 cost 等于 input tensor 的数据传输字节数量。

4.4 推导结果

推导得到的nd_sbp_signature如下：

{"bn_in_op2nd_sbp":{"in_0":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]},"in_1":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]},"out_0":{"sbp_parallel":[{"split_parallel":{"axis":"0"}}]}}}

示例代码中，如果一个输入是split，另一个是broadcast，推导的签名结果都是broadcast。如果推断的sbp签名是split，是否能减少数据搬运呢？

5 NdSbp 的推导过程

NdSbp 的推导主要包括3步

重点看一下有效签名的获取。主要是两步：

GetNdSbpSignatureList: 获取全部签名
FilterNdSbpSignatureListByLogicalShape: 过滤不合适的签名

5.1 NdSbp 签名的候选集

GetNdSbpSignatureList 核心是两步：

GetSbpSignaturesIf: 得到一维的签名（和 1D SBP 的情况相同）
DfsGetNdSbpSignature: 根据一维签名拓展到多维

这个过程，如果深入到数据细节去看，会涉及 input/output、ranks、NdSbp 等多个维度，有点抽象复杂。

如果从官方文档 2D SBP中说明的 ranks 和 NdSbp 的物理含义出发，会更容易理解。
以ranks=[[0, 1, 2], [3, 4, 5]]为例（ranks=[r1, r2]），这是一个二维的设备阵列。算子的每个输入、输出也都有两个 sbp，NdSbpSignature 中的 value 是二维的，有两个槽位。假设 Op 的 1D Sbp 有 n 个签名。
从形式上看，NdSbpSignature 是先按 bn 组织数据。但是从数据分布的过程看，是先按SbpSignature组织数据。一个 NdSbpSignature 等价于 SbpSignature 数组。
NdSbp中的每个槽位，都表示一个 1D Sbp 的数据分布（所有的 input/output一起分布）。比如第 0 个槽位，就是在r1和r2这两个 sub group 之间分布数据，这个分布必须是一个有效的 1D SbpSignature（所有的 input/output一起分布）。第 1 个槽位，对于r1，就是将分配给它的数据子集，再根据一个 SbpSignature 进行分布（所有的 input/output一起分布）。
所以，只需要按 SbpSignature整体 填满两个槽位就行。每个槽位各有 n 种可能，一共有n*n个候选签名。这样生成的候选集是完整的，不会漏掉候选项。这就是direct product of 1D sbp signatures的含义。

6 模块间协作关系

SbpSignature 推导的实现用了大量 functional 的代码。应该为了不同模块间的信息屏蔽，或者父类、子类之间的逻辑复用、信息传递等目的，很多信息都封装到 function 中，需要时再检索、转换。

下图展示了不同模块之间的部分关系：

7 断点

启动命令

source /mnt/oneflow/build/source.sh
gdb --args python3 /mnt/oneflow/test.py
# set breakpoints
# run

断点示例

set breakpoint pending on
break oneflow::AddNOp::GetSbp
break oneflow::(anonymous namespace)::CheckAndConstructOp
break oneflow::Operator::FillBlobParallelDesc
break oneflow::Operator::InferNdSbpSignature
break oneflow::one::ConsistentTensorMetaInferArgs::MakeNdSbpConstraints
break oneflow::one::ConsistentTensorMetaInferArgs::MakeNdSbpInferHints
break oneflow::one::(anonymous namespace)::Interpret
break oneflow::ParallelDesc::MaybeInit
break oneflow::one::functional::impl::AddFunctor::operator()