文档HOT
下载
中文
注册

入门样例

该样例实现了从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

下面进行代码解析：

  1. 导入Python模块。
    from tbe import tik

    “tbe.tik”：提供了所有TIK相关的python函数，具体请参考CANN软件安装后文件存储路径的“python/site-packages/tbe/tik”。

  2. 构造TIK DSL容器。
    通过TIK类构造函数构造TIK DSL容器。 
    1
    		 
    tik_instance = tik.Tik()
  3. 定义数据。

    通过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)
  4. 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)
  5. 将加载到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参与计算。

    结合add算子，我们通过vec_add实现计算过程。

    • 【接口定义】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)
  6. 将data_C_ub中的计算结果搬运到data_C中。同样使用data_move实现。
    tik_instance.data_move(data_C, data_C_ub, 0, 1, 128 //16, 0, 0)
  7. 将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])
  8. 返回tik实例。
    return tik_instance
父主题： TIK样例