关于c++:OneFlow源码阅读8eager模式下的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 的推导。

本节以下文字摘自 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。

以下文字次要参考官网文档 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 都是数组。

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]

为了进步性能，从 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 推导过程中用到。

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 持有所有推导过的后果。

理论的推导在 ConsistentTensorInferCache::Infer 中进行。

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

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。

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。

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

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

调用 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":{}}}}]}

分两步过滤不适合的签名

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。

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 的数据传输字节数量。

推导失去的 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，是否能缩小数据搬运呢？

NdSbp 的推导次要包含 3 步

调用 GetValidNdSbpSignatureList 获取无效的签名
剔除不能蕴含 nd_sbp_constraints 的签名
贪婪搜寻较优的签名

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

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

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 的含意。

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

下图展现了不同模块之间的局部关系：

启动命令

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

oneflow v0.8.0
SBP Signature
2D SBP
placement api

关于c++:OneFlow源码阅读8eager模式下的SBP-Signature推导

1 SBP Signature

1.1 NdSbp 及 NdSbpSignature

2 placement.hierarchy

3 tensor add 是哪个算子？

4 一维 SBP 的推导过程

4.1 ConsistentTensorInferCache 中的推导筹备

4.1.1 推导 output 的 shape 和 dtype

4.1.2 结构 UserOp

4.2 Operator 中的推导筹备

4.3 SbpSignature 的推导

4.3.1 SbpSignature 的候选集

4.3.2 过滤不适合的签名

4.3.3 签名排序

4.4 推导后果

5 NdSbp 的推导过程

5.1 NdSbp 签名的候选集

6 模块间协作关系

7 断点

8 参考资料