下载
EN
注册

TIK算子实现流程

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

图1 算子实现流程(AICore)

如果算子过大(比如算子某个作用域内变量嵌套层数过大的场景),占用空间超限,编译时可能会出现因系统资源不足导致的报错,可尝试通过命令 “ulimit -s”临时增大堆栈空间来解决,例如堆栈空间默认为8MB的情况下,通过ulimit -s 16384将其增大至16MB。

代码结构如下所示:

# 导入依赖的Python模块
from tbe import tik
import tbe.common.platform as tbe_platform
from tbe.common.utils import para_check
import numpy as np

# 算子定义入口函数,即为算子运行时调用的函数,定义可参见TIK入口函数
@para_check.check_input_type(dict, dict, dict, str)
def operation_name(input_x, input_y, output_z, kernel_name):

    # 定义目标机,即设置昇腾AI处理器的版本,默认运行目标为AI Core
    # soc_version请配置为用户当前实际使用的昇腾AI处理器的型号
    tbe_platform.set_current_compile_soc_info(soc_version)
    # 构建TIK DSL容器
    tik_instance = tik.Tik(disable_debug=False)

    # 进行数据定义
    data_x = tik_instance.Tensor("dtype", shape, name="data_x", scope=tik.scope_gm)
    ... ...
    data_x_ub = tik_instance.Tensor("dtype", shape, name="data_x_ub", scope=tik.scope_ubuf)
    # 将数据搬入指定的Core
    tik_instance.data_move(data_x_ub, data_x, ...)

    # 进行数据计算
    tik_instance.vec_xxx(...)
    # 将数据搬出指定的Core
    tik_instance.data_move(data_x, data_x_ub, ...)
    # 将算子编译成目标文件,TIK编译函数定义可参见TIK编译函数
    tik_instance.BuildCCE(kernel_name="simple_add",inputs=[data_A,data_B],outputs=[data_C])
    # 返回TIK示例
    return tik_instance