TIK算子泛化

简介

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

对TIK算子进行泛化有如下两种实现方式：

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

两种实现方式的对比：

方式一：采用静态编译的方式，每次编译出的.o文件只能进行固定shape的计算。在编译的时候决定所有tensor的大小，循环顺序和指令的位置，充分利用空间，减少不必要的scalar操作，可以获得极致的性能。
此种方式开发的TIK算子即可以用于网络中，也可以使用AscendCL接口进行单算子的调用。
方式二：有限的.o解决所有问题，通过运行参数进行分支判断等操作，但会增加scalar操作，导致性能有所下降。
此种方式开发的TIK算子仅可使用AscendCL接口进行单算子的调用，无法应用到网络中。

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

下面我们仅介绍方式一的算子泛化实现注意点，方式二的详细实现方法请参见TIK自定义算子动态Shape专题。

实现

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

算子声明。
因为算子要支持任何合法的数据类型与数据形状，所以这些信息要从算子声明中获取，算子的声明定义如下所示：
```
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的名称、个数及顺序需要与算子原型定义保持一致，可选输入也需要在此处定义，在计算逻辑中去判断是否有数据传入，并进行相应处理。

获取数据类型占用空间和数据形状，并为输入输出张量在内存（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)

因为输入数据大小是从算子声明中动态获取的，所以后续计算时的循环处理次数、指令计算参数等都是不固定的，所以数据在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

开启多核进行计算。

为了充分利用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)

将对应的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算子已实现完毕，此算子的完整样例代码可参见进阶样例。

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