TVM学习(七)算子

作者:安平博,Xilinx高级工程师;来源:AI加速微信公众号

relay算子

上一章梳理了一遍TVM前端流程,前端完成了tensorflow算子到TVM算子的转换。这一章CNN网络中最普遍的卷积运算为例,看一下relay算子的表示。在python/tvm/relay/frontend/tensorflow.py文件中convert_map有:

_convert_map = {
…
'Conv2D'                            : _conv('conv'),
…
}

在_conv函数中会根据layout对weights,inputs,outputs进行重排,然后调用AttrCvt来获得op。

out = AttrCvt(op_name=_dimension_picker('conv',

                                      surfix="_transpose" if opname == 'conv_transpose' else ""),

            ignores=['explicit_paddings'],

            transforms={

                'kernel_shape': 'kernel_size',

                'data_format': 'data_layout',

                'dilations': ('dilation', (0, 0)),

                'group': ('groups', 1)},

            custom_check=_dimension_constraint())([inputs_data, inputs[1]], attr)

AttrCvt的调用位于python/tvm/relay/frontend/common.py文件夹中,根据注释看出这个类主要实现算子转换,实际上是根据传入的op_name映射到relay的算子。首先会对传入的attrs进行检测,如果有不符合的属性会报错或者warning,如果属性有相应转换就进行新属性替换,最后调用get_relay_op。在这个函数中可以看到,依据op_name在全局字典_op中搜索相应op,然后返回。所有的op都位于python/tvm/relay/op包中,conv在op/nn中定义。nn.py中包含如下调用关系:conv2d -> _make.conv2d()。在_make.py中实际上实现了C++类到python类的注册,就是一行代码:

tvm._ffi._init_api("relay.op.nn._make", __name__)

_init_api函数在python/tvm/_ffi/registry.py中,我们可以看一下:

def _init_api(namespace, target_module_name=None):

   

    target_module_name = (

        target_module_name if target_module_name else namespace)

    if namespace.startswith("tvm."):

        _init_api_prefix(target_module_name, namespace[4:])

    else:

        _init_api_prefix(target_module_name, namespace)

 

 

def _init_api_prefix(module_name, prefix):

    module = sys.modules[module_name]

 

    for name in list_global_func_names():

        if not name.startswith(prefix):

            continue

 

        fname = name[len(prefix)+1:]

        target_module = module

 

        if fname.find(".") != -1:

            continue

        f = get_global_func(name)

        ff = _get_api(f)

        ff.__name__ = fname

        ff.__doc__ = ("TVM PackedFunc %s. " % fname)

        setattr(target_module, ff.__name__, ff)

实际上是通过名字获取C++注册的函数,然后设置给到_make.py文件中。这样就相当于_make.py文件中定义了conv2d相关的函数了。

接下来我们转移到C++中,看一看conv算子的实现。

在src/relay/op/nn/convolution.cc中有以下代码:

TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d").set_body_typed([](Expr data, Expr weight, Array strides, Array padding,  Array dilation, int groups, IndexExpr channels, Array kernel_size, String data_layout, String kernel_layout, String out_layout, DataType out_dtype) {

      return MakeConv(data, weight, strides, padding, dilation, groups, channels, kernel_size, data_layout, kernel_layout, out_layout, out_dtype, "nn.conv2d");

    });



RELAY_REGISTER_OP("nn.conv2d")

    .describe(R"code(2D convolution layer (e.g. spatial convolution over images).

 

This layer creates a convolution kernel that is convolved

with the layer input to produce a tensor of outputs.

 

- **data**: This depends on the `layout` parameter. Input is 4D array of shape

            (batch_size, in_channels, height, width) if `layout` is `NCHW`.

- **weight**: (channels, in_channels, kernel_size[0], kernel_size[1])

- **out**:  This depends on the `layout` parameter. Output is 4D array of shape

            (batch_size, channels, out_height, out_width) if `layout` is `NCHW`.

 

)code" TVM_ADD_FILELINE)

    .set_attrs_type()

    .set_num_inputs(2)

    .add_argument("data", "Tensor", "The input tensor.")

    .add_argument("weight", "Tensor", "The weight tensor.")

    .set_support_level(2)

    .add_type_rel("Conv2D", Conv2DRel)

.set_attr("FInferCorrectLayout", ConvInferCorrectLayout);

C++通过宏定义TVM_REGISTER_GLOBAL将算子注册到一个全局对象中。可以看一下这个宏定义:

#define TVM_REGISTER_GLOBAL(OpName) \

  TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register(OpName)

可以看到注册实现在Registry类中,其中有一个Register函数,这个函数会通过全局manager来将算子注册进去。Set_body会将lamda以及普通函数设置给registry中的统一函数封装形式。

Set_body中将通过MakeConv构建一个conv算子,然后注册到registry中。在MakeConv中,首先根据传入的conv参数,包括strides,kernel,layout等,构建atrrs对象,然后根据op的名字从已经注册过的conv算子中得到conv的算子,然后再将attrs和op一起打包到call类中。

Op算子是通过RELAY_REGISTER_OP注册到一个公共AttrRegistry中的。在一个op类中实际上并没有包含这个op的计算过程,只是纳入了这个算子的输入输出以及属性的信息。

TOPI算子

TOPI是TVM自身的一个算子库,这些算子可以通过te来进行表达,类似于numpy的方式。比如对于numpy有np.sum(),同样tvm也可以有te.sum这样的表示。这为通过tvm语言来构造计算图结构提供了方便。我们用官方的一个例子来深入追踪一下topi算子的代码。

n = te.var("n")
m = te.var("m")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), "k")
B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")
s = te.create_schedule(B.op)
print(tvm.lower(s, [A], simple_mode=True))

输出的函数是这样的:

primfn(A_1: handle) -> ()  attr = {"global_symbol": "main", "tir.noalias": True}  
buffers = {A: Buffer(A_2: Pointer(float32), float32, [n: int32, m: int32], [stride: int32, stride_1: int32], type="auto")}  
buffer_map = {A_1: A} {  attr [B: Pointer(float32)] "storage_scope" = "global";  
allocate(B, float32, [n]);  
for (i: int32, 0, n) {    
B[i] = 0f32    
for (k: int32, 0, m) {      
B[i] = ((float32*)B[i] + (float32*)A_2[((i*stride) + (k*stride_1))])    }  
}
}

Var就类似于tensorflow中variable,创建了一个变量。其调用链为python/tvm/te/operation.py -> python/tvm/tir/expr.py -> src/tir/ir/expr.cc。var继承了PrimExpr类,建立var的时候创建了VarNode。VarNode中保存了变量的类型,名字等信息。

Placeholder也类似tensorflow中的占位符,实际上最终创建了一个PlaceholderOp,保存了名字,shape,dtype信息。

接下来重点看一下compute这个算子。Python调用位于python/tvm/te/operation.py中。这个主要是实现lamba函数算子的转换。

def compute(shape, fcompute, name="compute", tag="", attrs=None):

fcompute是对应着lamba表达式。首先从fcompute的__code__中提取出变量名称和数目信息,然后对应每个输入变量和对应的shape信息一起创建IterVar。IterVar有点像for循环中的循环变量,这里做了IterVar的抽象。然后将IterVar传入fcompute创建了函数体。然后就是根据fcompute的类型建立ComputeOp或者TensorComputeOp。这里我们追踪一下ComputeOp的实现。在src/te/operation/compute_op.cc中。也是构建一个ComputeOpNoe。并记录下数据,表达式信息。

dim_var = [tvm.tir.IterVar((0, s), x, 0) for x, s in zip(arg_names, shape[:out_ndim])]

body = fcompute(*[v.var for v in dim_var])

    if isinstance(body, _tensor.TensorIntrinCall):

        for i, s in enumerate(shape[out_ndim:]):

            var_name = "ax" + str(i)

            dim_var.append(tvm.tir.IterVar((0, s), var_name, 4))

        op_node = _ffi_api.TensorComputeOp(name,

                                           tag,

                                           dim_var,

                                           body.reduce_axis,

                                           out_ndim,

                                           body.intrin,

                                           body.tensors,

                                           body.regions,

                                           body.scalar_inputs)

    else:

        if not isinstance(body, (list, tuple)):

            body = [body]

        body = convert(body)

        op_node = _ffi_api.ComputeOp(

            name, tag, attrs, dim_var, body)

 

num = op_node.num_outputs

outputs = tuple(op_node.output(i) for i in range(num))

return outputs[0] if num == 1 else outputs

最新文章