功能实现
计算实现
计算实现包含依赖python模块导入、算子函数声明、算子入参校验、计算逻辑实现以及调度与编译。计算实现时,可以通过一定的方法进行精度与性能上的提升。
导入Python模块
进行TBE DSL算子开发时,首先需要在算子实现文件中导入昇腾AI软件栈提供的Python模块,代码示例如下所示,算子实现文件的命名请参见算子定义命名规则。
如果您进行代码实现时依赖了其他自行引入的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/lang/cce/”目录下的python函数,使用方法请参见TBE DSL API。
- “te.tvm”:引入TVM后端代码生成机制。
具体的接口定义可查看ATC安装路径下“/python/site-packages/te/tvm”目录下的python函数,使用方法请参见https://docs.tvm.ai/。
- “te.platform.fusion_manager.fusion_manager”:提供了实现算子的UB自动融合的接口。
具体的接口定义可查看ATC安装路径下“/python/site-packages/te/platform/fusion_manager.py”文件中的fusion_manager函数的定义。
- “topi.generic”:提供了自动算子调度接口。
具体的接口定义可查看ATC安装路径下“/python/site-packages/topi/generic”目录下的所有python函数,使用方法请参见TBE DSL API。
算子函数声明
算子的代码实现中包括两个函数:算子接口函数与算子compute函数,算子的compute函数会在算子接口函数中被调用。
下面详细介绍这两个函数的声明规则。
- 算子接口函数声明
如下所示,一个算子的接口函数中包含了算子输入信息、算子输出信息以及内核名称,函数的声明信息需要与算子信息定义文件中的信息对应。
def operationname(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="KernelName")
- 算子接口函数名称operationname当前版本请与算子实现文件名称保持一致,命名规则请参见算子定义命名规则。
- 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等信息,字典格式,此字段为预留位。
- attribute1,attribute2...:算子的属性,此处需要为算子的属性赋默认值,算子属性的名称、个数与顺序需要与算子信息定义文件/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,包含shape和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函数。
- 首先在算子接口函数operatorname( )中,获取算子输入tensor的shape以及dtype,并可自行实现基本的校验功能。
- 获取算子输入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()
- (可选)在算子实现函数中添加算子输入/输出及属性基本信息校验,有助于在算子编译阶段,提前发现问题。
例如,对于Add算子,首先校验两个输入的dtype是否一致,然后校验输入的数据类型是否在允许的数据类型列表中,代码实现如下所示。
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))
开发者可根据算子特点自定义实现校验函数。
- 获取算子输入tensor的shape以及dtype,为后续定义输入tensor的张量占位符做准备。
- 然后根据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对象被覆盖的情况。
- 在算子接口定义函数中调用compute函数进行计算过程的描述。
例如:
res = add_compute(data_x, data_y, output_z, kernel_name)
输入tensor为使用tvm.placeholder定义的占位tensor,其他为算子接口函数透传的参数。
- 算子compute函数的实现。
在compute函数中,完成算子的计算过程,计算过程的实现主要根据算子分析中的TBE DSL API进行代码开发。
入门示例中介绍了简单的add算子的实现过程,下面我们以计算公式较复杂的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进行计算,可以提高中间计算结果的精度,最后的结果输出时需要将数据类型转换成原数据类型。
- 当算子的计算过程比较繁琐时,可以通过抽调内部函数的方法保持每个模块的简洁性和可读性。
调度与编译
如下代码所示,当定义完计算逻辑后,需要在算子接口实现函数中实现调度与编译。
通过调用auto_schedule接口,便可以自动生成相应的调度,此处通过TVM的打印机制可以看到相应计算的中间表示。配置信息包括是否需要打印IR、是否编译以及算子内核名以及输入、输出张量。
with tvm.target.cce(): schedule = generic.auto_schedule(result) config = { "print_ir": True, "need_build": True, "name": kernel_name, "tensor_list": [input_data, result] "bool_storage_as_bit":True } te.lang.cce.cce_build_code(schedule, config)
- 使用“generic”的“auto_schedule”接口,自动生成相应的调度(schedule),“auto_schedule”接口的参数为算子的输出张量。
schedule可以理解为:描述的计算过程如何在硬件上高效执行。就是把相关的计算和硬件设备上的相关指令对应起来。schedule对象中包含一个“中间表示”(IR),它用一种类似伪代码来描述计算过程,可以通过“print_ir”参数把它打印出来进行查看。
- “need_build”:是否进行build,默认是True。
- “name”:编译生成的算子二进制文件名称,只能是大小写字母、数字、“_”的组合,且必须是字母或者“_”开头,长度小于或等于200个字符。
- “tensor_list”(张量列表)中保存输入张量、输出张量,这个顺序需要严格按照算子本身的输入、输出数据顺序排列。
注意:输入tensor需要是placeholder接口返回的tensor对象,此tensor对象的内存地址不能被覆盖。
例如:"tensor_list": [tensor_a, tensor_b, res],tensor_a与tensor_b是输入张量,res为输出张量。
- “bool_storage_as_1bit”:Bool类型存储时是否按照1 bit存储。True:按照1bit存储,False:按照8bit进行存储,默认值为True。
当te.lang.cce.vcmp(lhs, rhs, op, mode)接口的mode为bool时,需要设置此参数为False。
- 根据调度和配置使用“te.lang.cce”提供的“cce_build_code”接口来进行算子编译,算子编译过程会根据输入的数据形状、类别、算子参数等编译出专用内核,这个过程在离线模型生成器转换模型时发生。
- schedule:生成的算子计算schedule对象。
- config:编译参数配置的map。
编译完成后,会生成算子目标文件.o文件与算子描述文件.json文件。