Add oneflow.nn.functional.depend api #9807

PYNing · 2023-01-30T08:35:49Z

增加一个Python OP，用于：
（1）防止指定OP在静态图优化时被消除或重排序；
（2）作为用户增加静态图控制边的接口，实现对执行序的约束或修改。

该OP存在于其他具有静态图特性的框架，例如：
Mindspore（https://www.mindspore.cn/docs/zh-CN/r1.9/api_python/ops/mindspore.ops.Depend.html）
Tensorflow （https://www.tensorflow.org/api_docs/python/tf/control_dependencies）

特性：

（1）为避免Eager Mode下的性能损失，Python接口判别在Eager Mode还是在Grpah Mode下运行，Eager Mode直接返回输入；
（2）为避免Grpah Mode下的性能损失，增加可配置开关的Pass，用于消除多添加的OP，并相应的添加底层的控制边;
（3）考虑self-loop导致的死锁的情况；
（4）Pass考虑了多个depend OP连锁的情况，以及可能重复添加控制边的情况；
（5）Kernel直接复用已有代码；
（6）包含了单元测试（考虑用户多种可能的用法）和文档

效果：

以单元测试的第一个例子（test_depend_graph_case0）为例
网络定义

class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(128, 128)

    def forward(self, x):
        # to ensure "x * 2" be executed before "self.linear(x)" in graph mode
        # base use case
        x1 = x * 2
        x = nn.functional.depend(x, x1)
        x2 = self.linear(x)
        return x2

不使用nn.functional.depend时，job_TestGraph_0_plan.dot的截图

从图可知，OP “model-scalar_mul-0”和“model.linear-matmul-1”并无执行顺序的约束，且从ID大小推测后者将先与前者执行，与用户定义的OP执行顺序不一致

使用nn.functional.depend后，job_TestGraph_0_plan.dot的截图

从图可知，OP “model-scalar_mul-0”与“model.linear-matmul-1”之间增加了一条控制边，且从ID大小推测前者将先与后者执行，达到用户控制OP执行顺序的目的。且由于存在控制边，防止了 “model-scalar_mul-0”被其他Pass消除

PS: 第一次向OF贡献算子，希望能被接纳~
如有要补充的请及时告知~

wangxicoding · 2023-01-31T04:00:06Z

oneflow/core/job/job_build_and_infer_ctx.cpp

@@ -1005,6 +1005,9 @@ Maybe<void> LazyJobBuildAndInferCtx::Complete() {
    // pinned identity can be pruned since GenerateOptimizerOpConfs pass has
    // already construct a complete computational graph
    JUST(DoPass("PrunePinnedIdentityOpPass"));
+    // prune depend OP and and add ctrl_in_op to op_conf accordingly
+    // to express the same semantics and avoid performance loss
+    JUST(DoPass("PruneDependOpPass"));


为啥放到这个位置，而不是更前面

已更新代码，将PruneDependOpPass提前到PruneAmpWhiteIdentityOpPass前。
理由：
（1）将PruneDepend尽早执行，可以发掘更多的算子优化空间（如删除Depend OP后可能满足FuseAddToOutputPass的执行条件）；
（2）但在前面的部分Pass在删除或更新OP时未考虑控制边的转移或保持（如EliminateDeadNodesPass、AutoMixedPrecision）。如果放在它们之前执行，新添加的控制边可能丢失导致失效。

经阅读前面的Pass代码和测试，将PruneDependOpPass的执行提前到PruneAmpWhiteIdentityOpPass之前比较合适。

wangxicoding · 2023-01-31T07:56:11Z

oneflow/core/autograd/gradient_funcs/depand.cpp

+    if (ctx->in_requires_grad) { in_grads->at(0) = out_grads.at(0); }
+    if (ctx->depend_tensor_requires_grad) {
+      in_grads->at(1) =
+          JUST(functional::Constant(ctx->depend_tensor_shape, Scalar(0), out_grads.at(0)->dtype(),


如果要实现反向的话，dtype和device应该和depend_tensor一样吧

已更新代码，depend_tensor梯度的dtype和device与depend_tensor一致。

wangxicoding · 2023-01-31T09:09:40Z

oneflow/core/job_rewriter/prune_depend_op_pass.cpp

+    // GetRelativeNodes() considers the chain of multiple depend OP Nodes and processes them
+    // from top to down, so skip the intermediate nodes
+    if (!IsDependOPNodeAtTop(op_node, del_nodes)) { continue; }
+    const std::vector<RelativeNodes> relatives = GetRelativeNodes(op_node, del_nodes);


这一段逻辑有点晦涩，有没有一些graph之类的注释，更直观些

wangxicoding · 2023-02-02T07:13:17Z

Pass的代码看了下，目前没看出啥问题，但目前的算法逻辑还是太晦涩了。
其实可以按DAG拓扑排序遍历的方式来去掉depend op。入度为0时遍历执行操作，若为depend op，删op，更新。一次遍历即可

PYNing · 2023-02-03T04:13:48Z

Pass的代码看了下，目前没看出啥问题，但目前的算法逻辑还是太晦涩了。其实可以按DAG拓扑排序遍历的方式来去掉depend op。入度为0时遍历执行操作，若为depend op，删op，更新。一次遍历即可

已按此思路重构Pass的代码

wangxicoding · 2023-02-03T09:37:20Z

oneflow/core/job_rewriter/prune_depend_op_pass.cpp

+
+      // Step 1.3 process src nodes
+      const OpNode* cur_src_node = GetNodeFromInputEdge(cur_node);
+      if (IsDependyOp(dst_node->op().op_conf()) && cur_node == GetNodeFromInCtrlEdge(dst_node)) {


如果可以边遍历边改图，这段逻辑可以去掉

因缺少API支持，这种写法比较困难。
参考已有的较为简单的Pass，比如EliminateDeadNodesPass 和 PruneAmpWhiteIdentityOpPass，它们理论上可以边遍历边改图，但没有这样做，而是走产生OpGraph - > 分析OpGraph 并记录变更->根据变更修改Job对象的流程。

wangxicoding · 2023-02-03T10:26:52Z

oneflow/core/job_rewriter/prune_depend_op_pass.cpp

+    CHECK(src_node);
+    const OpNode* nearest_depend_node = node_info.second.nearest_depend_node;
+    const auto& old_lbi = nearest_depend_node->op().BnInOp2Lbi(nearest_depend_node->op().SoleObn());
+    const auto& new_lbi = src_node->op().BnInOp2Lbi(src_node->op().SoleObn());


src_node可能不止一个输出吧，比如：

a, b = op0() c = op1() b = depend(b, c) d = op2(b)

得通过连接边来判断

已更新代码，用于处理src_node多个输出的情况，并针对这种情况追加了单例测试（test_depend_graph_case7）

wangxicoding

LGTM，看着没啥子大问题了。
不过当前depend op不支持对source op添加控制边来控制执行顺序。比如下面三个op没有输入，就没法在用户侧通过depend op来控制，这种场景可以以后再考虑

a = op0()
b = op1()
c = op2()

wangxicoding · 2023-02-06T07:25:36Z

oneflow/core/job_rewriter/prune_depend_op_pass.cpp

+    CHECK(src_node && depend_node_nearest_dst && depend_node_nearest_src);
+    const auto& old_lbi =
+        depend_node_nearest_dst->op().BnInOp2Lbi(depend_node_nearest_dst->op().SoleObn());
+    const auto new_lbi = GetNewLbi(src_node, depend_node_nearest_src);


是不是直接用depend_node_nearest_src的输入就可以了。不过目前这样也没有问题就是了

可以。不过，148~168行的逻辑（Step 1.3）涉及对src_node的更新，去掉对src_node的记录会显得这段逻辑不那么自然……

BBuf · 2023-02-07T02:43:36Z

oneflow/core/functional/functional_api.yaml

@@ -2764,6 +2764,10 @@
  signature: "Tensor (Tensor input) => IsFinite"
  bind_python: True

+- name: "depend"
+  signature: "Tensor (Tensor input, Tensor depend_tensor) => Depend"


函数签名要优化下么？比如Tensor depend_tensor是不是直接叫depend就好了，这里会考虑和一个list of tensor建立控制边么

（1）Tensor depend_tensor 已重命名为 depend；
（2）已支持传入depend的类型为Tensor或List[Tensor]，并为List[Tensor]的情形追加了测试样例。

strint · 2023-02-07T08:32:16Z

python/oneflow/framework/function_util.py

@@ -532,6 +532,17 @@ def set_prune_amp_white_identity_ops(func_desc, value=True):
    func_desc.job_config_proto.prune_amp_white_identity_ops = value


+@oneflow_function_config("prune_depend_ops")
+def set_prune_depend_ops(func_desc, value=True):


这个是不是没有用？

graph 的控制接口现在都在 https://oneflow.readthedocs.io/en/master/graph.html#config-options-on-a-graph

这个 python/oneflow/framework/function_util.py 是计划移除的

是的，没有用。之前有询问过这里的代码，答复是“function_util.py 里面是 0.4 之前的接口，代码会清理掉。”
我可以删掉这段。

graph 的控制接口现在都在 https://oneflow.readthedocs.io/en/master/graph.html#config-options-on-a-graph

这个 python/oneflow/framework/function_util.py 是计划移除的

我是这样想的
（a）预计这个OP很少被使用
（b）优化这个OP的Pass相对安全
（c）避免config太多，用户看文档太花时间

综合考虑就没有在config里添这个Pass开关了。

如有必要，我可以添加

如有必要，我可以添加

那可以先不加

PYNing and others added 6 commits January 18, 2023 15:14

add a new OP: oneflow.nn.functional.depend

7fcdb35

reformat .py file and remove unused import

7245f9f

Merge branch 'Oneflow-Inc:master' into add_depend_op

1b5db60

Fix bugs in prune_depend_op_pass

1de065a

Add UnitTest and DocTest for Depend OP

ddddbf6

Merge branch 'Oneflow-Inc:master' into add_depend_op

137cfd8

PYNing requested review from hjchen2, BBuf, jackalcooper and daquexian as code owners January 30, 2023 08:35

PYNing added feature op test documentation api python labels Jan 30, 2023

PYNing requested a review from oneflow-ci-bot January 30, 2023 08:51

PYNing added 3 commits January 30, 2023 09:41

fix license and python format

354b007

fix C++ format

2379f64

fix error in static analysis

ad334bc

wangxicoding reviewed Jan 31, 2023

View reviewed changes

PYNing added 9 commits January 31, 2023 10:50

adjust pass order for potential better optimization

80e094e

Fix autograd of depend OP

fd13735

Merge branch 'Oneflow-Inc:master' into add_depend_op

cd592da

Simplify logic in PruneDependOpPass

994226d

fix logic in PruneDependOpPass and refine comments

5ff01d2

Merge branch 'Oneflow-Inc:master' into add_depend_op

0ba2e63

fix cpp format error

25cc35c

add log for PruneDependOpPass debug

684c882

fix logic of PruneDependOpPass

feed392

Refactor code of PruneDependOpPass for Readability

56fb105

PYNing added 2 commits February 3, 2023 11:33

Merge branch 'Oneflow-Inc:master' into add_depend_op

8c34c5f

Refactor code of PruneDependOpPass for Readability

fa2efad

wangxicoding reviewed Feb 3, 2023

View reviewed changes

PYNing added 4 commits February 6, 2023 05:54

enhance: consider source node has multiple outputs

e72c1da

add two more tests to TestDependGraph

ba854bd

add one more test case to TestDependGraph

5d35a37

reformat python file

439e229

wangxicoding approved these changes Feb 6, 2023

View reviewed changes

Merge branch 'Oneflow-Inc:master' into add_depend_op

0e42ddf

BBuf approved these changes Feb 7, 2023

View reviewed changes

PYNing added 3 commits February 7, 2023 03:51

rename the second parameter

4a7ee4f

support multiple tensors form different OP

3c4d87c

Merge branch 'Oneflow-Inc:master' into add_depend_op

7829fb2

strint reviewed Feb 7, 2023

View reviewed changes

PYNing enabled auto-merge (squash) February 7, 2023 11:22

PYNing added 2 commits February 8, 2023 00:03

Merge branch 'master' into add_depend_op

7ce7f52

Merge branch 'master' into add_depend_op

4dce37b

PYNing merged commit f4be05d into Oneflow-Inc:master Feb 8, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add oneflow.nn.functional.depend api #9807

Add oneflow.nn.functional.depend api #9807

PYNing commented Jan 30, 2023 •

edited

Loading

wangxicoding Jan 31, 2023

PYNing Jan 31, 2023

wangxicoding Jan 31, 2023

PYNing Jan 31, 2023

wangxicoding Jan 31, 2023

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

wangxicoding commented Feb 2, 2023

PYNing commented Feb 3, 2023

wangxicoding Feb 3, 2023

PYNing Feb 6, 2023 •

edited

Loading

wangxicoding Feb 3, 2023

PYNing Feb 6, 2023

wangxicoding left a comment

wangxicoding Feb 6, 2023

PYNing Feb 7, 2023

BBuf Feb 7, 2023

PYNing Feb 7, 2023 •

edited

Loading

strint Feb 7, 2023

strint Feb 7, 2023

PYNing Feb 7, 2023

PYNing Feb 7, 2023

strint Feb 7, 2023

Add oneflow.nn.functional.depend api #9807

Add oneflow.nn.functional.depend api #9807

Conversation

PYNing commented Jan 30, 2023 • edited Loading

特性：

效果：

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

wangxicoding commented Feb 2, 2023

PYNing commented Feb 3, 2023

Choose a reason for hiding this comment

PYNing Feb 6, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wangxicoding left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

PYNing Feb 7, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

PYNing commented Jan 30, 2023 •

edited

Loading

PYNing Feb 6, 2023 •

edited

Loading

PYNing Feb 7, 2023 •

edited

Loading