mindspore源码学习3---- 动态图执行流程

再次回到梦开始的地方: Cell的call方法:

image.png

如果我们当前不是GRAPH_MODE, 那我们就会跳过378的判断, 继续向下, 走到391行, 用_pynative_exec来创建graph, 继而执行与释放资源等.
同样经过一轮跟第1篇类似的操作, 我们可以追踪到c++代码的这里:

image.png

进去看一下new_graph和end_graph都在做什么:

image.png

看上图, 大概就是调用了另一个叫做GradExecutor的类来执行new和end操作, 内部实现呢:

image.png

image.png

看起来init的时候内部有个什么stack用来存储需要运行的cell的样子. 这个后面再深究吧.

image.png

在end的时候, 还会对这个stack做pop操作, 应该是跟之前的init对应的.

回到Cell的call方法, 它调用的run_construct实现在哪儿呢?

image.png

可以看到, 这里会调用我们手写的layer的construct函数, 而我们的网络会调用一些现成的layer或者算子, 因为layer最终还是调用到算子上, 所以我们来看一下算子是如何被调用的.

因为算子跟layer相似, 都是在init函数中创建算子对象, 在construct函数把对象当成函数调用, 就像这样:

image.png

因此算子具体执行了什么逻辑, 也是需要看算子的call方法, 所有的算子的基类都是primitive这个类, 我们注意到primitive这个类就是有call方法的:

image.png

会调用_run_op这个方法, 然后调用real_run_op:

image.png

这个real_run_op同样也是从c++侧定义好, 供python侧调用的:

image.png

这还是祖传手艺了, 继续去c++侧找具体实现:

image.png

在mindspore\ccsrc\pipeline\pynative\pynative_execute.cc这个文件下的RealRunOp方法:

image.png

又是我们的老朋友PynativeExecutor了, 继续看它的forward_executor和RunOpS方法:

image.png

继续看RunOpInner的实现:

image.png

看起来前4步都是在准备一些需要的输入数据和输出格式, 最后一步是执行计算, 比如说第四步时需要进行shape和dtype的推断:

image.png

继续看:

image.png

image.png

image.png

image.png

image.png

回调到python侧:

image.png

以Add算子为例的话, 就会调用到它的infer_value, 但是实际上这里的x和y不是常量的Tensor, 都是None,

image.png

实际的计算还是在第5步进行. 那我们还是关注一下第5步的实现:

image.png

继续看, 在backendpolicy里面随便选一个, 就看RunOpInVM吧

image.png

继续看到RunPyComputeFunction

image.png

那么vm是拿什么函数来执行的呢?

image.png

可以看到, 这个vm有个注册器, 只要注册在里面的算子就可以找到它的vm的实现, 从而调用到, 以Add为例的话, 可以该算子测试时是有python本地实现的:

image.png

但我们此时安装的mindspore并不包含test包, 在我们的包的安装目录下是找不到上图这段代码的.
也就是说我们走的不是vm的逻辑, 而是走的其他逻辑. 哈哈哈, 我这里绕了个大弯路, 但是我觉得写出来也是有意义的, 因为我通过设置环境变量来打印了更多的log, 后面还是发现在pycharm中实际走的是哪一个了! mindspore有一个叫做GLOG_v的环境变量, 默认为3, 也就是只打印warning和error, 我们可以把它设置为0, 就可以看到info和debug信息了! 如何设置呢? 很简单, 就在pycharm, run, edit configurations就可以了:

image.png

增加一个环境变量GLOG_v, 值设为0:

image.png

重新跑一次我们的用例, 就可以看到详细的log了

image.png

通过log其实可以看到, 我们走的是这个!

image.png

在这个方法中, 会调用session来进行计算, 这里的session根据我们的设置不同, 可以是cpu/gpu或者ascend.

image.png

我在自己家里的电脑, 自然是cpu, 继续看log也可以证明这一点.

image.png

在执行算子kernel之前, 还会检查已经注册过的算子中是否有当前的这个名字(Add). 这里又有很长的故事了, 以后有机会再写.

cpu_session会调用runtime的run方法

image.png

继而调用每一个kernel的launch方法

image.png

注意这个kernel_mod是一个基类, 每一个kernel都是它的子类:

image.png

最终会调用到cpukernel下Add算子的实现:

image.png

总结一下就是:
动态图的执行从python侧Cell的call方法开始, 调用到算子的call方法, 层层封装调到了c++侧. 通过不同的设置可以选择走ms或者ge或者vm, 我们已经看到vm在我们用户安装包里是不会走到的, 我的情况是走的ms, 然后又会选择不同的后端(cpu/gpu/ascend)来选择不同的算子, 在选择算子这里需要算子提前注册好. 然后最终调用到算子的launch方法, 实现真正的计算.

当然在这个过程中我没有关心资源管理的问题, 后续有机会再写文章详细探讨