---
title: TIK算子泛化
description: "在TIK简介的示例中，我们仅支持固定的数据输入个数和数据类型，所有的tensor申请的大小和TIK API参数配置为固定值。但是我们若想编写“一类”算子，要能够适应任何合法的数据类型、数据形状和数据排布格式，这种场景就称之为算子的泛化。由于当前版本TIK对数据排布格式不敏感，所以算子的泛化仅考虑数据类型和数据形状的泛化即可。"
url: https://www.hiascend.com/document/detail/zh/canncommercial/latest/others/tbeaicpudevg/atlasopdev_10_0059.html
sourcePath: /source/zh/canncommercial/900/others/tbeaicpudevg/atlasopdev_10_0059.html
indexId: 7744383c54959a9bf28334d029f972fb3372602291b6032cc3415edd43a253c973
---
# TIK算子泛化

#### 简介

在TIK简介的示例中，我们仅支持固定的数据输入个数和数据类型，所有的tensor申请的大小和TIK API参数配置为固定值。但是我们若想编写“一类”算子，要能够适应任何合法的数据类型、数据形状和数据排布格式，这种场景就称之为算子的泛化。由于当前版本TIK对数据排布格式不敏感，所以算子的泛化仅考虑数据类型和数据形状的泛化即可。

对TIK算子进行泛化有如下两种实现方式：1. 方式一：算子的输入shape和数据类型从算子的方法声明中获取，这样就可以做到根据不同的输入shape和数据类型，自动计算出tensor的申请大小，数据循环处理的次数，指令执行的参数等，从而使算子可以适配不同的数据类型和shape。编译的时候将对应的shape和数据类型作为入参传进去，同一份代码根据不同的输入编译出对应的.o，编译出的.o仅仅解决对应shape和数据类型的计算，运行的时候不需要额外的运行参数，只需要传入输入、输出的地址。
2. 方式二：按照最大的内存需求申请空间，编译的时候将shape范围作为编译参数传入，编译出的.o可以支持一类shape计算，运行的时候除了输入、输出的地址，还需要将输入输出shape作为运行参数传进来。


两种实现方式的对比：1. 方式一：采用静态编译的方式，每次编译出的.o文件只能进行固定shape的计算。在编译的时候决定所有tensor的大小，循环顺序和指令的位置，充分利用空间，减少不必要的scalar操作，可以获得极致的性能。
  此种方式开发的TIK算子即可以用于网络中，也可以使用AscendCL接口进行单算子的调用。

2. 方式二：有限的.o解决所有问题，通过运行参数进行分支判断等操作，但会增加scalar操作，导致性能有所下降。
  此种方式开发的TIK算子仅可使用AscendCL接口进行单算子的调用，无法应用到网络中。


开发人员可以根据实际的情况去选择适合自己的编译和运行方式。

下面我们以方式一为例，介绍算子泛化实现注意点。


#### 实现

下面仍以运行在AI Core上的Vadd算子为例，讲解进行算子泛化时的注意点。

1. 算子声明。
  因为算子要支持任何合法的数据类型与数据形状，所以这些信息要从算子声明中获取，算子的声明定义如下所示：

```
def vadd_sample(input_x, input_y, output_z, kernel_name):
```


  其中input_x、input_y是Vadd算子的两个输入，字典格式，包含shape、ori_shape、format、ori_format与dtype信息；output_z为算子的输出，也是字典格式，为预留位。

  其中输入tensor与输出tensor的名称、个数及顺序需要与算子原型定义保持一致，可选输入也需要在此处定义，在计算逻辑中去判断是否有数据传入，并进行相应处理。

2. 获取数据类型占用空间和数据形状，并为输入输出张量在内存（Global Memory）中开辟空间。
  TIK支持的dtype包含uint8、int8、uint16、int16、float16、uint32、int32、float32、uint64、int64，这些dtype都能从类型名称的字符串中获取所占用空间，所以可以定义get_bit_len函数，根据输入dtype名称，计算输出dtype占用空间，单位为“bit”。函数实现如下所示：

```
def get_bit_len(dtype):
index = 0
for i in dtype:
if i.isdigit():
break
index += 1
return int(dtype[index:])
```


  然后动态的根据输入来获取输入shape和dtype，若算子的计算过程中需要其他参数，可同样从输入中获取。

```
class Vadd():
def __init__(self, input_x, input_y, output_z, kernel_name="vadd_sample"):
# 获取输入张量的形状和数据类型，此处仅获取input_x即可，对于Vadd算子，开发者需确保两个输入的dtype与shape一致
self.shape_x = input_x.get("shape")
self.dtype_x = input_x.get("dtype")
self.shape_y = input_y.get("shape")
self.dtype_y = input_y.get("dtype")
self.shape_z = output_z.get("shape")
self.dtype_z = output_z.get("dtype")
self.kernel_name = kernel_name
# 请根据实际AI处理器型号进行设置
soc_version="xxx"
# 目标核类型为默认的AI Core
tbe_platform.set_current_compile_soc_info(soc_version)
self.tik_instance = tik.Tik()
```


  为输入、输出张量在内存中申请空间。

```
self.input_x_gm = self.tik_instance.Tensor(
self.dtype_x, self.shape_x, name="input_x_gm", scope=tik.scope_gm)
self.input_y_gm = self.tik_instance.Tensor(
self.dtype_y, self.shape_y, name="input_y_gm", scope=tik.scope_gm)
self.output_z_gm = self.tik_instance.Tensor(
self.dtype_z, self.shape_z, name="output_z_gm", scope=tik.scope_gm)
```


3. 因为输入数据大小是从算子声明中动态获取的，所以后续计算时的循环处理次数、指令计算参数等都是不固定的，所以数据在Unified Buffer中进行计算前需要计算好相关值，例如一共有多少个AI Core，Unified Buffer中能存放多少个元素，每个迭代计算多少个元素等，为后续计算做准备。

```
# 获取AI Core的个数
self.aicore_num = tbe_platform.get_soc_spec("CORE_NUM")
# UB上数据读取和写入必须32B对齐，此参数用来计算Tensor划分和数据搬运指令
block_byte_size = 32
# 获取UB空间大小，单位为Bytes
ub_size_bytes = tbe_platform.get_soc_spec("UB_SIZE")
# 计算输入的数据类型对应多少个字节
dtype_bytes_size = get_bit_len(self.dtype_x) // 8
# 根据输入的数据类型计算一个block可以存放多少个对应的元素
self.data_each_block = block_byte_size // dtype_bytes_size
# ub_size_bytes//dtype_bytes_size，得到UB总共能存放多少个数；有2个输入张量要存放，所以每个张量最多能保存的元素数量是ub_size_bytes // dtype_bytes_size // 2
# UB的数据读写必须32B对齐，故再在上一步的基础上做个对齐，假设上一步得到的中间结果为data_per_tensor，则对齐操作为：data_per_tensor//data_each_block*data_each_block，得出在UB上，每个张量应该有多少个元素
self.ub_tensor_size = (
ub_size_bytes // dtype_bytes_size // 2 // self.data_each_block * self.data_each_block)
# 计算输入的元素总个数
self.input_num = functools_reduce(lambda x, y: x * y, self.shape_x)
# 计算每个AI Core需要处理的数据量，当前只考虑均分场景，且均分后32 Bytes对齐
self.data_num_each_core = self.input_num // self.aicore_num
# 计算每次repeat能计算多少个元素，vector指令每个repeat最多计算8个block（256Bytes），此时mask取的对应数据类型的最大值。
self.vector_mask_max = 8 * self.data_each_block
```


4. 开启多核进行计算。
  为了充分利用AI Core，需要使用for_range函数开启多核并行计算。存储在Unified Buffer中的Tensor定义，需要写在for_range多核循环内，如下所示：

```
def vadd_compute(self):
with self.tik_instance.for_range(0, self.aicore_num, block_num=self.aicore_num) as index:
# 创建在Unified Buffer上的tensor，shape就是上一步骤计算出来的“ub_tensor_size”
self.input_x_ub = self.tik_instance.Tensor(self.dtype_x, (self.ub_tensor_size,),name="input_x_ub",scope=tik.scope_ubuf)
self.input_y_ub = self.tik_instance.Tensor(self.dtype_y, (self.ub_tensor_size,),name="input_y_ub",scope=tik.scope_ubuf)
```


5. 将对应的Global Memory中的数据搬运到Unified Buffer，并进行计算，需要注意每次搬运的偏移量为已经处理过的数据个数。

```
# index为AI Core的索引序号
move_offset = index * self.data_num_each_core
# 每个AI Core负责自己的数据分片
self.vadd_compute_each_core(move_offset, self.data_num_each_core)
```


  下面详细讲解每个AI Core的计算逻辑，即上述中的“vadd_compute_each_core”函数。

  此函数的实现也是泛化算子实现的难点，因为UB一次能放下248KB数据，也就是最多124KB的张量A与124KB的张量B。而且一条vec_add指令计算的数据量也是有限的，所以需要进行多次循环计算。

```
def vadd_compute_each_core(self, move_offset, move_num):
# 计算循环次数
loop_time = move_num // self.ub_tensor_size
# 如果数据量连一次循环都不够，直接走尾块计算即可
if loop_time > 0:
# 经典的for_range循环
with self.tik_instance.for_range(0, loop_time) as loop_index:
# move_offset要在计算过程中持续刷新，相当于一个指针在内存中移动
move_offset = loop_index * self.ub_tensor_size
# 把偏移量和计算量传递给循环计算函数
self.vadd_compute_each_loop(move_offset, self.ub_tensor_size)
move_offset = loop_time * self.ub_tensor_size
# 对不够塞满UB的最后一点数据进行处理
last_num = move_num % self.ub_tensor_size
if last_num > 0:
self.vadd_compute_each_loop(move_offset, last_num)
```


  循环处理函数“vadd_compute_each_loop”的实现如下所示：

```
def vadd_compute_each_loop(self, move_offset, move_num):
# 把数据从内存中搬进UB
burst_len = math.ceil(move_num / self.data_each_block)
self.tik_instance.data_move(self.input_x_ub, self.input_x_gm[move_offset], 0, 1, burst_len, 0, 0)
self.tik_instance.data_move(self.input_y_ub, self.input_y_gm[move_offset], 0, 1, burst_len, 0, 0)
# 调用vec_add执行计算任务
# 首先计算数据总量够填满多少个vec_add指令
vadd_loop = move_num // (self.vector_mask_max * 255)
add_offset = 0
if vadd_loop > 0:
# 循环执行vec_add计算
with self.tik_instance.for_range(0, vadd_loop) as add_index:
add_offset = add_index * self.vector_mask_max * 255
self.tik_instance.vec_add(self.vector_mask_max,
self.input_x_ub[add_offset],
self.input_x_ub[add_offset],
self.input_y_ub[add_offset],
255, 8, 8, 8)
add_offset = vadd_loop * vector_mask_max * 255
# 计算上一步剩余的数据量够填满多少个迭代，并调用一次vec_add进行计算
repeat_time = (move_num % (self.vector_mask_max * 255) // self.vector_mask_max)
if repeat_time > 0:
self.tik_instance.vec_add(self.vector_mask_max,
self.input_x_ub[add_offset],
self.input_x_ub[add_offset],
self.input_y_ub[add_offset],
repeat_time, 8, 8, 8)
add_offset += repeat_time * self.vector_mask_max
# 计算上一步剩余的数据量有多少，直接作为mask值调用最后一次vec_add
last_num = move_num % self.vector_mask_max
if last_num > 0:
self.tik_instance.vec_add(last_num,
self.input_x_ub[add_offset],
self.input_x_ub[add_offset],
self.input_y_ub[add_offset],
1, 8, 8, 8)
# 把计算结果从UB中搬回内存
self.tik_instance.data_move(self.output_z_gm[move_offset],
self.input_x_ub, 0, 1, burst_len, 0, 0)
```


  至此，泛化的Vadd算子已实现完毕，此算子的完整样例代码可参见进阶样例。