背景

众所周知，我们在设计、描述AI模型的时候，大多是用AI框架提供的Python API去构建我们的模型，比如MindSpore的nn.Cell，nn.Conv等等。但你有没有想过，最终你的硬件运行这个模型的时候，使用的是什么代码格式？很显然，底层硬件看不懂Python代码。了解计算机知识的同学应该知道答案：AI框架和AI芯片的驱动软件会把最开始的Python定义的模型进行多次编译，使它转化成一份硬件可以执行的代码。而且不同的AI框架、不同的计算硬件，在转化过程中产生的中间件的格式也不相同，比如大家应该见过这些格式的AI模型：ckpt、pth、mindir、om、onnx、air等等。

AI模型格式转换

MindSpore IR就是昇思MindSpore模型在编译过程中的一种中间表达格式，全场景统一部署依赖于MindSpore IR，不同设备都可以加载MindSpore IR模型或者基于MindSpore IR模型进行编译运行。如果我们有一个MindSpore IR格式的模型，我们就可以对它进行改造，然后使用原有的编译方法对它修改后的MindSpore IR模型进行编译优化，再部署在不同的硬件设备上，而不用对不同的硬件做定制化的语法改造。为了看懂本文下面的章节，建议大家看两遍MindSpore官网对于MindSpore IR的介绍[1]，了解一下MindSpore IR图的构图格式以及形态。

那么我们怎样对MindSpore IR模型进行修改呢？昇思MindSpore提供了FuncGraph[2]这个C++类来表示MindIR模型，同时提供了很多API对FuncGraph进行修改，比如图节点的增删查改和节点复制、节点遍历等等，详细情况可以参考gitee目录。接下来我们详细讲讲怎么构造一个FuncGraph来表示计算流程。

1. 如何得到一个FuncGraph？

可以参考昇思MindSpore提供的ut的具体操作[3]。

我们可以使用mindspore的load()[4]接口加载一个mindir后缀的模型文件得到一个模型的FuncGraph：

func_graph = load('xxx.mindir')

也可以直接在C++脚本中new一个FuncGraph：

mindspore::FuncGraphPtr func_graph = std::make_shared<FuncGraph>();

2. 如何在创建的FuncGraph里面插入计算表达式？

假设我们现在要在图里面实现一个计算表达式：x+1，x是这个图的输入变量。我们想让这个图输出x+1的结果，那么我们该怎么写代码呢？

首先，我们要用函数GetParameterByName创建一个变量 x:

auto x = func_graph->GetParameterByName("x");

然后再创建一个值为1的Tensor：

// 设置shape为(1, 1)
ShapeVector tensor_shape{1, 1};
// 初始化Tensor，此时Tensor的值为0
tensor::TensorPtr const_one_tensor = std::make_shared<Tensor>(mindspore::kNumberTypeInt32, tensor_shape);
// 把Tensor的值赋为1
int *int_data = reinterpret_cast<int *>(const_one_tensor->data_c());
int_data[0] = 1;
// 包装成ValueNode
mindspore::ValueNodePtr const_one_tensor_vnode = std::make_shared<mindspore::ValueNode>(const_one_tensor);

由于我们要实现一个加法，所以需要用到加法算子：

mindspore::PrimitivePtr add_prim = mindspore::prim::kPrimAdd;

我们还需要给这个算子原语设置属性，否则运行的时候会提示属性缺失，最基本的属性是"is_load"：

add_prim->set_attr("is_load", MakeValue(true));

注意，通过查看mindspore/core/ops/core_ops.h，可以知道kPrimAdd实际上是一个预先声明好的全局变量std::make_shared<Primitive>("Add")，也就是说如果你的图里面有多处使用kPrimAdd，那么它们指向的实际上是同一个实例化primitive。从而当你的图里面需要用到多个相同算子、但算子的属性不同时，不能直接用mindspore::prim:kprimxxx的方式使用算子，而应该用std::make_shared<Primitive>("xxx")的方式实例化新的算子primitive。

然后，我们还需要把这个primitive包装成一个valueNode，然后加到这个图里面：

mindspore::ValueNodePtr add_v_node = std::make_shared<mindspore::ValueNode>(add_prim);
(void)func_graph->AddValueNode(add_v_node);

最后，我们再创建一个CNode，来实现这个“加”的过程，加法算子的valueNode放最前面，加数放后面：

CNodePtr add_c_node = func_graph->NewCNode({add_v_node, x, const_one_tensor_vnode});

最后，我们还需要给这个CNode设置abstract，告诉编译器这个CNode的输出形状和数据类型，以便让编译器检查整图算子之间输入输出的一致性。由于这个加法算子的输出形状和数据类型和const_one_tensor是一致的，所以直接取const_one_tensor的abstract作为CNode的abstract：

add_c_node->set_abstract(const_one_tensor->ToAbstract());

最后的最后，我们还需要引入一个return算子，把Add的结果作为图的输出：

mindspore::ValueNodePtr return_v = std::make_shared<Primitive>("Return"));
(void)func_graph->AddValueNode(return_v);
mindspore::CNodePtr return_c_node = fg_clone->NewCNode({return_v, add_c_node});
return_c_node ->set_abstract(const_one_tensor->ToAbstract());

注意，在上述过程中，我们添加到FuncGraph里面的自变量节点x和常量Tensor 1都是作为ValueNode插入的，而计算过程“加”和“返回”都要使用CNode来表示。

3.如何在FuncGraph中插入子图？

以上讲了如何创建一个 f(x)=x+1的计算图，如果我们现在有另外一个计算函数g(x)，并且要把这个计算逻辑加到原来的图里面，得到一个fg(x)=g(x)+1的计算图，该怎么做呢？

首先，假设g(x)对应的FuncGraph图为 g_graph，那么我们可以用Partial算子把func_graph的自变量x传到g(x)里面去：

// 创建Partial算子的primitive
mindspore::ValueNodePtr partial_vnode = std::make_shared<Primitive>("Partial", kSideEffectPropagate);
(void)func_graph->AddValueNode(partial_vnode);
// 把子图g_graph作为一个vnode
mindspore::ValueNodePtr subgraph_node = std::make_shared<mindspore::ValueNode>(g_graph);
subgraph_node ->set_abstract(g_graph->ToAbstract());
(void)func_graph->AddValueNode(subgraph_node);
// 创建Cnode，把x赋给g_graph
mindspore::CNodePtr partial_cnode = func_graph->NewCNode({partial_vnode, subgraph_node, x});

然后，再根据2.2里面的流程把partial_cnode和常数1相加，再返回。

CNodePtr add_c_node = func_graph->NewCNode({add_v_node, partial_cnode, const_one_tensor_vnode});

4. 如何在FuncGraph中插入switch控制流节点？

switch控制流结构一般指的是这种计算结构：如果x>1，那么f(x)=x+1；如果x <=1，那么f(x)=x-1。我们可以使用switch算子来实现这个计算流程。switch算子的输入格式一般是{条件判断节点，分支1，分支2}。

首先，我们需要创建条件判断节点：

ValueNodePtr greater_v_node = std::make_shared<Primitive>("Greater");
(void)func_graph->AddValueNode(greater_v_node);
CNodePtr greater_c_node = func_graph->NewCNode({greater_v_node, x, const_one_tensor_vnode});
ShapeVector the_shape{1, 1};
// 注意，此处要使用bool类型
tensor::TensorPtr greater_tensor = std::make_shared<Tensor>(mindspore::kNumberTypeBool, the_shape);
greater_c_node->set_abstract(greater_tensor->ToAbstract());
(void)func_graph->AddNode(greater_c_node);

然后创建分支1和分支2。假设f(x)=x+1对应的子图是sub_graph_1，f(x)=x-1对应的子图是sub_graph_2，我们要用2.3小节的方式把sub_graph_1和sub_graph_2用Partial算子包装后传到switch算子里面：

// 获得主图的manager
auto mgr = mindspore::Manage(func_graph);
// 把两个子图加到主图里面
mgr->AddFuncGraph(sub_graph_1);
mgr->AddFuncGraph(sub_graph_2);
// 假设创建partial分支的函数为AddPartialBranch
mindspore::CNodePtr switch_partial_1 = AddPartialBranch(sub_graph_1, x);
mindspore::CNodePtr switch_partial_2 = AddPartialBranch(sub_graph_2, x);
// 声明switch算子
mindspore::ValueNodePtr switch_v_node = std::make_shared<Primitive>("Switch"));
(void)func_graph->AddValueNode(switch_v_node);
// 把条件控制算子、两个分支加到switch算子里面
mindspore::CNodePtr switch_c_node = 
func_graph->NewCNode({switch_v_node, greater_c_node, switch_partial_1, switch_partial_2});
switch_c_node->set_abstract(switch_partial_1->ToAbstract());
func_graph->AddNode(switch_c_node);

这样就完成了switch控制流算子的构建，最后，还要加一个call算子才能让switch算子执行：

mindspore::CNodePtr call_cnode = func_graph->NewCNode({switch_c_node});     
func_graph->AddNode(call_cnode);

最后再把call node传给return node或者下面的其它节点。

参考链接

[1]https://www.mindspore.cn/docs/zh-CN/r1.9/design/mindir.html

[2]https://gitee.com/mindspore/mindspore/blob/r1.9/mindspore/core/ir/func_graph.h

[3]https://gitee.com/mindspore/mindspore/tree/r1.9/mindspore/core/ir

[4]https://gitee.com/mindspore/mindspore/tree/r1.9/tests/ut/cpp/ir

[5]https://www.mindspore.cn/docs/zh-CN/r1.9/api_python/mindspore/mindspore.load.html