更新时间:2021-03-18 GMT+08:00
分享

TIK介绍

什么是TIK

TIK(Tensor Iterator Kernel)是一种基于Python语言的动态编程框架,呈现为一个Python模块,运行于Host CPU上。

开发者可以通过调用TIK提供的API基于Python语言编写自定义算子,然后TIK编译器会编译为昇腾AI处理器应用程序的二进制文件。

TIK的优势

TIK算子开发方式是一种灵活的开发方式。TIK代码在算子开发效率和算子性能自动优化上有着一定的优势:

  1. 自动内存分配和自动数据依赖规划,让用户用串行编程的思路既可以写出高性能的并行计算的算子。
  2. 通过手动调度可以更加精确的控制数据搬运和计算流程,从而实现更高的性能,将昇腾AI处理器的能力发挥到极致。

TIK算子开发流程

基于TIK API编写Python程序的通用步骤,如图1所示。

图1 算子实现流程

主要步骤为:

  1. Python模块导入。

    from te import tik↵

    “te.tik”:提供了所有TIK相关的python函数,具体请参考ATC安装目录下的“python/site-packages/te.egg/te/tik”。

  2. 构建TIK DSL容器。

    1
    2
    from te import tik
    tik_instance = tik.Tik()
    

  3. 在AI Core的外部存储和内部存储中定义输入数据、输出数据。

    1
    2
    3
    4
    5
    6
    7
    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)
    

  4. 将外部存储中的数据搬入AI Core内部存储(比如Unified Buffer)中。

    1
    2
    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. 进行计算。

    1
    2
    3
    repeat = tik_instance.Scalar('int32')
    repeat.set_as(1)
    tik_instance.vec_abs(128, data_C_ub[0], data_A_ub[0], data_B_ub[0], repeat, 8, 8, 8)
    

  6. 搬出到外部存储。

    1
    tik_instance.data_move(data_C, data_C_ub, 0, 1, 128 //16, 0, 0)
    

  7. 将TIK DSL容器中的语句,编译成昇腾AI处理器可执行的代码。

    1
    2
    #其中kernel_name决定了算子昇腾AI处理器上可执行的二进制文件的名称,inputs为从外部存储中加载的数据,outputs对应计算后搬运到外部存储中的数据。
    tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
    

分享:

    相关文档

    相关产品

close