TVM自定义算子+测试

定义算子

TVM分为编译器前端和编译器后端。编译器前端主要解决的事情是将其他框架（pytorch等）转换成自己的symbolic的表示方式，所有的基本操作都会转换成sym.基本算子的形式，然后通过注册的形式，使得sym.基本算子调用后端的算子实现。编译器后端的算子实现分为两部分，第一部分是compute的部分，compute定义了基本的计算规则，schedule定义了基本计算规则在不同的硬件上的实际执行的方式。

编译器后端

注册算子基本上在topi/python/topi/nn里面。

算子的测试

变量的声明

例如对于一个基本的conv2d的算子而言，首先要声明其输入。tvm里面的声明有点像占位符的感觉，得提前说明所有tensor的大小，类型。例如声明Input和Filter。

// An highlighted block
A = tvm.placeholder(shape=(n, ic, ih, iw), dtype=dtype, name='A');
B = tvm.placeholder(shape=(oc, ic, kh, kw), dtype=dtype, name='B')

A表示一个数据布局为(batch, input_channel, input_height,input_width)的四维Tensor。B表示一个数据布局为(channel_multiply, input_channel, kernel_height,kernel_width)的四维Tensor.

调用底层算子

调用底层的topi算子，返回了一个output,类型是Tvm.tensor，存储的是计算的规则。也就是调用compute的规则之后生成的格式。s表示调用默认的schedule后生成的计算规则，也就是for具体循环的排列。具体可以调用 tvm.lower来查看。

// An highlighted block
output = conv2d_nchw(Input = A, Filter = B, stride = (stride_h, stride_w), padding = (pad_h, pad_w), dilation = (dilation_h, dilation_w))
s = tvm.create_schedule(output.op)
print(tvm.lower(s, [A,</