如何查看MindSpore的IR图

1、为什么需要IR图

中间表示（Intermediate Representation）是编译器用于表示源代码的数据结构或代码，通常深度学习框架中都会有IR这种表达，用于将用户脚本语言翻译成底层语言，介于源语言和目标语言之间的程序表示。相当于IR定义了一套从脚本到机器语言的翻译逻辑，便于编译器进行程序分析与优化。好的中间表示有利于AI模型的编译优化和执行，是AI框架进行高效训练和推理的基础。

MindSpore作为2020年3月华为正式开源的一套全场景AI计算框架，同样提供一种基于图表示的函数式中间表示，即：MindIR。通过统一的算子IR定义，消除不同后端的模型差异。MindIR是基于图表示的函数式表达，接近于ANF函数式的语义。

2、IR图目前有哪些类型

从训练的角度看，目前业界的AI框架有三种执行模式：Eager执行模式、图执行模式和Staging(混合)执行模式，其中高性能模式下（Graph执行模式和Staging执行模式）都要基于图层IR。

MindIR更偏向于图执行模式，根据Python脚本的AST进行构图。将AST语义转换成ANF语义。可以覆盖控制流等复杂网络的各种语法表达，同时兼顾到编译性能。由于ANF是一种函数式语言，所以会存在副作用的问题。目前MindSpore在图模式下也已支持消除副作用。

MindIR具有以下特点：

（1）是基于图的。MindSpore框架同时支持静态图和动态图。其中IR的表示主要是针对图模式。其中的函数是可以被递归调用，也可以被当做参数传到其他的函数中，或者从其他函数中返回，使得MindSpore可以表达一系列的控制流结构。

（2）是纯函数的。纯函数是指函数的结果只依赖函数的参数。若函数依赖或影响外部的状态，比如，函数会修改外部全局变量，或者函数的结果依赖全局变量的值，则称函数具有副作用。若使用了带有副作用的函数，代码的执行顺序必须得到严格的保证，否则可能会得到错误的结果，比如对全局变量的先写后读变成了先读后写。值得说明的是，MindSpore图模式已支持副作用的表达，能够将副作用的表达转换为纯函数的表达，从而在保持ANF函数式语义不变的同时，确保执行顺序的正确性，从而实现自由度更高的自动微分。对比Jaxpr IR来说，它是一种强类型、纯函数的中间表示，其输入、输出都带有类型信息，函数输出只依赖输入，不依赖全局变量。所以Jax的静态模式下并不支持副作用的表达，即有全局变量使用的限制。

（3）是支持闭包表示的。反向模式的自动微分，需要存储基本操作的中间结果到闭包中，然后再去进行组合连接。所以有一个自然的闭包表示尤为重要。闭包是指代码块和作用域环境的结合，在MindIR中，代码块是以函数图呈现的，而作用域环境可以理解为该函数被调用时的上下文环境。

（4）是强类型的。每个节点需要有一个具体的类型，这个对于性能最大化很重要。在机器学习应用中，因为算子可能很耗费时间，所以越早捕获错误越好。因为需要支持函数调用和高阶函数，相比于TensorFlow的数据流图，MindIR的类型和形状推导更加复杂且强大。

3、如何得到一段脚本对应的IR图，如何查看IR图

通过context.set_context(save_graphs=True)来保存各个编译阶段的中间代码。被保存的中间代码有两种格式，一个是后缀名为.ir的文本格式，一个是后缀名为.dot的图形化格式。当网络规模不大时，建议使用更直观的图形化格式来查看，当网络规模较大时建议使用更高效的文本格式来查看。

DOT文件可以通过graphviz转换为图片格式来查看，例如将dot转换为png的命令是dot -Tpng *.dot -o *.png。

下面通过一个用例来生成和查看IR图。新建一个test.py的Python文件，如下。

from mindspore import Tensor, context, Parameter, ms_function
import mindspore as ms
import mindspore.nn as nn
import numpy as np
from mindspore.ops import functional as F


context.set_context(mode=context.GRAPH_MODE, save_graphs=True)

def test():
  class TestNet(nn.Cell):
    def __init__(self):
      super(TestNet, self).__init__()
      self.weight = Parameter(Tensor(np.array(0), ms.int32), name="param")

    def construct(self, x):
      out = 0
      i = 0
      while i < 3:
        F.assign(self.weight, i)
        out = x * self.weight + out
        i = i + 1
      return out

  net = TestNet()
  input_x = Tensor(np.array(1), ms.int32)
  res = net(input_x)
  print("res:", res)

设置了保存IR图的命令后（即：save_graphs=True），执行脚本：pytest -s test.py

可以在当前路径下得到多个以.ir为后缀的文件。不同阶段会有不同的IR图对应，从而便于框架开发者和用户了解具体每一步的优化逻辑，便于问题定位。编译流程在前后端会存在更多的优化流程，这些优化流程以现有IR为输入，又以新生成的IR为输出，被称为优化器。优化器负责分析并改进中间表示，极大程度的提高了编译流程的可拓展性，也降低了优化流程对前端和后端的破坏。

对应初级开发者来说，将得到的IR图信息即dot文件，转换成图片格式更便于理解。

首先我们需要了解MindSpore中的ANF文法。MindIR中的ANode对应于ANF的原子表达式，ValueNode用于表示常数值，ParameterNode用于表示函数的形参，CNode则对应于ANF的复合表达式，表示函数调用。

<ANode> ::= <ValueNode> | <ParameterNode>
<ParameterNode> ::= Parameter
<ValueNode> ::= Scalar | Named | Tensor | Type | Shape
               | Primitive | MetaFuncGraph | FuncGraph
<CNode> ::= (<AnfNode> …)
<AnfNode> ::= <CNode> | <ANode>

在得到多个IR文件后，可以逐一打开文件查看文件内容。例如会在路径下看到如下的IR文件。

00_parse_xx.ir
01_symbol_resolve_xx.ir
02_combine_like_graphs_xx.ir
03_inference_opt_prepare_xx.ir
04_abstract_specialize_xx.ir
05_auto_monad_xx.ir
06_inline_xx.ir
07_py_pre_ad_xx.ir
08_pynative_shard_xx.ir
09_pipeline_split_xx.ir
10_optimize_xx.ir
11_py_opt_xx.ir
12_auto_monad_reorder_xx.ir
13_eliminate_forward_cnode_xx.ir
14_eliminate_ad_related_special_op_node_xx.ir
15_validate_xx.ir
16_distribute_split_xx.ir
17_task_emit_xx.ir
18_execute_xx.ir
...

查看00_parse.ir文件可以看到，得到了5张subgraph，subgraph之间存在调用关系，根图为construct_wrapper.1。定义了一个param，也就是外部传入的x。在parse阶段主要做了脚本符号上的简单解析。IR图中还将代码脚本行对应显示，便于理解。

#IR entry      : @construct_wrapper.1
#attrs         :
#Total params  : 1

%para1_x : <null>

#Total subgraph : 5

subgraph attr:
is_while_header : 1
subgraph instance: ↵construct.9 : 0xaaaaf3825310
subgraph @↵construct.9(%para2_фi, %para3_фout) {
  %0([CNode]2) = resolve(CommonOPS: 'Namespace:mindspore._extends.parse.trope', while_cond)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  %1([CNode]3) = resolve(Ast: 'Namespace:mindspore._extends.parse.trope', lt)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(0)/
  %2([CNode]2) = %1(%para2_фi, 3)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  %3([CNode]2) = %0(%2)
      : (<null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  %4([CNode]17) = Switch(%3, @↻construct.18, @↓construct.19)
      : (<null>, <null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  %5([CNode]20) = %4()
      : () -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  Return(%5)
      : (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
}

subgraph attr:
subgraph instance: construct.6 : 0xaaaaf261a490
subgraph @construct.6(%para4_x) {
  %0([CNode]4) = resolve(SymbolStr: 'Namespace:test', F)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  %1([CNode]7) = resolve(ClassMember: 'Namespace:test..<TestNet::281473602514480>', weight)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
  %2([CNode]21) = call @↵construct.9(0, 0)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
  Return(%2)
      : (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
}

subgraph attr:
subgraph instance: ↻construct.18 : 0xaaaaf3825cb0
subgraph @↻construct.18() {
  %0([CNode]5) = getattr($(@construct.6:[CNode]4), assign)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(116)/        F.assign(self.weight, i)/
  %1([CNode]8) = %0($(@construct.6:[CNode]7), $(@↵construct.9:para2_фi))
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(116)/        F.assign(self.weight, i)/
  %2([CNode]10) = stop_gradient(%1)
      : (<null>) -> (<null>)
      # scope: (Default)
  %3([CNode]11) = resolve(Ast: 'Namespace:mindspore._extends.parse.trope', add)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(0)/
  %4(i) = %3($(@↵construct.9:para2_фi), 1)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(118)/        i = i + 1/
  %5([CNode]12) = resolve(Ast: 'Namespace:mindspore._extends.parse.trope', add)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(0)/
  %6([CNode]13) = resolve(Ast: 'Namespace:mindspore._extends.parse.trope', mul)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(0)/
  %7([CNode]14) = %6($(@construct.6:para4_x), $(@construct.6:[CNode]7))
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(117)/        out = x * self.weight + out/
  %8(out) = %5(%7, $(@↵construct.9:para3_фout))
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(117)/        out = x * self.weight + out/
  %9([CNode]15) = call @↵construct.9(%4, %8)
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
  %10([CNode]16) = Depend(%9, %2) primitive_attrs: {side_effect_propagate: 1} cnode_attrs: {topo_sort_rhs_first: true}
      : (<null>, <null>) -> (<null>)
      # scope: (Default)
  Return(%10)
      : (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
}

subgraph attr:
subgraph instance: ↓construct.19 : 0xaaaaf38264a0
subgraph @↓construct.19() {
  Return($(@↵construct.9:para3_фout))
      : (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(119)/      return out/
}

subgraph attr:
subgraph instance: construct_wrapper.1 : 0xaaaaf382c740
subgraph @construct_wrapper.1() {
  %0([CNode]22) = call @construct.6(%para1_x)
      : (<null>) -> (<null>)
      # scope: (Default)
  Return(%0)
      : (<null>)
      # scope: (Default)
      # In file /.../mindspore/mindspore/test.py(115)/      while i < 3:/
}

从MindSpore的pipeline流程中，可以对应看到具体一些优化pass的流程。

// Parse the python ast to ANF graph
  (void)actions.emplace_back(std::make_pair("parse", ParseAction));

  // Resolve the python func
  (void)actions.emplace_back(std::make_pair("symbol_resolve", SymbolResolveAction));

  auto multi_graphs = parallel::CostModelContext::GetInstance()->is_multi_subgraphs();
  if (!multi_graphs && pipeline::GetJitLevel() != "O0") {
    (void)actions.emplace_back(std::make_pair("combine_like_graphs", CombineLikeGraphs));
  }

  (void)actions.emplace_back(std::make_pair("meta_unpack_prepare", MetaUnpackPrepareAction));
  // Evaluate type and shape, and specialize.
  (void)actions.emplace_back(std::make_pair("abstract_specialize", AbstractSpecializeAction));
  // Auto-monad for side-effects handling.
  (void)actions.emplace_back(std::make_pair("auto_monad", AutoMonadAction));
  // Do data structure simplifications and inline.
  (void)actions.emplace_back(std::make_pair("inline", OptInlineAction));
  // Add pre-ad, post-inline python pass stub.
  (void)actions.emplace_back(std::make_pair("py_pre_ad", PreAdActionPyStub));
  // Handle the pynative shard.
  (void)actions.emplace_back(std::make_pair("pynative_shard", PynativeShardAction));
  // Do PipelineSplit action.
  (void)actions.emplace_back(std::make_pair("pipeline_split", PipelineSplitAction));

这样，我们可以将用例脚本到中间语言的解析过程与框架的解析流程对应上。

其中optimize阶段包含很多优化pass，有兴趣研究还可以将相关设置打开，从而把更多的IR图保存下来，对应到optimize中的pass。相关设置开关：

export MS_DEV_DUMP_IR_CONFIG="DISABLE_BACKEND#ENABLE_PASS_IR#LINE_LEVEL1"
export MS_DEV_TRAVERSE_SUBSTITUTIONS_MODE=1

test.py中的用例涵盖了常量折叠和副作用保序等优化pass。常量折叠是指多个量进行计算时，如果能够在编译时刻直接计算出其结果，那么变量将由常量替换。

    def construct(self, x):
      out = 0
      i = 0
      while i < 3:
        F.assign(self.weight, i)
        out = x * self.weight + out
        i = i + 1
      return out


第一遍循环时，out的初始值为0，那么out = x * self.weight + out可以简化为：out = x * self.weight。
同理， i = i + 1可以简化为 i = 1

在多个优化pass处理之后的10_optimize.ir中，我们可以看到IR文件得到进一步简化。IR文件中只有1张subgraph。这是由于脚本中while表达的控制流条件是常量条件，框架那边可以在编译时期对其展开处理，从而简化流程。同时可以看到由于我们使用了全局变量parameter，即会存在副作用相关的表达，需要保证全局变量的读写顺序，IR图中的Load和UpdateState等节点就是为了保证其读写顺序，在此不做展开，有机会单独讨论副作用相关操作，感兴趣的朋友可以关注04_abstract_specialize.ir和05_auto_monad.ir两张IR图之间的差异。

#IR entry      : @22_1_construct_wrapper.112
#attrs         :
check_set_strategy_valid_once_only : 1
#Total params  : 2

%para1_x : <Tensor[Int32], ()>
%para2_param : <Ref[Tensor(I32)], (), ref_key=:param>  :  has_default

#Total subgraph : 1

subgraph attr:
check_set_strategy_valid_once_only : 1
subgraph instance: 22_1_construct_wrapper.112 : 0xaaaaf38f94c0
subgraph @22_1_construct_wrapper.112() {
  %0([CNode]57) = Assign(%para2_param, Tensor(shape=[], dtype=Int32, value=0), U) primitive_attrs: {output_names: [output], side_effect_mem: true, input_names: [ref, value]}
      : (<Ref[Tensor(I32)], (), ref_key=:param>, <Tensor[Int32], (), value=...>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/ops/function/parameter_func.py(55)/    return assign_(variable, value)/
  %1([CNode]108) = UpdateState(U, %0)
      : (<UMonad>, <Tensor[Int32], ()>) -> (<UMonad>)
      # scope: (Default)
  %2([CNode]106) = Load(%para2_param, %1)
      : (<Ref[Tensor(I32)], (), ref_key=:param>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
  %3([CNode]107) = UpdateState(%1, %2)
      : (<UMonad>, <Tensor[Int32], ()>) -> (<UMonad>)
      # scope: (Default)
  %4([CNode]57) = Assign(%para2_param, Tensor(shape=[], dtype=Int32, value=1), %3) primitive_attrs: {output_names: [output], side_effect_mem: true, input_names: [ref, value]}
      : (<Ref[Tensor(I32)], (), ref_key=:param>, <Tensor[Int32], (), value=...>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/ops/function/parameter_func.py(55)/    return assign_(variable, value)/
  %5([CNode]102) = UpdateState(%3, %4)
      : (<UMonad>, <Tensor[Int32], ()>) -> (<UMonad>)
      # scope: (Default)
  %6([CNode]100) = Load(%para2_param, %5)
      : (<Ref[Tensor(I32)], (), ref_key=:param>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
  %7([CNode]101) = UpdateState(%5, %6)
      : (<UMonad>, <Tensor[Int32], ()>) -> (<UMonad>)
      # scope: (Default)
  %8([CNode]57) = Assign(%para2_param, Tensor(shape=[], dtype=Int32, value=2), %7) primitive_attrs: {output_names: [output], side_effect_mem: true, input_names: [ref, value]}
      : (<Ref[Tensor(I32)], (), ref_key=:param>, <Tensor[Int32], (), value=...>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/ops/function/parameter_func.py(55)/    return assign_(variable, value)/
  %9([CNode]97) = UpdateState(%7, %8)
      : (<UMonad>, <Tensor[Int32], ()>) -> (<UMonad>)
      # scope: (Default)
  %10([CNode]95) = Load(%para2_param, %9)
      : (<Ref[Tensor(I32)], (), ref_key=:param>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
  %11([CNode]14) = Mul(%para1_x, %10) primitive_attrs: {output_names: [output], input_names: [x, y]}
      : (<Tensor[Int32], ()>, <Tensor[Int32], ()>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/test.py(117)/        out = x * self.weight + out/
  %12([CNode]14) = Mul(%para1_x, %6) primitive_attrs: {output_names: [output], input_names: [x, y]}
      : (<Tensor[Int32], ()>, <Tensor[Int32], ()>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/test.py(117)/        out = x * self.weight + out/
  %13([CNode]14) = Mul(%para1_x, %2) primitive_attrs: {output_names: [output], input_names: [x, y]}
      : (<Tensor[Int32], ()>, <Tensor[Int32], ()>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/test.py(117)/        out = x * self.weight + out/
  %14([CNode]78) = Add(%12, %13) primitive_attrs: {output_names: [output], input_names: [x, y]}
      : (<Tensor[Int32], ()>, <Tensor[Int32], ()>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/ops/function/math_func.py(280)/    return tensor_add(x, y)/
  %15([CNode]78) = Add(%11, %14) primitive_attrs: {output_names: [output], input_names: [x, y]}
      : (<Tensor[Int32], ()>, <Tensor[Int32], ()>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/ops/function/math_func.py(280)/    return tensor_add(x, y)/
  %16([CNode]96) = UpdateState(%9, %10)
      : (<UMonad>, <Tensor[Int32], ()>) -> (<UMonad>)
      # scope: (Default)
  %17([CNode]16) = Depend(%15, %16) primitive_attrs: {side_effect_propagate: 1}
      : (<Tensor[Int32], ()>, <UMonad>) -> (<Tensor[Int32], ()>)
      # scope: (Default)
  Return(%17)
      : (<Tensor[Int32], ()>)
      # scope: (Default)
      # In file /.../mindspore/test.py(115)/      while i < 3:/
}

参考：

MindSpore IR（MindIR） — MindSpore master documentation

ir指令、立即数的作用_AI框架中图层IR的分析_何老师Matt的博客-CSDN博客

机器学习系统：设计和实现 — 机器学习系统：设计和实现 1.0.0 documentation (openmlsys.github.io)

Abstract A Simple Graph-Based Intermediate Representation.