Add算子(TBE DSL方式)
功能描述
本样例使用TBE DSL方式进行了Add算子的实现,并对Add算子进行了单算子功能验证及单算子网络验证。
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算子的设计规格如下:
表14-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)
算子适配插件实现
将原始Tensorflow的Add算子解析并映射为适配昇腾AI处理器的Add算子,算子属性的映射可直接调用AutoMappingFn( )接口进行实现,完整代码可参考sample样例中的“framework/tf_plugin/add_plugin.cpp”文件。
算子原型定义
原型定义的关键点是推理输出Tensor的shape以及对算子输入的内在关联关系进行校验。
Add算子推理输出shape的原理为:首先获取两个输入的shape,然后将两个输入shape广播为相同的shape,输出shape取两个输入中每个维度的大值。InferShapeAndTypeAdd函数的代码如下所示:
“op_proto/add.cpp”
bool InferShapeAndTypeAdd(Operator& op, const string& input_name1, const string& input_name2, const string& output_name) { // vOutputDesc.push_back(op.GetInputDesc(0)); TensorDesc vOutputDesc = op.GetOutputDesc(output_name); DataType input_dtype = op.GetInputDesc(input_name1).GetDataType(); Format input_format = op.GetInputDesc(input_name1).GetFormat(); // 针对shape维度大小进行交换 ge::Shape shapeX = op.GetInputDesc(input_name1).GetShape(); ge::Shape shapeY = op.GetInputDesc(input_name2).GetShape(); std::vector<int64_t> dimsX = shapeX.GetDims(); std::vector<int64_t> dimsY = shapeY.GetDims(); if (dimsX.size() < dimsY.size()) { std::vector<int64_t> dimsTmp = dimsX; dimsX = dimsY; dimsY = dimsTmp; } // 对小的shape进行1补齐 if (dimsX.size() != dimsY.size()) { int dec = dimsX.size() - dimsY.size(); for (int i = 0; i < dec; i++) { dimsY.insert(dimsY.begin(), (int64_t)1); } } // 设置输出的shape维度,取两个输入shape中每个维度的大值 std::vector<int64_t> dimVec; for (size_t i = 0; i < dimsX.size(); i++) { if ((dimsX[i] != dimsY[i]) && (dimsX[i] != 1) && (dimsY[i] != 1)) { return false; } int64_t dims = dimsX[i] > dimsY[i] ? dimsX[i] : dimsY[i]; dimVec.push_back(dims); } ge::Shape outputShape = ge::Shape(dimVec); vOutputDesc.SetShape(outputShape); vOutputDesc.SetDataType(input_dtype); vOutputDesc.SetFormat(input_format); op.UpdateOutputDesc(output_name, vOutputDesc); return true; }
算子信息定义
Add算子的信息定义文件请参见“tbe/op_info_cfg/ai_core/<soc_version>/add.ini”。