如何添加TpuLang接口

Viewed 377
developer tpulang
zhang1
1 Answers

一. 选择一个TpuLang接口任务，比如Conv接口

二. 查看接口定义以及参数说明

以Conv_v2为例，接口定义如下：

     def conv_v2(tensor_i,
                 weight,
                 bias = None,
                 stride = None,
                 dilation = None,
                 pad = None,
                 group = 1,
                 input_zp = None,
                 weight_zp = None,
                 out_dtype = None,
                 out_name = None):
        # pass
参数说明
* tensor_i：Tensor类型，表示输入Tensor，4维NCHW格式。
* weight：Tensor类型，表示卷积核Tensor，4维[oc, ic, kh, kw]格式。其中oc表示输出Channel数，ic表示输入channel数，kh是kernel_h，kw是kernel_w。
* bias：Tensor类型，表示偏置Tensor。为None时表示无偏置，反之则要求shape为[1, oc, 1, 1]。
* dilation：List[int]，表示空洞大小，取None则表示[1,1]，不为None时要求长度为2。List中顺序为[长，宽]
* pad：List[int]，表示填充大小，取None则表示[0,0,0,0]，不为None时要求长度为4。List中顺序为[上， 下， 左， 右]
* stride：List[int]，表示步长大小，取None则表示[1,1]，不为None时要求长度为2。List中顺序为[长，宽]
* groups：int型，表示卷积层的组数。若ic=oc=groups时，则卷积为depthwise conv
* input_zp：List[int]型或int型，表示输入偏移。取None则表示0，取List时要求长度为ic。
* weight_zp：List[int]型或int型，表示卷积核偏移。取None则表示0，取List时要求长度为ic，其中ic表示输入的Channel数。
* out_dtype：string类型或None，表示输出Tensor的类型。输入tensor类型为float16/float32时，取None表示输出tensor类型与输入一致，否则取None表示为int32。取值范围：/int32/uint32/float32/float16
* out_name：string类型或None，表示输出Tensor的名称，为None时内部会自动产生名称。
返回值
返回一个Tensor，该Tensor的数据类型由out_dtype确定。

可以看到conv接口有11个输入，在参数说明中有参数的详细描述。
其中tensor_i，weight，bias为输入tensor，其余为输入参数。
返回值为tensor，其data type和name分别由out_dtype，out_name确定。
11个输入中，tesor_i和weight是必要参数，其他参数可以不给，其默认值在参数说明中说明。

三. 查看top层conv定义

查看TopOps.td，找到对应的conv对应的arguments和results定义：

let arguments = (ins
    AnyTensor:$input,
    AnyTensor:$filter,
    AnyTensorOrNone:$bias,
    I64ArrayAttr:$kernel_shape,
    I64ArrayAttr:$strides,
    I64ArrayAttr:$pads, // top,left,bottom,right
    DefaultValuedAttr<I64Attr, "1">:$group,
    OptionalAttr<I64ArrayAttr>:$dilations,
    OptionalAttr<I64ArrayAttr>:$inserts,
    DefaultValuedAttr<BoolAttr, "false">:$do_relu,
    DefaultValuedAttr<F64Attr, "-1.0">:$relu_limit
  );

  let results = (outs AnyTensor:$output);

可以看到输入为3个tensor。bias为tensor或none。
kernel_shape，strides，pads，dilations，inserts为I64ArrayAttr，int64 array类型。
group为I64Attr，int64类型。
do_relu为BoolAttr，bool类型。
relu_limit为F64Attr，float64类型。
输出为一个tensor类型。

四.添加conv_v2接口

在tpu-mlir/python/transform/TpuLang.py中添加接口

1）添加接口定义
def conv_v2(input: Tensor,
        weight: Tensor,
        bias: Tensor = None,
        stride: List[int] = None,
        dilation: List[int] = None,
        pad: List[int] = None,
        group=1,
        input_zp: Union[int, List[int]] = None,
        weight_zp: Union[int, List[int]] = None,
        out_dtype: str = None,
        out_name: str = None):

输入参数与定义一致，并在接口添加了输入类型的约束。

2)设置默认参数
    dilation = [1, 1] if dilation is None else dilation
    stride = [1, 1] if stride is None else stride
    pad = [0, 0, 0, 0] if pad is None else pad

如果输入为None，则给定默认值。

3) 将参数换为attr
    attr = {
        "kernel_shape": ArrayAttr(weight.shape[2:]),
        "strides": ArrayAttr(stride),
        "dilations": ArrayAttr(dilation),
        "pads": ArrayAttr(pad),
        "do_relu": Attr(False, "bool"),
        "group": Attr(group)
    }

该接口需要将参数转换为tpu-mlir中可以识别的参数，通过ArrayAttr和Attr可以将参数转为tpu-mlir可以识别的参数。ArrayAttr和Attr的第一个参数为数据，第二个参数为类型，默认类型为“int64”。

4) 设置量化参数
    input.quantization(zero_point=input_zp)
    weight.quantization(zero_point=weight_zp)

设置输入或输出tensor的量化参数(scale, zero_point)。

5) 生成输出tensor
    o_dtype = "int32"
    if out_dtype is not None:
        o_dtype = out_dtype
    elif input.dtype == "float32" or input.dtype == "float16":
        o_dtype = input.dtype
    def _shape_inference():
        kh_ext = dilation[0] * (weight.shape[2] - 1) + 1
        kw_ext = dilation[1] * (weight.shape[3] - 1) + 1
        oh = (input.shape[2] + pad[0] + pad[1] - kh_ext) // stride[0] + 1
        ow = (input.shape[3] + pad[2] + pad[3] - kw_ext) // stride[1] + 1
        return [input.shape[0], weight.shape[0], oh, ow]
    output = Tensor(_shape_inference(), dtype=o_dtype, name=out_name)

根据参数，新建输出tensor，此处需要对输入tensor进行type_inference和shape_inference(未来可能去掉type_inference和shape_inference)

6) 注册op并返回output
    inputs = [input, weight, bias]
    TpuLang.insert_op(Top.ConvOp, inputs=inputs, outputs=[output], params=attr)
    return output

给定op_name, inputs, outputs, params。将op注册到tpulang中。

五. 添加conv_v2单元测试

1) 注册单元测试
class TPULANG_IR_TESTER(object):
    # This class is built for testing single operator transform.
    def __init__(self):
        self.test_function = {
            #############################
            # TpuLang Test Case, Alphabetically
            #############################
            "Conv2d": self.test_Conv2d,
            "HModel": self.test_Model,
        }

在self.test_function中注册Conv2d测试

2) 创建测试用例
def test_Conv2d(self, case_name):
    """Conv 2D"""
    @tpulang
    def _test_convolution(
        input_shape:List[int], kernel_shape:List[int], stride:List[int]=[1,1],
        dilation:List[int]=[1,1], pad:List[int]=None, group=1, dtype="float32",
        zp:List[int]=[None,None]
    ):
        x_data = rand_data(input_shape, dtype)
        x = tpul.Tensor(dtype=dtype, shape=input_shape, data=x_data)
        conv = self.conv_op(x, kernel_shape, stride, pad, group=group, dilation=dilation, zp=zp, dtype=dtype)
        tpul.compile(case_name, [x], [conv], False, 2)

    _test_convolution([1, 3, 28, 28], [12, 3, 1, 1], group=3)
    _test_convolution([1, 3, 32, 32], [12, 3, 3, 3], stride=[2,2], pad=[1,1,1,1])
    _test_convolution([1, 3, 32, 32], [12, 3, 3, 3], stride=[2,2], pad=[1,1,1,1], dtype="int8", zp=[5, -8])

新建测试用例步骤如下：

  1. _test_convolution参数可以根据测试参数自主选择，尽可能覆盖较多的参数。
  2. 创建输入tensor：x = tpul.Tensor(...)
  3. 调用接口函数conv = self.conv_op(...)。此处对conv_v2接口进行封装。
  4. 调用编译命令tpul.compile(name，inputs, outputs,...)

这里的conv_op定义如下：

def conv_op(self, x, kshape, stride, pad=None, group=1, dilation=[1,1], zp=[None, None], dtype="float32"):
    oc = kshape[0]
    weight = self.coeff_tensor(kshape, dtype)
    out_dtype =  dtype if dtype == 'float32' else 'int32'
    bias = self.coeff_tensor(oc, out_dtype)
    conv = tpul.conv_v2(x, weight, bias=bias, stride=stride, pad=pad,
                        dilation=dilation, group=group, input_zp=zp[0],
                        weight_zp=zp[1], out_dtype=out_dtype)

其封装了conv_v2接口，新建了weight和bias coeff tensor。

3) 调用测试命令，测试conv_v2接口
  1. 进入docker环境
  2. source envsetup.sh
  3. 执行命令 test_tpulang.py --case Conv2d

结果如下：

test_tpulang.py Conv2d

no found SET_CHIP_NAME environment value, set bm1684x as default
Test: Conv2d
Save mlir file: Conv2d_0_origin.mlir
[Running]: tpuc-opt --init --shape-infer --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Conv2d_0_origin.mlir -o Conv2d_0.mlir 
[Success]: tpuc-opt --init --shape-infer --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Conv2d_0_origin.mlir -o Conv2d_0.mlir 
Mlir file generated:Conv2d_0.mlir
[CMD]: model_runner.py --input Conv2d_0_in_f32.npz --model Conv2d_0.mlir --output Conv2d_0_top_outputs.npz
Save mlir file: Conv2d_1_origin.mlir
[Running]: tpuc-opt --init --shape-infer --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Conv2d_1_origin.mlir -o Conv2d_1.mlir 
[Success]: tpuc-opt --init --shape-infer --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Conv2d_1_origin.mlir -o Conv2d_1.mlir 
Mlir file generated:Conv2d_1.mlir
[CMD]: model_runner.py --input Conv2d_1_in_f32.npz --model Conv2d_1.mlir --output Conv2d_1_top_outputs.npz
Save mlir file: Conv2d_2_origin.mlir
[Running]: tpuc-opt --init --shape-infer --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Conv2d_2_origin.mlir -o Conv2d_2.mlir 
[Success]: tpuc-opt --init --shape-infer --canonicalize --mark-FLOPs --save-weight --mlir-print-debuginfo Conv2d_2_origin.mlir -o Conv2d_2.mlir 
Mlir file generated:Conv2d_2.mlir
[CMD]: model_runner.py --input Conv2d_2_in_f32.npz --model Conv2d_2.mlir --output Conv2d_2_top_outputs.npz
====== TEST Conv2d Success ======
zhang1

如果有问题和疑问，请直接留言回复

zhang•