--- title: TIK函数 description: "经过前几章的学习，相信读者已经掌握了编写TIK算子的数据搬运和矢量计算的技能，以及了解了TIK作为代码生成器的本质，也已经可以编写一些最基本的TIK算子。" url: https://www.hiascend.com/document/detail/zh/canncommercial/latest/others/tbeaicpudevg/atlasopdev_10_0057.html sourcePath: /source/zh/canncommercial/900/others/tbeaicpudevg/atlasopdev_10_0057.html indexId: d3c5b3826ca31352811a022b19626eea422a2c34a0639ec2ae764bf002afbe5a73 --- # 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，如果不返回，则不会产生相应的作用。