一个算子在 MindSpore 框架中的执行流程
本文分析了一个算子在 MindSpore 框架中的执行流程。MindSpore 中设计了 Primitive 对算子进行了封装和抽象，一般来说封装和抽象是出于差异，这种差异来自于底层执行设备的差异，比如有 CPU，GPU，Ascend 等执行设备，每种执行设备上的计算逻辑，内存分配，通信逻辑各不相同。“没有什么问题是加一层抽象不能解决的”，如果有，咱们加两层哈哈。本文着重分析一个算子在 MindSpore 框架中的执行流程，对 Primitive 的设计论述相对较少，但是通过观察一个算子在框架中的执行流程，我们可以形象的感知到 Primitive 的作用。后续将会写一篇文章分析 Primitive 的设计。
原文发布在 GitLink 上的 MindSpore 评注解读上面，可以的话点一点链接！
链接：https://forum.gitlink.org.cn/forums/7330/detail
Python ReLU 代码
首先写一个最简单的算子，计算一个向量的 ReLU。
import mindspore
import mindspore.ops as P
from mindspore import Tensor
from mindspore import dtype as mstype
import numpy as np
input_x = Tensor(np.random.randn(2,3), mstype.float32)
relu = P.ReLU()
output_x = relu(input_x)
print(input_x)
print(output_x)
输出如下，大于 0 的部分保留，小于 0 的部分设为 0：
[[ 1.5206318 -0.35908994 -0.54122275]
[ 0.32850873 -0.6513135 -2.8261368 ]]
[[1.5206318 0. 0. ]
[0.32850873 0. 0. ]]
Python 端源码分析
ReLU 继承了 Primitive，Primitive 又继承自 C++ 导出的 Primitive_。
调用 ReLU 的背后，实际上执行的是 _run_op 这个函数
# mindspore/python/mindspore/ops/primitive.py
def __call__(self, *args):
should_elim, output = self.check_elim(*args)
for arg in args:
if isinstance(arg, Parameter) and arg.has_init:
arg.init_data()
if should_elim:
return output
return _run_op(self, self.name, args)
_run_op 调用了 C++ 导出的 real_run_op
# mindspore/python/mindspore/ops/primitive.py
@_wrap_func
def _run_op(obj, op_name, args):
"""Single op execution function supported by ge in PyNative mode."""
output = real_run_op(obj, op_name, args)
return output
由此我们进入到 C++ 端的实现中。
C++ 端源码分析
首先看看 pybind11 导出的函数。
// mindspore/ccsrc/pipeline/jit/init.cc
(void)m.def("real_run_op", &mindspore::pynative::RealRunOp, "Run op pynatively.");
RealRunOp 内部实现如下，先构造 info，然后再执行。info 里面会保存参数，保存 Python 端的 Primitive 对象。
// mindspore/ccsrc/pipeline/pynative/pynative_execute.cc
py::object RealRunOp(const py::args &args) {
CheckPyNativeContext();
const auto &executor = PynativeExecutor::GetInstance();
MS_EXCEPTION_IF_NULL(executor);
OpExecInfoPtr op_exec_info = executor->forward_executor()->GenerateOpExecInfo(args);
MS_EXCEPTION_IF_NULL(op_exec_info);
py::object ret = py::none();
PynativeExecutorTry(executor->forward_executor()->RunOpS, &ret, op_exec_info);
return ret;
}
传递到 ForwardExecutor 内部的 RunOpInner。首先执行检查，判断参数正确性，判断算子是否存在，特殊处理混合精度算子。因为是 PyNative，所以是动态图，也就是执行的时候构图，可以从下面看到 ConstructForwardGraph，最后通过 GetOpOutput 执行算子。
std::function<void(py::object *, const OpExecInfoPtr &)> RunOpS = [this](auto &&PH1, auto &&PH2) {
RunOpInner(std::forward<decltype(PH1)>(PH1), std::forward<decltype(PH2)>(PH2));
};
void ForwardExecutor::RunOpInner(py::object *ret, const OpExecInfoPtr &op_exec_info) {
MS_EXCEPTION_IF_NULL(ret);
MS_EXCEPTION_IF_NULL(op_exec_info);
MS_LOG(DEBUG) << "RunOp name: " << op_exec_info->op_name;
if (kSummaryOperators.count(op_exec_info->op_name)) {
MS_LOG(DEBUG) << "PyNative not support Operator " << op_exec_info->op_name;
return;
}
if (op_exec_info->op_name == prim::kPrimMixedPrecisionCast->name()) {
RunMixedPrecisionCastOp(op_exec_info, ret);
return;
}
// 1.Set cast for inputs
SetCastForInputs(op_exec_info);
// 2.Construct graph, first step abs will update by node
auto cnode = ConstructForwardGraph(op_exec_info);
// 3.Get inputs abstract
abstract::AbstractBasePtrList args_spec_list;
GetInputsArgsSpec(op_exec_info, &args_spec_list);
// 4.Get output abstract
bool prim_cache_hit = false;
GetOpOutputAbstract(op_exec_info, args_spec_list, &prim_cache_hit);
// 5.Get output
GetOpOutput(op_exec_info, args_spec_list, cnode, prim_cache_hit, ret);
}
GetOpOutput 内部还要处理反向梯度，需要将算子的信息记录下来。最终执行到 RunOpWithInitBackendPolicy 里面，使用指定的后端来运行算子。后端指的是运行引擎，比如 vm，ge，tsd 等，不同的运行引擎应该有不同的优化策略。
void ForwardExecutor::GetOpOutput(const OpExecInfoPtr &op_exec_info,
const abstract::AbstractBasePtrList &args_spec_list, const CNodePtr &cnode,
bool prim_cache_hit, py::object *ret) {
MS_EXCEPTION_IF_NULL(op_exec_info);
const auto &prim = op_exec_info->py_primitive;
MS_EXCEPTION_IF_NULL(prim);
// Infer output value by constant folding
MS_EXCEPTION_IF_NULL(ret);
py::dict output = abstract::ConvertAbstractToPython(op_exec_info->abstract, true);
if (!output[ATTR_VALUE].is_none()) {
*ret = output[ATTR_VALUE];
grad()->RecordGradOpInfo(op_exec_info);
MS_LOG(DEBUG) << "Get output by constant folding, output is " << py::str(*ret);
return;
} else if (prim->is_const_prim()) {
*ret = py::cast("");
grad()->RecordGradOpInfo(op_exec_info);
MS_LOG(DEBUG) << "Get const prim";
return;
}
// Add output abstract info into cache, the const value needs to infer evert step
if (grad()->enable_op_cache() && !prim_cache_hit && !IsDynamicShape(op_exec_info)) {
AbsCacheKey key{prim->name(), prim->Hash(), prim->attrs()};
auto &out = prim_abs_list_[key];
out[args_spec_list].abs = op_exec_info->abstract;
out[args_spec_list].attrs = prim->evaluate_added_attrs();
}
// Run op with selected backend, nop is no need run backend
ValuePtr out_real_value = nullptr;
if (op_exec_info->is_nop_prim) {
DoNopOutput(op_exec_info, &out_real_value);
*ret = BaseRefToPyData(out_real_value);
} else {
auto result = RunOpWithInitBackendPolicy(op_exec_info);
py::object out_real = result;
if (result.size() == 1 && op_exec_info->abstract != nullptr &&
!op_exec_info->abstract->isa<abstract::AbstractSequence>()) {
out_real = result[0];
}
// Get output value
if (grad()->grad_flag()) {
out_real_value = PyObjToValue(out_real);
}
*ret = out_real;
}
if (grad()->need_construct_graph() && !grad()->in_cell_with_custom_bprop_()) {
MS_EXCEPTION_IF_NULL(cnode);
const auto &obj_id = GetId(*ret);
cnode->set_abstract(op_exec_info->abstract);
node_abs_map_[obj_id] = op_exec_info->abstract;
grad()->SaveOutputNodeMap(obj_id, *ret, cnode);
grad()->DoOpGrad(op_exec_info, cnode, out_real_value);
// Dynamic shape should update to top cell
if (IsDynamicShape(op_exec_info)) {
grad()->top_cell()->set_dynamic_shape(true);
}
} else {
node_abs_map_.clear();
}
// Record op info for judge whether the construct of cell has been changed
grad()->RecordGradOpInfo(op_exec_info);
grad()->UpdateForwardTensorInfoInBpropGraph(op_exec_info->op_info, out_real_value);
}
接下来进入到根据策略执行算子的地方。如果跟进去看，RunOpInVM，最后的逻辑是运行了 python 的计算函数。如果要往 C++ 走，应该还要看 RunOpInMs，即在 MindSpore 中执行算子。
py::object ForwardExecutor::RunOpWithBackendPolicy(MsBackendPolicy backend_policy, const OpExecInfoPtr &op_exec_info) {
py::object result;
if (backend_policy == kMsBackendVmOnly) {
#ifndef ENABLE_TEST
if (kVmOperators.find(op_exec_info->op_name) != kVmOperators.end()) {
result = RunOpInVM(op_exec_info);
} else {
result = RunOpInMs(op_exec_info);
}
#else
result = RunOpInVM(op_exec_info);
#endif
}
return result;
}
在 RunOpInMs 里面，可以看到运行时 Runtime 又分为了 Ms 和 MindRT 两种，我们暂且看 Ms 这一条分支。先获取一个 Session，然后运行算子。Session 需要通过 MS_REG_SESSION 宏来注册一个设备对应的 Session 类。这个类是一个注册工厂，在启动的时候，创建对应的类，注册到工厂当中，后面可以通过 GetCurrentSession 运行时获取对应的 Session 类。
VectorRef outputs;
if (!enable_mind_rt) {
auto cur_session = GetCurrentSession(cur_target, device_id);
MS_EXCEPTION_IF_NULL(cur_session);
cur_session->RunOp(&op_run_info, &outputs);
} else {
auto cur_mind_rt_backend = GetMindRtBackend(cur_target, device_id);
MS_EXCEPTION_IF_NULL(cur_mind_rt_backend);
mindspore::ScopedLongRunning long_running;
cur_mind_rt_backend->RunOp(&op_run_info, &outputs);
}
// mindspore/ccsrc/backend/common/session/session_factory.h
#define MS_REG_SESSION(DEVICE_NAME, SESSION_CLASS) \
static const SessionRegistrar g_session_registrar__##DEVICE_NAME##_##_reg( \
DEVICE_NAME, []() { return std::make_shared<SESSION_CLASS>(); });
通过搜索使用了 MS_REG_SESSION 宏的地方，我们可以看到有 cpu, gpu, ascend 还有其他 session 类调用了。我们专注于看 cpu 这一条分支。上面分析到了，通过跟踪 CPUSession 的父类 SessionBasic，再看它的类成员 Executor，我们可以看到最后其实调用的是每个 Session 的 RunOpImpl 这个函数。于是我们只要专注看 CPUSession 的 RunOpImpl 即可。
调用 ProcessInputTensorsForHeterogeneous 如果数据不在 CPU 上，那么将数据同步到 CPU 上。构建 Op 计算图，创建输出向量，再将输入和输出向量绑定到 device::cpu::CPUKernelRuntime 对象上，最后调用 Run 方法执行。
void CPUSession::RunOpImpl(const GraphInfo &graph_info, OpRunInfo *op_run_info,
std::vector<tensor::TensorPtr> *input_tensors, VectorRef *outputs,
const std::vector<int64_t> &tensors_mask) {
MS_EXCEPTION_IF_NULL(input_tensors);
MS_EXCEPTION_IF_NULL(op_run_info);
ProcessInputTensorsForHeterogeneous("CPU", *input_tensors);
const auto &kernel_graph = BuildOpImpl(*op_run_info, graph_info, *input_tensors, tensors_mask);
EraseValueNodeTensor(tensors_mask, input_tensors);
// Remove reorder after PS feature finish adapting push/pull in auto_monad.
auto execution_order = kernel_graph->execution_order();
Reorder(&execution_order);
kernel_graph->set_execution_order(execution_order);
// runtime init
if (!runtime_.Init()) {
MS_LOG(EXCEPTION) << "Kernel runtime init error.";
}
runtime_.AssignKernelGraphAddress(kernel_graph.get());
std::map<tensor::TensorPtr, session::KernelWithIndex> tensor_to_node;
runtime_.CreateOutputTensors(kernel_graph.get(), *input_tensors, outputs, &tensor_to_node);
runtime_.BindInputOutput(kernel_graph.get(), *input_tensors, outputs);
bool ret = runtime_.Run(*kernel_graph, false);
if (!ret) {
MS_LOG(EXCEPTION) << "Run Op failed";
}
UpdateDynamicOutputShape(tensor_to_node);
// update output abstract of dynamic op to op_run_info
if (op_run_info->output_is_dynamic_shape) {
UpdateOutputAbstract(kernel_graph, op_run_info);
}
SetOutputFlags(*outputs);
runtime_.RunOpClearMemory(*kernel_graph);
}
device::cpu::CPUKernelRuntime 是 MindSpore Primitive 设计的一部分，Primitive 通过对 “计算设备” 进行封装和抽象，抽象出了 KernelRuntime 算子执行的逻辑就在这里面，我们已经快接近真相了。CPUKernelRuntime 的 Run 方法里面为了处理 Profile，安全，调试等，写了不少条件编译，我们只看最核心的逻辑。前面我们可以看到 kernel_graph 已经设置好了执行顺序，因此在这里我们只需要依次获取 kernel，然后逐个执行，调用 Kernel 的 Launch 方法。
for (const auto &kernel : kernels) {
auto kernel_mod = AnfAlgo::GetKernelMod(kernel);
MS_EXCEPTION_IF_NULL(kernel_mod);
// akg kernel do not support dynamic shape by now
kernel::NativeCpuKernelMod *cpu_kernel = nullptr;
if (session::AnfRuntimeAlgorithm::GetKernelType(kernel) != KernelType::AKG_KERNEL) {
cpu_kernel = dynamic_cast<kernel::NativeCpuKernelMod *>(kernel_mod);
MS_EXCEPTION_IF_NULL(cpu_kernel);
}
if (common::AnfAlgo::IsDynamicShape(kernel)) {
AnfAlgo::InferShape(kernel);
auto args = kernel::GetArgsFromCNode(kernel);
if (cpu_kernel != nullptr && cpu_kernel->Resize(args->op, args->inputs, args->outputs, args->depend_tensor_map) ==
kernel::KRET_RESIZE_FAILED) {
MS_LOG(EXCEPTION) << "Node " << kernel->fullname_with_scope() << " Resize failed!";
}
}
std::vector<kernel::AddressPtr> kernel_inputs;
std::vector<kernel::AddressPtr> kernel_workspaces;
std::vector<kernel::AddressPtr> kernel_outputs;
GetRuntimeAddressFromNode(kernel, &kernel_inputs, &kernel_outputs, &kernel_workspaces);
bool ret = true;
try {
ret = kernel_mod->Launch(kernel_inputs, kernel_workspaces, kernel_outputs, 0);
} catch (std::exception &e) {
MS_LOG(EXCEPTION) << e.what() << trace::DumpSourceLines(kernel);
}
if (!ret) {
MS_LOG(EXCEPTION) << "Launch kernel failed." << trace::DumpSourceLines(kernel);
}
static_cast<CPUMemoryManager *>(mem_manager_.get())->DecreaseAddressRefCount(kernel);
}
再往下翻，我们找到 ReLU 的 CPU kernel 实现，其实就是判断是否大于 0，大于 0 的保留，小于 0 的设置为 0。
template <typename T>
void Relu(ArithmeticSelfCpuKernelFunc *content, const T *in, T *out, size_t size) {
auto task = [in, out](size_t start, size_t end) {
for (size_t i = start; i < end; i++) {
out[i] = std::greater<T>()(in[i], 0) ? in[i] : 0;
}
};
ParallelLaunchAutoSearch(task, size, content, &content->parallel_search_info_);
}
MS_KERNEL_FACTORY_REG_BY_CREATOR(NativeCpuKernelMod, ReLU,
[]() { return std::make_shared<ArithmeticSelfCpuKernelMod>(kReLU); });
总结
自此，我们分析了一个算子从 Python 前端最终运行到 C++ CPU Kernel 的全流程，其他分支流程大同小异，基本遵循着 “构图，图优化，分配内存，绑定输入输出，运行” 这样一个模式。