入门示例
使用TBE DSL进行算子开发,开发者只需要关注算子本身计算逻辑,无需关注调度策略,开发方式简单便捷。
目标
本节以add算子为例带您快速熟悉一个TBE DSL算子的编写流程。
Add算子实现了两个数据相加,返回相加结果的功能,如下所示:
算子分析
使用TBE DSL方式开发Add算子前,我们需要确定算子功能、输入、输出,算子开发方式、算子类型以及算子实现函数名称等。
- 明确算子的功能以及数学表达式。
add算子的数学表达式为:
z=x+y
计算过程是:将两个输入参数相加,得到最终结果z并将其返回。
- 明确输入和输出。
- add算子有两个输入:x与y,输出为z。
- 本样例中算子的输入支持的数据类型为float16、float32、 int32,算子输出的数据类型与输入数据类型相同。
- 算子输入支持所有shape,输出shape与输入shape相同。
- 算子输入支持的format为:NCHW,NC1HWC0,NHWC,ND。
- 确定算子开发方式及使用的计算接口。
- 计算过程只涉及加法操作,查看TBE DSL API,初分析可使用te.lang.cce.vadd(lhs, rhs)接口实现x+y。
- 由于te.lang.cce.vadd(lhs, rhs)接口要求两个输入tensor的shape需要相同,所以需要首先获取两个输入tensor中较大的shape值,然后调用te.lang.cce.broadcast(var, shape, output_dtype=None)接口将入参广播成指定的shape大小。
- 明确算子实现文件名称、算子实现函数名称以及算子的类型(OpType)。
- 算子类型需要采用大驼峰的命名方式,即采用大写字符区分不同的语义。
- 算子文件名称和算子函数名称,可选用以下任意一种命名规则:
- 用户自定义,此时需要在算子信息定义中配置opFile.value与opInterface.value。
- 不配置算子信息定义中的opFile.value与opInterface.value,FE会将OpType按照如下方式进行转换后进行算子文件名和算子函数名的匹配。转换规则如下:
- 首字符的大写字符转换为小写字符。
例如:Abc -> abc
- 小写字符后的大写字符转换为下划线+小写字符。
例如:AbcDef -> abc_def
- 紧跟数字以及大写字符后的大写字符,作为同一语义字符串,查找此字符串后的第一个小写字符,并将此小写字符的前一个大写字符转换为下划线+小写字符,其余大写字符转换为小写字符。若此字符串后不存在小写字符,则直接将此字符串中的大写字符转换为小写字符。
例如:ABCDef -> abc_def;Abc2DEf -> abc2d_ef;Abc2DEF -> abc2def;ABC2dEF -> abc2d_ef。
- 首字符的大写字符转换为小写字符。
因此本例中,算子类型定义为Add;算子的实现文件名称及实现函数名称定义为add。
通过以上分析,得到Add算子的设计规格如下:
表6-1 Add算子设计规格算子类型(OpType)
Add
算子输入
name:x
shape:all
data type:
float16、float32、 int32
format:
NCHW,NC1HWC0,
NHWC,ND
name:y
shape:all
data type:
float16、float32、 int32
format:
NCHW,NC1HWC0,
NHWC,ND
算子输出
name:z
shape:all
data type:
float16、float32、 int32
format:
NCHW,NC1HWC0,
NHWC,ND
算子实现使用主要DSL接口
te.lang.cce.broadcast(var, shape, output_dtype=None)
te.lang.cce.vadd(lhs, rhs)
算子实现文件/实现函数名称
add
算子代码实现
Add算子仅支持float16, float32, int32三种数据类型,所以需要对算子的输入数据进行校验;由于Add算子允许两个输入数据的shape不同,但算子计算接口te.lang.cce.vadd( )要求两输入shape相同,因此需要对算子两个输入的shape进行广播并对其进行校验,算子实现代码示例如下所示:
SHAPE_SIZE_LIMIT = 2147483648 # 比较两个输入的每个维度的大小,取每个维度的大值,生成out_shape def _produce_shapes(shape1, shape2): shape1 = list(shape1) shape2 = list(shape2) flag = 0 if len(shape1) < len(shape2): shape1, shape2 = shape2, shape1 flag = 1 output_shape_len = len(shape1) dec = output_shape_len - len(shape2) for i in range(dec): shape2 = [1] + shape2 out_shape = [] for i in range(output_shape_len): if (shape1[i] != shape2[i]) and (shape1[i] != 1) and (shape2[i] != 1): raise RuntimeError("input shapes not match!") out_shape.append(shape1[i] if shape1[i] > shape2[i] else shape2[i]) if flag == 1: shape1, shape2 = shape2, shape1 return shape1, shape2, out_shape # 将shape转换为list def _shape_to_list(shape): result = [] for i in shape: if isinstance(i, tvm.expr.Var): result.append(i) else: result.append(i.value) return result # 实现Add算子的计算逻辑 @fusion_manager.register("add") def add_compute(input_x, input_y, output_z, kernel_name="add"): shape_x = _shape_to_list(input_x.shape) shape_y = _shape_to_list(input_y.shape) shape_x, shape_y, shape_max = _produce_shapes(shape_x, shape_y) # shape_max取shape_x与shape_y的每个维度的大值 shape_size = reduce(lambda x, y: x * y, shape_max[:]) if shape_size > SHAPE_SIZE_LIMIT: raise RuntimeError("the shape is too large to calculate") input_x = te.lang.cce.broadcast(input_x, shape_max) # 将input_x的shape广播为shape_max input_y = te.lang.cce.broadcast(input_y, shape_max) # 将input_y的shape广播为shape_max res = te.lang.cce.vadd(input_x, input_y) # 执行input_x + input_y return res # 返回计算结果的tensor # 算子定义函数 def add(input_x, input_y, output_z, kernel_name="add"): # 获取算子输入tensor的shape与dtype shape_x = input_x.get("shape") shape_y = input_y.get("shape") check_tuple = ("float16", "float32", "int32") input_data_type = input_x.get("dtype").lower() if input_data_type not in check_tuple: raise RuntimeError("only support %s while dtype is %s" % (",".join(check_tuple), input_data_type)) # shape_max取shape_x与shape_y的每个维度的最大值 shape_x, shape_y, shape_max = _produce_shapes(shape_x, shape_y) if shape_x[-1] == 1 and shape_y[-1] == 1 and shape_max[-1] == 1: # 如果shape的长度等于1,就直接赋值,如果shape的长度不等于1,做切片,将最后一个维度舍弃(按照内存排布,最后一个维度为1与没有最后一个维度的数据排布相同,例如2*3=2*3*1,将最后一个为1的维度舍弃,可提升后续的调度效率)。 shape_x = shape_x if len(shape_x) == 1 else shape_x[:-1] shape_y = shape_y if len(shape_y) == 1 else shape_y[:-1] shape_max = shape_max if len(shape_max) == 1 else shape_max[:-1] # 使用TVM的placeholder接口对第一个输入tensor进行占位,返回一个tensor对象 data_x = tvm.placeholder(shape_x, name="data_1", dtype=input_data_type) # 使用TVM的placeholder接口对第二个输入tensor进行占位,返回一个tensor对象 data_y = tvm.placeholder(shape_y, name="data_2", dtype=input_data_type) # 调用compute实现函数 res = add_compute(data_x, data_y, output_z, kernel_name) # 自动调度 with tvm.target.cce(): schedule = generic.auto_schedule(res) # 编译配置 config = {"name": kernel_name, "tensor_list": (data_x, data_y, res)} te.lang.cce.cce_build_code(schedule, config)