更新时间:2021-03-18 GMT+08:00
分享

计算实现

计算实现包含依赖python模块导入、算子函数声明、算子入参校验、计算逻辑实现以及调度与编译。计算实现时,可以通过一定的方法进行精度与性能上的提升。

导入Python模块

进行TBE DSL算子开发时,首先需要在算子实现文件中导入昇腾AI软件栈提供的Python模块,代码示例如下所示,算子实现文件的命名请参见1

如果您进行代码实现时依赖了其他自行引入的Python模块,请自行进行依赖导入。

import te.lang.cce
from te import tvm
from te.platform.fusion_manager import fusion_manager
from topi import generic

其中:

  • “te.lang.cce”:引入TBE支持的特定域语言接口,包括常见的运算vmuls、vadds、matmul等。

    具体的接口定义可查看ATC安装路径下“/python/site-packages/te/te/lang/cce/”目录下的python函数,使用方法请参见TBE DSL API

  • “te.tvm”:引入TVM后端代码生成机制。

    具体的接口定义可查看ATC安装路径下“/python/site-packages/te/te/tvm”目录下的python函数,使用方法请参见https://docs.tvm.ai/

  • “te.platform.fusion_manager.fusion_manager”:提供了实现算子的UB自动融合的接口。

    具体的接口定义可查看ATC安装路径下/python/site-packages/te/te/platform/fusion_manager.py”文件中的fusion_manager函数的定义。

  • “topi.generic”:提供了自动算子调度接口。

    具体的接口定义可查看ATC安装路径下“/python/site-packages/topi/topi/generic”目录下的所有python函数,使用方法请参见TBE DSL API

算子函数声明

算子的代码实现中包括两个函数:算子接口函数与算子compute函数,算子的compute函数会在算子接口函数中被调用。

下面详细介绍这两个函数的声明规则。

  • 算子接口函数声明

    如下所示,一个算子的接口函数中包含了算子输入信息、算子输出信息以及内核名称,函数的声明信息需要与算子信息定义文件中的信息对应。

    def operationname(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="KernelName")
    • 算子接口函数名称operationname当前版本请与算子实现文件名称保持一致,命名规则请参见1
    • input_x1, input_x2:算子的输入tensor,每个tensor需要采用字典的形式进行定义,包含shape、ori_shape、format、ori_format与dtype信息,例如:

      dict input_x1 = {'shape' : (2,2), 'ori_shape' : (2,2), 'format': 'ND', 'ori_format':'ND', 'dtype' : 'float16'}

      输入tensor的名称、个数及顺序需要与算子信息定义文件/tbe/op_info_cfg/ai_core/<soc_version>/operationname.ini中的定义保持一致。

    • output_y:算子的输出tensor,包含shape和dtype等信息,字典格式,此字段为预留位。
    • attribute1attribute2...:算子的属性,此处需要为算子的属性赋默认值,算子属性的名称、个数与顺序需要与算子信息定义文件/tbe/op_info_cfg/ai_core/<soc_version>/operationname.ini中的定义保持一致。

      若算子无相关属性信息,此参数忽略。

    • kernel_name:算子在内核中的名称(即生成的二进制文件与算子描述文件的名称),用户自定义,保持唯一,只能是大小写字母、数字、“_”的组合,且必须是字母或者“_”开头,长度小于或等于200个字符。

    不带属性的sqrt算子的接口函数声明如下:

    def sqrt(input_x, output_y, kernel_name="sqrt"):

    带属性的reduce_sum算子的接口函数声明如下:

    def reduce_sum(x, y, axis=None, keep_dims=None, kernel_name="reduce_sum")
  • compute函数声明

    如下所示,compute函数的入参是计算过程中涉及的所有输入tensor、attribute,和输出tensor,以及内核名称。

    @fusion_manager.register("KernelName")
    def operationname_compute(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="KernelName")
    • 装饰器@fusion_manager.register("KernelName")是算子计算实现声明中必需的,其作用是整网运行时支持算子做UB自动融合,使得当前自定义算子可以在UB中根据UB融合规则自动与其他算子的计算进行组装。
    • input_x1, input_x2:compute函数的入参,为在 算子接口函数声明中声明的输入tensor对应的placeholder,包含shpe和dtype等信息。
    • output_y,attribute1=None,xxx等参数,都是从 算子接口函数声明中的算子接口函数中透传过来的,与算子接口函数的声明保持一致即可。

    例如,对于sqrt算子,compute函数定义如下:

    @fusion_manager.register("sqrt")
    def sqrt_compute(input_data, output_data, kernel_name="sqrt"):

    对于reduce_sum算子,算子接口和计算函数定义如下:

    @fusion_manager.register("reduce_sum")
    def reduce_sum_compute(x, y, axis=None, keep_dims=None, kernel_name="reduce_sum")

算子函数实现

完成算子函数声明后,就要具体实现算子接口函数和compute函数。

  1. 首先在算子接口函数operatorname( )中,获取算子输入tensor的shape以及dtype,并可自行实现基本的校验功能。

    1. 获取算子输入tensor的shape以及dtype,为后续定义输入tensor的张量占位符做准备。
      def add(input_x, input_y, output_z, kernel_name="add"):
          shape_x = input_x.get("shape")
          shape_y = input_y.get("shape")
          input_data_type = input_x.get("dtype").lower()
          input_data_type_y = input_y.get("dtype").lower()
    2. (可选)在算子实现函数中添加算子输入/输出及属性基本信息校验,有助于在算子编译阶段,提前发现问题。

      例如,对于Add算子,首先校验两个输入的dtype是否一致,然后校验输入的数据类型是否在允许的数据类型列表中,代码实现如下所示。

      1
      2
      3
      4
      5
      6
      7
          if input_data_type != input_data_type_y:
              raise RuntimeError(
                  "the input_x and input_y should have the same data type.")
          check_tuple = ("float16", "float32", "int32")
          if input_data_type not in check_tuple:
              raise RuntimeError("only support %s while dtype is %s" %
                                 (",".join(check_tuple), input_data_type))
      

      开发者可根据算子特点自定义实现校验函数。

  2. 然后根据shape与dtype定义好输入tensor的张量占位符。

    例如:
    data_input = tvm.placeholder(shape, name="data_input", dtype=dtype)

    使用TVM的placeholder接口对输入tensor进行占位,返回一个tensor对象,此位置中的数据在程序运行时才被指定。

    注意:

    调度与编译中的tensor_list的输入tensor需要是tvm.placeholder接口返回的tensor对象,所以此对象在后续计算过程实现中不能被替换。

    如下所示:

        #返回占位的data_input
        data_input = tvm.placeholder(shape, name='data', dtype=dtype)      
         if dtype == "float16":
             #将data_input的类型转换为float32,然后重新赋值给data_input,此时data_input的内容已经被改变
             data_input = te.lang.cce.cast_to(data_input, "float32")        
            ......
        with tvm.target.cce():
            schedule = generic.auto_schedule(res)
        config = {"print_ir":need_print, 
                        "need_build":need_build,
                        "name":kernel_name, 
                        "tensor_list":[data_input,res]} 
        te.lang.cce.cce_build_code(schedule,config)

    以上代码中,通过data_input = te.lang.cce.cast_to(data_input, "float32")转换数据类型后,placeholder返回的data_input对象已经被覆盖,编译配置tensor_list中的data_input已经不是原placeholder接口返回的tensor,此时算子实现代码编译时会出现以下错误:

    所以可重新定义一个tensor用于存储转换数据类型后的输入进行计算,如下所示:

    data_input1 = te.lang.cce.cast_to(data_input, "float32")

    或者如步骤3所示,计算过程在compute函数中进行,将placeholder返回的输入tensor通过形参传入compute函数进行计算,会生成新的地址用于计算,也可避免placeholder返回的tensor对象被覆盖的情况。

  3. 在算子接口定义函数中调用compute函数进行计算过程的描述。

    例如:

    res = add_compute(data_x, data_y, output_z, kernel_name)

    输入tensor为使用tvm.placeholder定义的占位tensor,其他为算子接口函数透传的参数。

  4. 算子compute函数的实现。

    在compute函数中,完成算子的计算过程,计算过程的实现主要根据算子分析中的TBE DSL API进行代码开发。

    下面我们以计算公式较复杂的relu算子为例,讲解算子的计算过程的实现以及部分DSL接口在使用过程中的注意事项。

    假设通过进行算子分析,得到relu算子的计算公式如下:

    计算实现代码如下所示:

    @fusion_manager.register("relu")
    def relu_compute(x, y, kernel_name="relu"):
        inp_dtype = x.dtype     # 获取输入数据的数据类型
        shape = x.shape         # 获取输入数据的形状
    
        # 若数据类型为float32、int32,使用vmax操作,避免精度损失
        if inp_dtype in ("float32", "int32"):
            tensor_zero = te.lang.cce.broadcast(tvm.const(CONST_ZERO, inp_dtype),shape)     # 返回形状与输入数据相同,每一个元素都为0,每一个元素的数据类型都为输入数据的数据类型的tensor
            data_res = te.lang.cce.vmax(x, tensor_zero)    # 取x与tensor_zero中的大值
        else:
            data_res = te.lang.cce.vrelu(x)      # 若数据类型为int8、float16,直接做relu操作。
    
        data_res = te.lang.cce.cast_to(data_res, inp_dtype)
    
        return data_res

    • 由于te.lang.cce.vrelu( )接口会将int8、uint8、int32、float32的数据类型转换为float16,而int32、float32进行数据类型转换时会造成精度损失,所以为了避免精度损失,对于这两种数据类型的输入,采用te.lang.cce.vmax( )接口取输入数据与0之间的大值。
    • TBE DSL中vmax接口要求两个输入tensor的shape相同,一般使用te.lang.cce.broadcast接口将输入tensor的shape广播到相同shape,一般取两个输入tensor的shape中每个维度的大值组成的shape。

    算子计算函数实现中的其他小技巧:

    • 若输入tensor数据类型不是float32,可以将其转换为float32进行计算,可以提高中间计算结果的精度,最后的结果输出时需要将数据类型转换成原数据类型。
    • 当算子的计算过程比较繁琐时,可以通过抽调内部函数的方法保持每个模块的简洁性和可读性。

(可选)op_select_format函数实现

开发者可以在算子实现文件中实现op_select_format函数,推导出算子的输入输出dtype与format,则后续进行算子信息定义时无需配置输入输出的dtype与format;若算子实现文件中不实现此函数,则后续进行算子信息定义时需要配置输入输出支持的dtype与format,算子信息定义的配置可参见算子信息定义

op_select_format函数的声明如下所示:

def op_select_format(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="xx"):

op_select_format函数的入参和算子接口函数保持一致(即算子的输入、输出、属性及kernel_name),出参为包含了当前算子输入输出支持的format和data type列表的字符串,字符串格式如下所示:

{
"input0": {
"name": "x",
"dtype": "float16,float16,int8,int8",
"format": "NC1HWC0_C04,NC1HWC0,NC1HWC0_C04,NC1HWC0"
},
"input1": {
"name": "y",
"dtype": "float16,float16,int8,int8",
"format": "FRACTAL_Z_C04,FRACTAL_Z,FRACTAL_Z_C04,FRACTAL_Z"
},
"output0": {
"name": "z",
"dtype": "float16,float16,int32,int32",
"format": "NC1HWC0,NC1HWC0,NC1HWC0,NC1HWC0"
}
}

例如,conv2d算子的op_select_format函数实现如下:

import json
def op_select_format(inputs, weights, bias, offset_w, outputs, strides,
                     pads, dilations, groups=1, data_format='NHWC',
                     offset_x=0, kernel_name="conv2d"):
    shape_x = inputs.get("ori_shape")
    format_x = inputs.get("ori_format")
    shape_y = weights.get("ori_shape")
    format_y = weights.get("ori_format")
    x_dict = dict(zip(list(format_x), shape_x))
    y_dict = dict(zip(list(format_y), shape_y))

    use_c04 = False
    if x_dict["C"] <= 4 and (y_dict["W"] != 1 or y_dict["H"] != 1):
        use_c04 = True
    res = {}
    if use_c04:
        res["input0"] = {
                "name":"x",
                "dtype":"float16, float16, int8, int8",
                "format": "NC1HWC0, NC1HWC0_C04, NC1HWC0, NC1HWC0_C04"
            }
        res["input1"] = {
            "name":"filter",
            "dtype":"float16, float16, int8, int8",
            "format": "FRACTAL_Z, FRACTAL_Z_C04, FRACTAL_Z, FRACTAL_Z_C04"
        }
        res["input2"] = {
            "name":"bias",
            "dtype":"float16, float16, int32, int32",
            "format": "ND, ND, ND, ND"
        }
        res["input3"] = {
            "name":"offset_w",
            "dtype":"int8, int8, int8, int8",
            "format": "ND, ND, ND, ND"
        }
        res["output0"] = {
            "name":"filter",
            "dtype":"float16, float16, int8, int8",
            "format": "NC1HWC0, NC1HWC0, NC1HWC0, NC1HWC0"
        }
    else:
        res["input0"] = {
            "name":"x",
            "dtype":"float16, int8",
            "format": "NC1HWC0, NC1HWC0"
        }
        res["input1"] = {
            "name":"filter",
            "dtype":"float16, int8",
            "format": "FRACTAL_Z, FRACTAL_Z"
        }
        res["input2"] = {
            "name":"bias",
            "dtype":"float16, int32",
            "format": "ND, ND"
        }
        res["input3"] = {
            "name":"offset_w",
            "dtype":"int8, int8",
            "format": "ND, ND"
        }
        res["output0"] = {
            "name":"filter",
            "dtype":"float16, int8",
            "format": "NC1HWC0, NC1HWC0"
        }

    return json.dumps(res, indent=4)

算子精度优化

算子精度优化主要有以下两种方法,详细精度优化方法可参见精度优化专题

  • 将入参的数据类型转成float32进行计算。

    采用float32的数据类型进行计算,可以提高中间计算过程的精度,从而提升最终结果的精度,尤其当中间计算过程较为复杂时效果比较明显。

    进行数据类型的转换需要注意以下两点:
    • 在最后输出结果时要将数据类型转换成原数据类型。
    • float16转成float32计算会导致运行性能降低,因此如果使用float16的数据类型进行计算的精度在可允许范围内,尽量不要转换数据类型。
  • 通过变换公式,避免使用精度误差较大的接口。

    昇腾310 AI处理器场景下,当前版本TBE DSL的vsqrt, vlog, vexp接口精度误差相对较大,对于对精度要求较高的场景,可通过牛顿迭代、泰勒展开式的方式对计算公式进行变换。

    例如:vsqrt需要使用3次牛顿迭代替换,vexp需要使用6阶泰勒展开替换,vlog需要5阶泰勒展开替换。

    进行公式展开时有以下注意事项:
    • 使用展开的公式会增加编译、运行时间,降低了性能,所以在精度误差可接受的情况下尽量不要进行公式展开。
    • 进行公式展开请遵循如下原则:

      首先确定替换前的公式精度达标区间,对于达标范围内的计算不需要进行展开;对于不达标的区间,进行公式展开只能满足固定区间内的精度达标。所以首先要确定展开后的公式精度达标区间,然后推导公式,将不达标的区间的计算映射到达标区间内计算。

    • 建议将重复调用的展开公式封装成函数进行计算,有以下两点优势:
      • 减少单函数的中间变量个数。
      • 出错时便于修改。

算子性能优化

算子性能优化主要有以下三种方法:

  • 避免使用运行时间较长的接口。

    当前版本的vrec、vsel、vcmp接口运行耗时较长,对于对性能要求较高的场景,开发者可以通过变换计算公式的方法替换掉耗时较长的接口。例如,计算1/exp(x),可以替换为 exp(-x),可以先求-x,再求指数,从而避免了求倒数的操作。

    注意:进行指令替换时,需要同时考虑精度。

  • tvm.const接口不单独使用。

    省略单独定义tvm.const的步骤,尤其是对于只使用一次的值。在使用的时候直接定义。例如:

    # 替换前
    cosh_one = tvm.const(NUM_ONE, "float32")
    tensor_one = te.lang.cce.broadcast(cosh_one, data_y.shape)
    # 替换后
    tensor_one = te.lang.cce.broadcast(tvm.const(NUM_ONE, "float32"),data_y.shape)
  • 减少总的计算次数。

    通过变换公式,减少计算次数也可以降低编译时间,提升性能。

    注意:变换后的公式要正确、精度要在可接受范围。

    例如,在计算(1/vsqrt(x))*data_dy时,可以直接使用data_dy/vsqrt(x):

    # 修改前
    vsqrt_res = te.lang.cce.vsqrt(num_to_vrsqrt)res = te.lang.cce.vdiv(tvm.const(NUM_ONE, "float32"), vsqrt_res)res = te.lang.cce.vmul(res, data_dy)
    # 修改后
    vsqrt_res = te.lang.cce.vsqrt(num_to_vrsqrt)res = te.lang.cce.vdiv(data_dy, vsqrt_res)

vcmp/vsel接口使用注意事项

在使用te.lang.cce.vcmp与te.lang.cce.vsel接口时,会出现某些shape的情况下输出结果不符合预期,这是因为vcmp接口中mode的默认值为bool,表示按照8bit进行存储。而config配置中若不配置"bool_storage_as_1bit"参数,此参数默认值为True,表示按照1bit进行存储,与mode为bool不匹配,所以要在算子接口实现函数的编译配置config中加入配置项:"bool_storage_as_1bit": False,例如:

with tvm.target.cce():
    schedule = generic.auto_schedule(res)
config = {"name": kernel_name,
          "tensor_list": [data_x, data_y, res],
          "bool_storage_as_1bit": False}
te.lang.cce.cce_build_code(schedule, config)
分享:

    相关文档

    相关产品

关闭导读