算子实现

代码结构

不同Tiling策略下，代码结构类似，如下所示：

import os
import warnings

from tbe import tik
from tbe.common.platform import set_current_compile_soc_info

#tiling mode 2 ub size
TILING_1_UB_SIZE = 112*1024

#max allow general mem size
TOTAL_GM_SIZE = 512*1024*1024   #GM_MAX_SIZE

# batch for N
MAX_BATCH = 1

# channel for C
MAX_CHANNEL = 1024

# width for W
MAX_WIDTH = 512

# height for H
MAX_HEIGHT = 1024

class BatchNorm(object):
    # 初始化函数
    def __init__(self, input0, gamma0, beta0, output0, kernel_name = "BatchNorm"): 
      """
      初始化函数实现
      """
    # 算子计算逻辑实现和编译
    def batchnorm_compute(self):
       # 不同Tiling策略下计算逻辑不同
       self.batchnorm_compute_tiling_c()
     
        # 算子编译
        self.tik_instance.BuildCCE(kernel_name=self.kernel_name,
                                   inputs=[self.input_gm,
                                           self.gamma_gm,
                                           self.beta_gm],
                                   outputs=[self.output_gm],
                                   flowtable=[self.input_n, self.input_c,
                                              self.input_h, self.input_w,
                                              self.inputtype, self.output_n,
                                              self.output_c, self.output_h,
                                              self.output_w, self.outputtype,
                                              self.gamma_c, self.gammatype,
                                              self.beta_c, self.betatype,
                                              self.param1, self.param2,
                                              self.param3, self.param4,
                                              self.param5, self.param6,
                                              self.param7, self.param8,
                                              self.param9, self.param10],
                                   config={"double_buffer_non_reuse": True,
                                           "out_of_bound_sync_check": True})
        return self.tik_instance

# 算子定义函数
def batch_norm(input_x, gamma, beta, output, kernel_name="BatchNorm"):
    obj = BatchNorm(input_x, gamma, beta, output, kernel_name)
    obj.batchnorm_compute()

如下几点需要关注：

• 不同Tiling策略下，代码结构类似，主要区别点在于不同Tiling策略下计算逻辑不同。
• 通过BuildCCE接口编译TIK算子时，可以传入flowtable参数。该参数为算子选择器计算好的Tiling参数，相当于在output容器后增加一段用户缓存Tiling参数的地址空间，flowtable列表中的Tiling参数个数和inputs的参数个数之和小于等于64，每个Tiling参数为TIK InputScalar类型，Tiling参数不在TIK中计算，将重复的Scalar运算全部转移到Host CPU上的算子选择器计算，减少AI Core上Scalar运算量，有利于流水线并行。