入门样例
该样例实现了从Global Memory中的A,B两处分别读取128个float16类型的数值,相加,并将结果写入Global Memory地址C中。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
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 |
下面进行代码解析:
- 导入Python模块。
from tbe import tik
“tbe.tik”:提供了所有TIK相关的python函数,具体请参考CANN软件安装后文件存储路径的“python/site-packages/tbe/tik”。
- 构造TIK DSL容器。
- 定义数据。
通过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类型的数据。
- 【接口定义】Tensor(dtype, shape, scope, name)
- 【参数分析】
- dtype:指定Tensor对象的数据类型。
- shape:指定Tensor对象的形状。
- scope:指定Tensor对象的所在buffer空间。scope_gm表示Global Memory中的数据;scope_ubuf表示Unified Buffer中的数据。
- name:指定Tensor名字,不同Tensor名字需要保持唯一。
- 【示例】
#在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)
- Global Memory中的数据搬运到Unified Buffer中。通过data_move实现数据搬运,即将data_A中的数据搬运到data_A_ub,data_B中的数据搬运到data_B_ub。
- 【接口定义】data_move (dst, src, sid, nburst, burst, src_stride, dst_stride, *args, **argv)
- 【参数分析】
- src, dst:源地址与目标地址。
- sid:sim ID,固定为0。
- burst, nburst:burst为每次搬运的数据大小(单位为32Byte),nburst为搬运次数。我们需要搬运的数据为128个float16类型的数据,占128*2Byte,小于Unified Buffer的大小(256KB),因此我们搬一次就可以把输入数据全部搬到Unified Buffer,搬运的次数nburst为1;由于burst单位为32Byte,每次搬运的数据大小burst为128*2/32Byte。
- src_stride, dst_stride:源/目的地址的stride,当需要带间隔跳跃式搬运时,需要设置这两个参数,示例中为连续搬运,因此两个参数都设置为0。
- 【示例】
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)
- 将加载到data_A_ub和data_B_ub的数据进行vec_add计算,并将计算结果写回到data_C_ub。
在实现计算之前,我们先来了解下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参与计算。
- 【接口定义】vec_add(mask, dst, src0, src1, repeat_times, dst_rep_stride, src0_rep_stride, src1_rep_stride)
- 【参数分析】
- src0,src1,dst:源操作数0,源操作数1和目的操作数,分别为data_A_ub、data_B_ub、data_C_ub。
- repeat_times:迭代次数。通过上面对TIK指令的基本了解,我们可以算出,对于128个float16的数据,通过1次repeat可以完成计算,因此repeat_times为1。
- dst_rep_stride, src0_rep_stride, src1_rep_stride:相邻迭代间,目的操作数/源操作数0/源操作数1头头间地址间隔,单位为32Byte,此处配置为8,表示一个迭代连续处理8*32Byte数据。
- mask:数据操作有效指示,128表示计算所有元素。
- 【示例】
tik_instance.vec_add(128, data_C_ub[0], data_A_ub[0], data_B_ub[0], 1, 8, 8, 8)
- 将data_C_ub中的计算结果搬运到data_C中。同样使用data_move实现。
tik_instance.data_move(data_C, data_C_ub, 0, 1, 128 //16, 0, 0)
- 将TIK DSL容器中的语句,编译成昇腾AI处理器可执行的代码。
通过BuildCCE将上述定义的TIK DSL容器编译生成昇腾AI处理器的可执行二进制文件。
- 【接口定义】BuildCCE(kernel_name, inputs, outputs, output_files_path=None, enable_l2=False)
- 【参数分析】
- kernel_name:指明编译产生的二进制代码中的AI Core核函数名称。
- inputs:存放程序的输入Tensor,必须是Global Memory的存储类型。
- outputs:存放程序的输出Tensor,必须是Global Memory的存储类型。
- output_files_path:指明编译产生文件的存储位置,默认为”./kernel_meta”。
- enable_l2:该参数暂不生效。
- 【示例】
tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
- 返回tik实例。
return tik_instance
父主题: TIK样例