(TVM开发代码学习)给Relay系统添加新算子

最新推荐文章于 2022-03-24 15:10:35 发布

吴正彪

最新推荐文章于 2022-03-24 15:10:35 发布

阅读量2k

点赞数 2

文章标签：深度学习编译器

本文链接：https://blog.csdn.net/weixin_44114416/article/details/122203220

版权

本文学习的是官网中的Adding an Operator to Relay — tvm 0.9.dev182+ge718f5a8a documentation (apache.org)

添加新的Relay算子需要以下步骤：

1、添加一个attribute node，声明在编译时已知的固定参数。

2、给算子写一个type relation，整合进Relay类的系统文件里。

3、在C++源码中用RELAY_REGISTER_OP宏定义来注册算子的数量、类型等信息。

4、定义算子的计算过程。

5、给算子计算和调度建立连接。

6、定义一个C++函数，产生一个CallNode类，并且注册进python api中。

7、给6中的python api写个简洁接口。

8、给算子做测试。

1 Defining an Attribute Node

attribute node:一般包括算子的基础属性，比如卷积算子的基础属性有stride和dilation。attributes应该定义在： include/tvm/relay/attrs/中。

本例实现一个numpy-like的累乘和累加算子。包含参数：

累乘函数的参数

其AttrsNode类节点的结构体ScanopAttrs已经定义在：include/tvm/relay/attrs/transform.h中。

定义了Scanop的属性类AttrsNode

可以看到，定义的累乘或者累加算子的属性有axis，dtype和exclusive(该布尔值决定是否把当前值纳入累乘计算里)。通过TVM_DECLARE_ATTRS宏定义实现了AttrsNode的注册。

2 Writing a Type Relation

所谓的type relation，就是这些关系表示为接受输入类型和输出类型列表（可以是不完整的类型列表），并返回满足关系的输入和输出类型列表的函数。这包括可以在编译时静态确定的形状信息。本质上，除了计算输出类型之外，运算符的关系还可以强制执行所有必要的类型规则（即通过检查输入类型）。

这话直接翻译得比较抽象，看代码。

累乘和累加运算符的type relation可以在 src/relay/op/tensor/http://transform.cc 中找到。

TVM_REGISTER_NODE_TYPE(ScanopAttrs);
bool ScanopRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
               const TypeReporter& reporter) {
  // types: [data, output]
  ICHECK_EQ(types.size(), 2) << "Expects two types, one for the input and another for the output";
  const auto* data = types[0].as<TensorTypeNode>();
  if (data == nullptr) {
    ICHECK(types[0].as<IncompleteTypeNode>())
        << "Scanop: expect input type to be TensorType but get " << types[0];
    return false;
  }

  const auto* param = attrs.as<ScanopAttrs>();

  auto dtype = param->dtype;
  if (dtype.is_void()) {
    dtype = data->dtype;
  }

  if (param->axis.defined()) {
    reporter->Assign(types[1], TensorType(data->shape, dtype));
  } else {
    auto prod = data->shape[0];
    for (size_t i = 1; i < data->shape.size(); ++i) {
      prod = prod * data->shape[i];
    }
    reporter->Assign(types[1], TensorType({prod}, dtype));
  }

  return true;
}

可以看到，这个type relation函数——ScanopRel()，输入了一个包含输入输出类型的二维数组types，以及Attrs对象attrs，还有一个TypeReporter对象。在函数体中，主要处理的对象就是这个TypeRepoter& reporter:

repoter涉及了输入输出类型与形状的绑定

reporter->Assign()报告了输入输出类型的强制转换记录，其中，Assign函数用于创造类型约束：

TypeReporterNode的Assign()

而TypeRepoter则重载了->运算符，将Assign()中创建类型约束实现了源类型到目标类型（即输入和输出类型）的转换。

实现了src类型到dst类型的强制转换

注意，在TypeReporter源代码（\include\tvm\ir\type_relation.h）中，官方还注释了其调用规则：

“If the input type information can be used to fully decide the IncompleteTypes, then the function should call reporter. Assign to report the new types, and return true. Otherwise, the function should return false.”

翻译：如果输入的类型信息可以用来完全确定不完整的类型，那么函数应该调用reporter，Assign以报告新类型，并返回true。否则，函数将返回false。这就是我们示例函数ScanopRel()做的事情。

感觉这一步官方介绍得不太清晰，我的理解也并不深入，关于算子的输入输出关系可能还得多看几个算子源代码的定义才能真正理解。

3 Relating the Arity and Attributes to an Operation

使用C++中的RELAY_REGISTER_OP宏注册以下信息：

参数数量
参数的名称和描述
支持的等级（1表示内部实现;较高的数字表示较少的内部支持或外部支持的算子）
上一步写好的type relation——ScanopRel()函数
其他优化时可能需要的注释

其注册代码在src/relay/op/tensor/http://transform.cc

注册算子属性信息，在步骤6中会调用到我们的cumsum或者cumprod属性

这一步就是单纯的注册好算子的信息进TVM的ffi系统里。代码中的TOpPattern 是向编译器提示运算符所做的计算模式的提示，这对于融合运算符可能很有用。kOpaque 告诉 TVM 不要费心去尝试融合这个操作符。

4 Defining the Compute of the Operation

这部分需要参考te和TOPI的教程，并查看在python/tvm/topi/scan.py里的CPU实现，GPU版本的实现参考python/tvm/topi/cuda/scan.py。

简单介绍一下，te就是tensor experssion，张量表达式，我理解为TVM内部记录计算过程的python语言格式。而topi就是TVM Operator Inventory，TVM算子库，里面存储了算子的计算步骤。具体可以查看：

Introduction to TOPI — tvm 0.9.dev182+ge718f5a8a documentation (apache.org)

CPU实现版本python/tvm/topi/scan.py：

先定义了一个scanop()函数，这个函数就是累加和累乘的模板，实现的是"逐元素"这一步骤。

累加累乘的实现中间函数scanop()

然后这个scanop()函数，函数中处理了一些输入输出寄存器，最后是返回一个extern()，称其为外部函数。

套娃+1,extern()

extern()函数定义在python/tvm/te的operation.py中。

在extern里，会调用python端的_ffi_api，将算子(算子会选择add或者multiply)通过ffi系统传递进去。TVM的FFI(Foreign Function Interface)的实现非常巧妙，以后有空开坑学习一下，这里咋们只要知道涉及这个ffi的东西就是实现了python中调用C++代码就行了。

extern()内部的ffi调用

注意上图的body，就是张量表达式定义时的lambda函数（如果不知道是啥意思,搜索关键字tvm.te.compute，简单来说这个lambda函数就是定义了张量表达式的计算)。这个lambda函数可以传进codegen中生成底层代码。

而在本例中，lambda函数是用在scanop()函数中定义的gen_ir()函数实现的，属于TIR模块的知识范畴。

gen_ir()生成lambda函数来定义计算，注意倒数第二行的binop

其中，binop会在后面的cumsum或者cumprod函数中，决定是加还是乘。会选择tir.generic.add

套娃再加1，这里选择generic.add

整理一下，其实这就是scanop()和extern()以及cumsum()套娃，scanop()完成“逐元素”这一步骤，cumsum或者cumprod函数实现调用ffi的乘或者加算子（实现步骤是参数binop取generic.add或generic.multiply），generic这一模块就是实现了ffi调用，使得python端能通过ffi系统调用c++源代码中的add算子，最后在scanop()中实现累加。