本节以一个简单的Add算子为例,带您快速熟悉TIK算子的编写流程。
本样例中,定义Add算子的两个输入的shape都为(128,),数据类型为float16,运行目标为昇腾AI处理器的AI Core。
算子分析阶段,需要明确算子的数学表达式,输入/输出,调度流程设计,算子的实现函数名称等。
y=x1+x2
计算过程是:将两个输入参数相加,得到最终结果y并将其返回。
因此本例中,算子类型定义为AddTik;算子的实现文件名称及实现函数名称定义为add_tik。
通过以上分析,得到Add算子的设计规格如下:
算子类型(OpType) |
AddTik |
|||
---|---|---|---|---|
算子输入 |
name:x1 |
shape: (128,) |
data type: float16 |
format: NCHW,NC1HWC0, NHWC,ND |
name:x2 |
shape: (128,) |
data type: float16 |
format: NCHW,NC1HWC0, NHWC,ND |
|
算子输出 |
name:y |
shape: (128,) |
data type: float16 |
format: NCHW,NC1HWC0, NHWC,ND |
TIK计算接口 |
vec_add |
|||
实现函数名称 |
add_tik |
下面以此AddTik算子为例描述TIK程序的编写步骤:
from tbe import tik import tbe.common.platform as tbe_platform import numpy as np
“tbe.tik”:提供了所有TIK相关的python函数,具体请参考CANN软件安装后文件存储路径的“python/site-packages/tbe/tik”。
def add_tik():
说明:本样例为了让大家对TIK有初步的认识,数据类型与形状都设置成了固定值,实际进行算子开发时,算子的数据是执行时才传入的,所以需要在算子实现函数定义时包含输入与输出数据基本信息,如TIK入口函数。
# soc_version请设置为实际昇腾AI处理器的型号
tbe_platform.set_current_compile_soc_info(soc_version)
“set_current_compile_soc_info”接口中的参数“core_type”用来指定核类型,默认值为“AiCore”,代表运行目标为AI Core。
tik_instance = tik.Tik(disable_debug=False)
因为此样例后续要进行功能调试,所以设置“disable_debug=False”使能调试开关。
data_A = tik_instance.Tensor("float16", (128,), name="data_A", scope=tik.scope_gm) data_B = tik_instance.Tensor("float16", (128,), name="data_B", scope=tik.scope_gm) data_C = tik_instance.Tensor("float16", (128,), name="data_C", scope=tik.scope_gm) data_A_ub = tik_instance.Tensor("float16", (128,), name="data_A_ub", scope=tik.scope_ubuf) data_B_ub = tik_instance.Tensor("float16", (128,), name="data_B_ub", scope=tik.scope_ubuf) data_C_ub = tik_instance.Tensor("float16", (128,), name="data_C_ub", scope=tik.scope_ubuf)
tik_instance.data_move(data_A_ub, data_A, 0, 1, 128*2 //32, 0, 0) tik_instance.data_move(data_B_ub, data_B, 0, 1, 128*2 //32, 0, 0)
tik_instance.vec_add(128, data_C_ub[0], data_A_ub[0], data_B_ub[0], 1, 8, 8, 8)
tik_instance.data_move(data_C, data_C_ub, 0, 1, 128*2 //32, 0, 0)
tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
其中,
return tik_instance
if __name__ == "__main__": # 调用TIK算子实现函数 tik_instance = add_tik() # 初始化数据,为128个float16类型的数字1的一维矩阵 data = np.ones((128,), dtype=np.float16) feed_dict = {"data_A": data, "data_B": data} # 启动功能调试 data_C, = tik_instance.tikdb.start_debug(feed_dict=feed_dict, interactive=True) # 打印输出数据 print(data_C)
python3 add_tik.py
输入data_A和data_B分别为128个float16类型的数字1的一维矩阵,展示的输出data_C如下所示:
[TIK]>c [2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]
以上仅对大致编程过程进行展示,后续会对上述实例中各接口、各参数的含义进行详细描述。