TIK函数

经过前几章的学习，相信读者已经掌握了编写TIK算子的数据搬运和矢量计算的技能，以及了解了TIK作为代码生成器的本质，也已经可以编写一些最基本的TIK算子。

在本章中，我们将主要介绍TIK中函数的概念和使用技巧。

TIK算子主要由TIK函数以及其他普通的Python函数组成，且必须包含一个TIK入口函数。TIK函数借用Python函数进行封装，它的形式类似于：函数名（参数列表）。其中参数列表包含一组以逗号分隔的参数定义，参数使用类似Python的方式定义，可以顺序指定，也可以通过变量名使用“name=value”的方式指定。因为TIK代码生成器的本质，所以我们在写TIK函数时可以使用一些技巧，前提是不破坏Python的语法，不然解释器语法检查通不过。

TIK入口函数

每个TIK算子都要有一个TIK入口函数，作为算子运行时的调用函数，TIK入口函数的声明如下所示：

def operationname(input_x1, input_x2, output_y, attribute1=None, attribute2=None,..., kernel_name="KernelName")

入口函数的函数名请与算子实现的python文件名保持一致，命名规则请参见算子定义命名规则。
input_x1, input_x2：算子的输入tensor，每个tensor需要采用字典的形式进行定义，包含shape、ori_shape、format、ori_format与dtype信息，例如：
dict input_x1 = {'shape' : (2,2), 'ori_shape' : (2,2), 'format': 'ND', 'ori_format':'ND', 'dtype' : 'float16'}

输入tensor的顺序及个数需要与算子原型定义保持一致，可选输入也需要在此处定义，在计算逻辑中去判断是否有数据传入，并进行相应处理。
output_y：算子的输出tensor，包含shape和dtype等信息，字典格式，此字段为预留位。
输出tensor的顺序及个数也需要与算子原型定义保持一致，可选输出也需要在此处定义。
attribute1，attribute2...：算子的属性，算子属性的顺序与个数需要与算子原型定义保持一致。

若算子无相关属性信息，此参数忽略；若算子的属性为可选值，此处需要为算子的属性赋默认值。
kernel_name：算子在内核中的名称（即生成的二进制文件与算子描述文件的名称），用户自定义，保持唯一，只能是大小写字母、数字、“_”的组合，且必须是字母或者“_”开头，长度小于或等于200个字符。

TIK入口函数中可以直接调用其他TIK函数，最终生成TIK算子对应的CCE算子。

TIK类构造函数

TIK类构造函数用于创建TIK DSL容器，并返回TIK实例，如下所示：

from tbe import tik
from tbe.common.platform import set_current_compile_soc_info
# 请根据实际昇腾AI处理器型号进行设置
soc_version="xxx"
set_current_compile_soc_info(soc_version,core_type="AiCore")
tik_instance = tik.Tik()

set_current_compile_soc_info用于设置目标昇腾AI处理器的型号及目标核的类型。若不调用此接口，则默认使用的昇腾AI处理器的型号为“Ascend 310”，目标核为“AI Core”，详细接口说明可参见set_current_compile_soc_info。
tik_instance = tik.Tik()用于构造TIK容器，并反馈TIK实例。
tik.Tik()函数中还可以控制是否开启TIK调试功能、设置编译错误信息打印级别。

TIK编译函数

TIK算子的计算逻辑实现完后，需要调用BuildCCE函数，将TIK描述语言编译成可以在昇腾AI处理器上执行的二进制代码，调用方式如下：

tik_instance.BuildCCE(kernel_name="KernelName", inputs=(data_input1_gm,data_input2_gm,), outputs=(data_output_gm,))
或者
tik_instance.BuildCCE(kernel_name="KernelName", inputs=[data_input1_gm,data_input2_gm], outputs=[data_output_g]))

inputs与outputs中是scope为“scope_gm”的Tensor组成的list或者tuple，且输入Tensor与输出Tensor的顺序需要与原型定义、TIK入口函数中的输入输出参数顺序一致。

TIK函数妙用

from tbe import tik
VLENINT32 = 64

def tik_func(tik_instance, data_input_ub, data_output_ub):
    tik_instance.vec_add(VLENINT32, data_output_ub, data_input_ub, data_output_ub, 4, 8, 8, 8)

def tik_vadd(data_input_gm, data_output_gm, dtype, kernel_name):
    tik_instance = tik.Tik()
    # 定义gm和ub
    data_input_gm = tik_instance.Tensor(dtype, (256,), name="data_input_gm", scope=tik.scope_gm)
    data_output_gm = tik_instance.Tensor(dtype, (256,), name="data_output_gm", scope=tik.scope_gm)
    data_input_ub = tik_instance.Tensor(dtype, (256,), name="data_input_ub", scope=tik.scope_ubuf)
    data_output_ub = tik_instance.Tensor(dtype, (256,), name="data_output_ub", scope=tik.scope_ubuf)

    # 数据搬运
    tik_instance.data_move(data_input_ub, data_input_gm, 0, 1, 32, 0, 0)
    # 矢量计算
    tik_func(tik_instance, data_input_ub, data_output_ub)
    # 数据搬运
    tik_instance.data_move(data_output_gm, data_output_ub, 0, 1, 32, 0, 0)

    # 生成CCE
    tik_instance.BuildCCE(kernel_name="tik_vadd", inputs=(data_input_gm,), outputs=(data_output_gm,))
    return tik_instance

if __name__ == "__main__":
    dtype = "int32"
    kernel_name = "tik_vadd"
    data_input_gm = None
    data_output_gm = None
    tik_instance = tik_vadd(data_input_gm, data_output_gm, dtype, kernel_name)

上面这个案例中，我们将vec_add提取了一层TIK函数，封装成了tik_func，需要传入tik实例（tik_instance），其实这里就能够比较容易判断出来TIK函数，如果传入tik_instance，就是TIK函数。

首先需要保证Python语法正确，所以需要传入矢量计算用到的data_input_ub和data_output_ub，然后这个TIK函数没有return data_output_ub，实际上我们对data_output_ub这个变量进行了修改操作，而IR中又显式地inline了。

这里就又体现了TIK代码生成器的本质了，因为Python解释器执行到了vec_add这个TIK语句，所以会生成对应的IR，并不是直接进行计算操作，所以没有将data_output_ub返回也不会有影响。

有时候，我们需要打印出算子的中间Tensor来调试算子功能，如果结合Python本身的语法，我们就可以写出非侵入式的test_tensor。

非侵入式测试

from tbe import tik

class Test:
    __instance=None
    def __init__(self):
        pass
    @classmethod
    def get_instance(cls):
        if not cls.__instance:
            cls.__instance = Test()
        return cls.__instance
    def initial_test(self, tik_instance, dtype, **kwargs):
        self.data_test = tik_instance.Tensor(dtype, kwargs['test_gm_shape'],
                                             name="data_test", scope=tik.scope_gm)
        self.data_test_ub = tik_instance.Tensor(dtype, kwargs['test_ub_shape'],
                                                name="data_test_ub", scope=tik.scope_ubuf)
    def get_test_gm(self):
        return self.data_test
    def get_test_ub(self):
        return self.data_test_ub

先封装一个测试类，测试类中定义ub和gm，个数等按需定义。

from tbe import tik
VLENINT32 = 64

# Test Use
from test_tensor import Test
test = Test.get_instance()

def tik_func(tik_instance, data_input_ub, data_output_ub):
    tik_instance.vec_add(VLENINT32, data_output_ub, data_input_ub, data_output_ub, 4, 8, 8, 8)
    # Test Use
    tik_instance.vec_add(VLENINT32, test.get_test_ub(), data_input_ub, data_output_ub, 4, 8, 8, 8)

def tik_vadd(data_input_gm, data_output_gm, dtype, kernel_name):
    tik_instance = tik.Tik()
    test.initial_test(tik_instance, dtype, test_gm_shape=(256,), test_ub_shape=(256,))
    data_input_gm = tik_instance.Tensor(dtype, (256,), name="data_input_gm", scope=tik.scope_gm)
    data_output_gm = tik_instance.Tensor(dtype, (256,), name="data_output_gm", scope=tik.scope_gm)
    data_input_ub = tik_instance.Tensor(dtype, (256,), name="data_input_ub", scope=tik.scope_ubuf)
    data_output_ub = tik_instance.Tensor(dtype, (256,), name="data_output_ub", scope=tik.scope_ubuf)

    tik_instance.data_move(data_input_ub, data_input_gm, 0, 1, 32, 0, 0)
    tik_func(tik_instance, data_input_ub, data_output_ub)
    tik_instance.data_move(data_output_gm, data_output_ub, 0, 1, 32, 0, 0)
    # Test Use
    tik_instance.data_move(test.get_test_gm(), test.get_test_ub(), 0, 1, 32, 0, 0)

    tik_instance.BuildCCE(kernel_name="tik_vadd",
                          inputs=(data_input_gm,),
                          # Test Use
                          outputs=(data_output_gm, test.get_test_gm(),))
    return tik_instance

if __name__ == "__main__":
    dtype = "int32"
    kernel_name = "tik_vadd"
    data_input_gm = None
    data_output_gm = None
    tik_instance = tik_vadd(data_input_gm, data_output_gm, dtype, kernel_name)

然后将Test类作为全局变量引进来，在入口函数里初始化，这样就可以利用Python的特性非侵入地在整个算子里进行使用，案例中标有Test Use的注释在算子交付的时候，需要修改或删去。而从TIK本质来看，上述实际上就是让Python解释器执行到这一TIK语句，即在算子的某处地方插入了新的IR，从而方便算子的debug。相反，如果我们不使用Test封装类，而是在入口函数里面定义ub或gm，则Python的语法特性反而会使得这个调试过程变得复杂，比如进入一个TIK函数，则需要把测试gm或者ub传进去，这样会改变函数调用接口，修改的地方会比较多（每个调用到改函数接口的地方都要进行霰弹式修改），调试过程显得复杂。

后面的章节我们也会介绍TIK功能调试。本例是说明调试的一种可能性，也为读者们在实际的编码过程中提供一种参考。

总结

TIK函数实际上是借壳于Python函数，没有严格意义上的TIK函数，所谓TIK函数，可以认为是在原来Python函数的基础上加上了TIK语句，其需要符合Python的语法规范，但是由于TIK本身代码生成器的本质，Python解释器执行到TIK语句的时候，就会生成相应的IR，所以TIK变量也不需要返回，但是如果是普通变量或TIK表达式（TIK变量为Expr类型），按照Python的语法，还是需要返回，因为其不会生成相应的IR，如果不返回，则不会产生相应的作用。

父主题： 算子代码实现（TBE TIK）