TBE Add 算子开发入门
1 TBE概述
TBE（Tensor Boost Engine，即张量加速引擎），是一款华为自研的算子开发工具，用于开发能够运行在NPU芯片上的算子，该工具是在业界著名的开源项目TVM（Tensor Virtual Machine）基础上扩展的，提供了一套Python API来进行开发。借助TBE算子开发框架，开发者可以基于Python语言开发自定义算子，从而降低了入门的难度，对于算子的业务逻辑比较简单的场景，适合用TBE进行算子开发。一般来说，TBE进行算子开发的方式有两种：DSL开发和TIK开发。DSL开发方式比较简单，通过算子定义文件，借助相关工具自动生成Python代码，并在编译时，自动生成特定目标的代码。
2 Add算子分析
Add算子实现两个张量x1,x2的和，数学表达式是 y = x1 + x2 。Add算子有两个输入与一个输出，支持的数据类型都为float16，数据排布格式为NCHW，要求两个输入的shape相同。本文以DSL方式实现算子，通过查看TBE DSL API，得知可以使用 tbe.dsl.vadd方法来实现业务逻辑 。这里将Add算子的OpType命名为AddDSL，算子实现文件及实现函数名称命名为add_dsl 。
3 Add算子实现
定义AddDSL算子的原型定义json文件add_dsl.json，用于生成AddDSL的算子开发工程。add_dsl.json文件的内容如下所示：
[
{
"op":"AddDSL",
"input_desc":[
{
"name":"x1",
"param_type":"required",
"format":[
"NCHW"
],
"type":[
"fp16"
]
},
{
"name":"x2",
"param_type":"required",
"format":[
"NCHW"
],
"type":[
"fp16"
]
}
],
"output_desc":[
{
"name":"y",
"param_type":"required",
"format":[
"NCHW"
],
"type":[
"fp16"
]
}
]
}
]
root@atlas500ai:/home/kzroot/mysoft/myAscendC/opdsl# ${INSTALL_DIR}/python/site-packages/bin/msopgen gen
-i ./add_dsl.json
-f tf
-c ai_core-Ascend310P3
-out ./AddDsl其中的 ${INSTALL_DIR} 为 /usr/local/Ascend/ascend-toolkit/latest 。执行成功显示如下所示：
AddDsl目录下生成算子工程，工程中包含各交付件的模板文件，编译脚本等，如下所示：
AddDsl
├── build.sh // 编译入口脚本
├── cmake // 编译解析脚本存放目录
├── CMakeLists.txt
├── framework // 算子适配插件相关文件存放目录
│ ├── CMakeLists.txt
│ └── tf_plugin
│ ├── CMakeLists.txt
│ └── tensorflow_add_dsl_plugin.cc // 算子适配插件实现文件
├── op_proto // 算子原型定义相关文件存放目录
│ ├── add_dsl.cc
│ ├── add_dsl.h
│ └── CMakeLists.txt
├── op_tiling
│ └── CMakeLists.txt
├── scripts // 自定义算子工程打包脚本存放目录
└── tbe
├── CMakeLists.txt
├── impl // 算子代码实现
│ └── add_dsl.py
└── op_info_cfg // 算子信息库存放目录
└── ai_core
└── ascend310p
└── add_dsl.ini修改 op_proto/add_dsl.cc文件，实现算子的输出描述推导函数及校验函数。具体如下所示：
#include "add_dsl.h"
namespace ge {
IMPLEMT_COMMON_INFERFUNC(AddDSLInferShape)
{
// 获取输出数据描述
TensorDesc tensordesc_output = op.GetOutputDescByName("y");
tensordesc_output.SetShape(op.GetInputDescByName("x1").GetShape());
tensordesc_output.SetDataType(op.GetInputDescByName("x1").GetDataType());
tensordesc_output.SetFormat(op.GetInputDescByName("x1").GetFormat());
// 直接将输入x1的Tensor描述信息赋给输出
(void)op.UpdateOutputDesc("y", tensordesc_output);
return GRAPH_SUCCESS;
}
IMPLEMT_VERIFIER(AddDSL, AddDSLVerify)
{
// 校验算子的两个输入的数据类型是否一致，若不一致，则返回失败。
if (op.GetInputDescByName("x1").GetDataType() != op.GetInputDescByName("x2").GetDataType()) {
return GRAPH_FAILED;
}
return GRAPH_SUCCESS;
}
COMMON_INFER_FUNC_REG(AddDSL, AddDSLInferShape);
VERIFY_FUNC_REG(AddDSL, AddDSLVerify);
} // namespace ge
实现AddDSL算子的计算逻辑。修改 tbe/impl/add_dsl.py 文件，其中已经自动生成了算子代码的框架，开发者需要在此文件中修改add_dsl_compute函数，实现此算子的计算逻辑。代码如下所示：
import tbe.dsl as tbe
from tbe import tvm
from tbe.common.register import register_op_compute
from tbe.common.utils import para_check
@register_op_compute("add_dsl")
def add_dsl_compute(x1, x2, y, kernel_name="add_dsl"):
"""
To do: Implement the operator by referring to the
TBE Operator Development Guide.
"""
# res = tbe.XXX(x1, x2)
# 调用dsl的vadd计算接口
res = tbe.vadd(x1, x2)
return res
@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
def add_dsl(x1, x2, y, kernel_name="add_dsl"):
"""
To do: Implement the operator by referring to the
TBE Operator Development Guide.
"""
data_x1 = tvm.placeholder(x1.get("shape"), dtype=x1.get("dtype"), name="data_x1")
data_x2 = tvm.placeholder(x2.get("shape"), dtype=x2.get("dtype"), name="data_x2")
res = add_dsl_compute(data_x1, data_x2, y, kernel_name)
# auto schedule
with tvm.target.cce():
schedule = tbe.auto_schedule(res)
# operator build
config = {"name": kernel_name,
"tensor_list": [data_x1, data_x2, res]}
tbe.build(schedule, config)
至此，AddDSL算子的所有交付件都已开发完毕。下面就可以进行编译工程了，修改build.sh脚本，配置算子编译所需环境变量ASCEND_TENSOR_COMPILER_INCLUDE ，将其设置为CANN软件头文件所在路径。执行如下命令可设置：
export ASCEND_TENSOR_COMPILER_INCLUDE=/usr/local/Ascend/ascend-toolkit/latest/include在算子工程目录下执行如下命令，进行算子工程编译。编译成功后，会在当前目录下创建build_out目录，并在build_out目录下生成自定义算子安装包 custom_opp_ubuntu_aarch64.run 。然后可以执行安装。
root@atlas500ai:/home/kzroot/mysoft/myAscendC/opdsl/AddDsl# ./build.sh
root@atlas500ai:/home/kzroot/mysoft/myAscendC/opdsl/AddDsl# ./build_out/custom_opp_ubuntu_aarch64.run
命令执行成功后，自定义算子包中的相关文件部署到CANN算子库中。下面编写测试用例。创建一个st目录，并创建一个测试用例的定义文件add_dsl.json ，具体内容如下所示：
[
{
"case_name":"Test_AddDSL_001",
"op": "AddDSL",
"error_threshold":[0.0001,0.0001],
"calc_expect_func_file":"./check_add.py:test_add",
"input_desc": [
{
"name": "x1",
"format": [
"NCHW"
],
"type": [
"float16"
],
"shape": [8,16],
"data_distribute": [
"uniform"
],
"value_range": [
[
-10.0,
10.0
]
]
},
{
"name": "x2",
"format": [
"NCHW"
],
"type": [
"float16"
],
"shape": [8,16],
"data_distribute": [
"uniform"
],
"value_range": [
[
-10.0,
10.0
]
]
}
],
"output_desc": [
{
"name": "y",
"format": [
"NCHW"
],
"type": [
"float16"
],
"shape": [8,16]
}
]
}
]其中的"calc_expect_func_file":"./check_add.py:test_add" 定义测试的校验函数，check_add.py文件内容如下所示：
def test_add(x1,x2,y):
res = x1['value'] + x2['value']
return [res,]配置ST测试用例编译所需环境变量。ST测试用例的主要功能为：使用AscendCL接口加载单算子模型文件并执行，所以生成并编译ST测试用例前需要配置AscendCL应用编译所需环境变量：
root@atlas500ai:/home/kzroot/mysoft/myAscendC/opdsl/st# export DDK_PATH=${INSTALL_DIR}
root@atlas500ai:/home/kzroot/mysoft/myAscendC/opdsl/st# export NPU_HOST_LIB=${INSTALL_DIR}/runtime/lib64/stub
root@atlas500ai:/home/kzroot/mysoft/myAscendC/opdsl/st# ${INSTALL_DIR}/python/site-packages/bin/msopst run
-i ./add_dsl.json
-soc Ascend310P3
-out ./out执行测试通过后，显示如下所示：
- 点赞
- 收藏
- 关注作者
评论（0）