该样例实现了从Global Memory中的A,B两处分别读取128个float16类型的数值,相加,并将结果写入Global Memory地址C中。
from tbe import tik
import tbe.common.platform as tbe_platform
def simple_add():
# 请根据实际昇腾AI处理器型号进行设置
soc_version="xxx"
tbe_platform.set_current_compile_soc_info(soc_version,core_type="AiCore")
tik_instance = tik.Tik()
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 //16, 0, 0)
tik_instance.data_move(data_B_ub, data_B, 0, 1, 128 //16, 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 //16, 0, 0)
tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
return tik_instance
下面进行代码解析:
from tbe import tik
“tbe.tik”:提供了所有TIK相关的python函数,具体请参考CANN软件安装后文件存储路径的“python/site-packages/tbe/tik”。
1
|
tik_instance = tik.Tik() |
通过Tensor在Global Memory中定义输入数据data_A和data_B,输出数据data_C。他们的Size分别为128个float16类型的数据。
通过Tensor在Unified Buffer中定义数据data_A_ub、data_B_ub、data_C_ub。他们的Size分别为128个float16类型的数据。
#在Global Memory中定义输入数据data_A和data_B,输出数据data_C。他们的Size分别为128个float16类型的数据。
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)
#在Unified Buffer中定义数据data_A_ub、data_B_ub、data_C_ub。他们的Size分别为128个float16类型的数据。
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 //16, 0, 0) tik_instance.data_move(data_B_ub, data_B, 0, 1, 128 //16, 0, 0)
在实现计算之前,我们先来了解下TIK指令的基本操作单位。
TIK的Vector指令每个cycle能处理256Byte的数据,并提供mask功能调整计算的数据,同时在时间上支持repeat操作,完成一连串的数据计算。
TIK指令操作分布在空间和时间两个维度,其中空间上最多处理256Byte数据(包括128个float16/uint16/int16、64个float32/uint32/int32或256个int8/uint8的数据),时间上支持repeat操作。1次repeat内部的数据,计算哪些数据,不计算哪些数据,由mask参数决定。针对float16数据,vetcor引擎一次计算128个elements,如mask=128,表示前128个elements参与计算。
结合add算子,我们通过vec_add实现计算过程。
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 //16, 0, 0)
通过BuildCCE将上述定义的TIK DSL容器编译生成昇腾AI处理器的可执行二进制文件。
tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
return tik_instance