概述

nntc中包含了丰富的算子库，可以满足大部分神经网络模型的编译需求。但在某些场景下，可能需要用户自定义算子来实现对tensor的计算。如：

1. nntc还未支持的算子，且无法通过其它算子组合实现

2. 算子为用户私有，未对公众开源

3. 使用多个算子API组合无法取得最佳计算性能，直接从tpu-kernel层自定义运算可以提高运行效率

自定义算子功能允许用户自由使用tpu_kernel中的接口，自定义tensor在tpu上的计算，并将这一计算过程封装为算子。与conv2d，relu等算子类似，自定义算子可以直接接入bmcompiler，再通过BMLang进行调用。

为了实现自定义算子功能，用户需要分别编写Device 代码、Plugin 代码并生成对应的动态库libkernel_module.so、libplugin.so供编译器识别，具体实现细节可以参考tpu-nntc中的demo。以下分别进行介绍：

Device 代码

Device 代码主要包括算子的定义，以及调用 tpu-kernel 实现计算的部分。同时，还需要编写设备上的分发函数代码，以保证编译出的程序能在 TPU 上正确识别。 算子计算涉及global layer，local layer，reshape相关操作。

1. 算子必须实现 global layer ， global layer的输入和输出数据都放在ddr中，数据需要从global mem搬运到local mem中执行运算，再将结果搬运至global mem。其优点是local mem可以任意使用，比较灵活；缺点是会产生较多的gdma搬运，tpu利用率较低。

2. 算子根据需要实现 local layer，local layer的输入和输出的数据都放在local mem中，可以与其他layer组合进行layer group优化，避免该layer计算时数据要搬入搬出到global mem中。其优点是可以节省gdma搬运, 运算效率高；缺点是比较复杂，local mem由bmcompiler提前分配好，不可任意使用，在部分算子中无法实现。

3. 算子必须实现reshape操作，需要根据算子参数，输入数据来确定输出数据的ttype，dtype，shape信息。

Device 代码最终会被编入libkernel_module.so 。

Plugin 代码

plugin代码主要包括layer的描述符接口和注册相关代码，供bmcompiler内部优化使用。首先需要继承 bmcompiler::LayerDescriptor ，设置bmcompiler用到的一些属性，其次需要添加函数注册信息register_tpu_layers。 bmcompiler会根据op_type分发调用相关接口函数。

根据不同算子的需求，Descriptor可能需要实现更多接口，可以参考 LayerDescriptor的详细定义(在bmcompiler_tpulayer_host.hpp中)：

class LayerDescriptor {
   public:
   // 通过该接口可以得到layer实例，进而得到layer的参数，输入及输出
   GraphNode* layer() const;

   // 返回layer的name，供日志显示
   virtual std::string get_name()  const { return "UNKNOWN"; }

   // 如果在tpukernel中使用了cdma或gde，必须返回true，以保证动态运行
   virtual bool must_dynamic_run() const { return false; }

   // 如果输出的shape无法根据输入的shape和layer参数确定(通常取决输入的数据， 比如where操作)，返回true
   virtual bool output_shape_is_dynamic() const { return false; }

   // 强制用户使用global layer编译
   virtual bool force_use_global_layer() const { return false; }

   // 该算子是否支持对h/w进行切分，如不支持，bmcompiler会在安排local layer时尝试将整个tensor放入到local mem中
   virtual bool can_split_local_h_w() const { return true; }

   // 该层是否可做原地操作
   virtual bool can_inplace(DATA_TYPE_T data_prec) const { return false; }

   // if the layer can support size of h/w dimension is dynamic
   virtual bool can_h_w_dynamic() const { return true; }

   // 是否仅更改的input的shape, output数据和input数据内容是一样的
   virtual bool just_reshape() const { return false; }

   // 输入输出能否共享同一块mem，如是，则输入输出的mem地址相同
   virtual bool is_share_mem_layer() const { return false; }

   // 针对global layer，是否需要buffer， 如果返回值大于0，则会申请相应大小的buffer
   virtual u64 get_global_buffer_size() const { return 0; }

   // 设置第in_idx个tensor的类型
   virtual TENSOR_TYPE_T input_type(int in_idx) const { return BMNET_NEURON; }

   //---------------------------------------------------------------------------
   // 以下参数仅用于local layer
   // 如果未实现local layer，可以忽略

   // 针对4维的split_shape大小，返回需要的buffer大小
   virtual u32 get_local_buffer_size(const int* split_shape) const { return 0; }

   // 该layer在反向推断data slice时，是否支持相同tensor可设置不同data slice
   virtual bool is_backslice_optimize_support_layer() const { return false; }

   // 当支持反向slice大小推断时，要提供以下参数
   // backward
   virtual int backward_hslice(int out_hslice) const { return out_hslice; }
   virtual int backward_wslice(int out_wslice) const { return out_wslice; }
   virtual int backward_hidx(int out_hidx) const { return out_hidx; }
   virtual int backward_widx(int out_widx) const { return out_widx; }
   virtual int backward_kh(int out_kh) const { return out_kh; }
   virtual int backward_stride_h(int out_stride_h) const { return out_stride_h; }
   virtual int backward_up_pad_h(int out_up_pad_h) const { return out_up_pad_h; }
   virtual int backward_down_pad_h(int out_down_pad_h) const { return out_down_pad_h; }

   // 该layer在做local layer时，第in_idx个输入tensor是否要在l2 mem上，默认是在local mem上。
   virtual bool should_input_on_l2mem(int in_idx) const { return false; };

   // forward
   virtual int forward_height(int in_height) const { return in_height; }
   };

plugin代码最终会被编译成libplugin.so。

BMLang接口

用户可以通过BMLang接口使用 custom_tpu_op 调用自定义算子，该接口的定义如下：

void custom_tpu_op(const std::vector<Tensor*> tensor_i, // [in]
                  std::vector<Tensor*> tensor_o, // [out]
                  int op_type, // op_type
                  void* param, // op_param
                  int param_size); // sizeof(param)

其python接口如下：

custom_tpu_op(tensor_i, # [in]
            tensor_o, # [out]
            op_type, # op_type
            param) # op_param

目前该接口仅支持 COMPILE_TPU 模式。

示例程序

验证Demo

在tpu-nntc内已包含三个自定义算子的示例，分别是abs_add、ceil_add、rpn。在具体开发工作前，可以先通过以下步骤验证 demo 的正确性和完整性。 首先请参考 “TFLite 量化模型编译示例-初始化tpu-nntc环境” 一节，确保你已经正确进入最新版本的docker，并初始化tpu-nntc的软件环境。

# 已经进入docker，并在/workspace目录下
# 下面初始化软件环境
cd /workspace/tpu-nntc
source scripts/envsetup.sh

# 切换到userlayer目录并初始化userlayer环境
cd $NNTC_TOP/userlayer/
source envsetup.sh

接下来通过以下步骤编译测试abs_add算子:

# 编译生成libkernel_module.so 和 libplugin.so
cd $NNTC_TOP/userlayer/abs_add/
mkdir build && cd build
cmake ..
make
# 设置plugin和kernel_module的路径
export BMCOMPILER_PLUGIN_PATH=$PWD/plugin/
export BMCOMPILER_KERNEL_MODULE_PATH=$PWD/device/libkernel_module.so
# 测试case0和case1的编译
python3 ../compile_case0.py
python3 ../compile_case1.py

如果命令行输出以下文字，说明模型编译成功：

=====================================
*** Store bmodel of BMCompiler...
=====================================

采用同样的方法，编译测试ceil_add算子：

cd $NNTC_TOP/userlayer/ceil_add/
mkdir build && cd build
cmake ..
make
export BMCOMPILER_PLUGIN_PATH=$PWD/plugin/
export BMCOMPILER_KERNEL_MODULE_PATH=$PWD/device/libkernel_module.so
python3 ../compile_case0.py

采用同样的方法，可以编译测试rpn算子：

cd $NNTC_TOP/userlayer/rpn_demo/
mkdir build && cd build
cmake ..
make
export BMCOMPILER_PLUGIN_PATH=$PWD/plugin/
export BMCOMPILER_KERNEL_MODULE_PATH=$PWD/device/libkernel_module.so
python3 ../compile_case0.py

代码结构

以下以abs_add算子为例，其目录结构如下：

abs_add
├── CMakeLists.txt
├── cmodel
│   └── CMakeLists.txt
├── compile_case0.py
├── compile_case1.py
├── device
│   ├── CMakeLists.txt
│   ├── device_absadd.c
│   └── device_entry.c
├── include
│   ├── bmtpu.h
│   └── host_absadd.h
└── plugin
   ├── CMakeLists.txt
   ├── host_absadd.cpp
   └── host_entry.cpp

以下对各部分代码的实现作简单介绍。

Device代码

第一步，定义算子的参数 param，在include/bmtpu.h中添加算子的标识符 TPU_ABSADD。

#ifndef __BMTPU_H__
#define __BMTPU_H__
#ifndef MAX_SHAPE_DIMS
#define MAX_SHAPE_DIMS 8
#endif

typedef enum {
TPU_ABSADD,
} tpu_op_type_t;

#endif //__BMTPU_H__

第二步，在device/device_absadd.c中定义算子需要的参数tpu_absadd_param_t，并实现global layer, local layer和reshape接口。

#include <inttypes.h>
#include "tpu_kernel.h"
#include "bmcompiler_defs.h"
#include "bmcompiler_tpulayer_device.h"

typedef struct {
   float b_val;
} tpu_absadd_param_t;

int absadd_global(const void* param, const int param_size,
               const tpu_tensor_t* in_tensors, const int in_num,
               const bm_addr_t buffer_addr, const bm_size_t buffer_size,
               tpu_tensor_t* out_tensors, const int out_num) {
   tpu_absadd_param_t* absadd_param = (tpu_absadd_param_t*)param;
   dim4 shape = {.n=in_tensors->shape[0],
                  .c=in_tensors->shape[1],
                  .h=in_tensors->shape[2],
                  .w=in_tensors->shape[3]};
   global_addr_t glb_in_addr = in_tensors[0].addr;
   global_addr_t glb_out_addr = out_tensors[0].addr;
   local_addr_t local_in_addr = 0;
   local_addr_t local_mid_addr = BANK_SIZE * 4;
   local_addr_t local_out_addr = BANK_SIZE * 8;
   tpu_gdma_cpy_S2L(local_in_addr, glb_in_addr, &shape, NULL, NULL, DT_FP32);
   tpu_bdc_abs(local_mid_addr, local_in_addr, &shape, NULL, NULL, DT_FP32);
   scalar_t C = {.f32 = absadd_param->b_val};
   tpu_bdc_fp_add_C(local_out_addr, local_mid_addr, C, &shape, NULL, NULL, DT_FP32);
   tpu_gdma_cpy_L2S(glb_out_addr, local_out_addr, &shape, NULL, NULL, DT_FP32);
   return 0;
}

int absadd_local(const void* param, const int param_size,
            const tpu_tensor_t* in_tensors, const int in_num,
            const u64 buffer_offset, const u64 buffer_size,
            tpu_tensor_t* out_tensors, const int out_num) {
   tpu_absadd_param_t* absadd_param = (tpu_absadd_param_t*)param;
   dim4 shape = {.n=in_tensors->slice_shape[0],
                  .c=in_tensors->slice_shape[1],
                  .h=in_tensors->slice_shape[2],
                  .w=in_tensors->slice_shape[3]};
   local_addr_t local_in_addr = in_tensors->addr;
   local_addr_t local_mid_addr = BANK_SIZE * 12;
   local_addr_t local_out_addr = out_tensors->addr;
   tpu_bdc_abs(local_mid_addr, local_in_addr, &shape, NULL, NULL, DT_FP32);
   scalar_t C = {.f32 = absadd_param->b_val};
   tpu_bdc_fp_add_C(local_out_addr, local_mid_addr, C, &shape, NULL, NULL, DT_FP32);
   return 0;
}

int absadd_reshape(const void* param, const int param_size,
               const tpu_tensor_t* in_tensors, const int in_num,
               tpu_tensor_t* out_tensors, const int out_num){
   out_tensors[0].dims = in_tensors[0].dims;
   out_tensors[0].shape[0] = in_tensors[0].shape[0];
   out_tensors[0].shape[1] = in_tensors[0].shape[1];
   out_tensors[0].shape[2] = in_tensors[0].shape[2];
   out_tensors[0].shape[3] = in_tensors[0].shape[3];
   out_tensors[0].dtype = in_tensors[0].dtype;
   out_tensors[0].ttype = in_tensors[0].ttype;
   return 0;
}

第三步，在 device/device_entry.c 中添加设备上的分发函数的接口。

#include "bmcompiler_tpulayer_device.h"
#include "bmtpu.h"

int absadd_global(const void* param, const int param_size,
const tpu_tensor_t* in_tensors, const int in_num,
const bm_addr_t buffer_addr, const bm_size_t buffer_size,
tpu_tensor_t* out_tensors, const int out_num);
int absadd_local(const void* param, const int param_size,
const tpu_tensor_t* in_tensors, const int in_num,
const bm_addr_t buffer_addr, const bm_size_t buffer_size,
tpu_tensor_t* out_tensors, const int out_num);
int absadd_reshape(const void* param, const int param_size,
const tpu_tensor_t* in_tensors, const int in_num,
tpu_tensor_t* out_tensors, const int out_num);

static inline int dummy_shape_infer_entry(
                     const void* param, const int param_size,
                     const tpu_tensor_t* in_tensors, const int in_num,
                     tpu_tensor_t* out_tensors, const int out_num){
   (void) param;
   (void) param_size;
   (void) in_tensors;
   (void) in_num;
   (void) out_tensors;
   (void) out_num;
   return BM_LAYER_NOT_SUPPORTED;
}

static inline int dummy_calculate_entry(
                     const void* param, const int param_size,
                     const tpu_tensor_t* in_tensors, const int in_num,
                     const bm_addr_t buffer_addr, const bm_size_t buffer_size,
                           tpu_tensor_t* out_tensors, const int out_num){
   (void) param;
   (void) param_size;
   (void) in_tensors;
   (void) in_num;
   (void) buffer_addr;
   (void) buffer_size;
   (void) out_tensors;
   (void) out_num;
   return BM_LAYER_NOT_SUPPORTED;
}

int tpu_global_calculate_entry(int op_type,
                     const void* param, const int param_size,
                     const tpu_tensor_t* in_tensors, const int in_num,
                     const bm_addr_t buffer_addr, const bm_size_t buffer_size,
                     tpu_tensor_t* out_tensors, const int out_num){
   if(op_type==TPU_ABSADD){
      return absadd_global(param, param_size, in_tensors, in_num, buffer_addr, buffer_size, out_tensors, out_num);
   }
   return dummy_calculate_entry(param, param_size, in_tensors, in_num, buffer_addr, buffer_size, out_tensors, out_num);
}

int tpu_shape_infer_entry(int op_type,
                     const void* param, const int param_size,
                     const tpu_tensor_t* in_tensors, const int in_num,
                     tpu_tensor_t* out_tensors, const int out_num){
   if(op_type==TPU_ABSADD){
      return absadd_reshape(param, param_size, in_tensors, in_num, out_tensors, out_num);
   }
   return dummy_shape_infer_entry(param, param_size, in_tensors, in_num, out_tensors, out_num);
}

Plugin代码

第一步，需要在 include/host_absadd.h 中通过继承 LayerDescriptor 实现 Descriptor 类。abs_add算子只需实现 get_name() 接口。

#include "bmcompiler_tpulayer_host.hpp"
// __BM_DESC__(TPU_ABSADD)
class ABSADDDescriptor: public bmcompiler::LayerDescriptor {
public:
   virtual std::string get_name() const override { return "ABSADD"; }
};

第二步，在 plugin/host_entry.cpp 中添加算子和接口注册的代码，以供编译器识别

#include "bmtpu.h"
#include "bmcompiler_tpulayer_host.hpp"

#include "host_absadd.h"
extern "C" {
int absadd_global(const void* param, const int param_size,
const tpu_tensor_t* in_tensors, const int in_num,
const bm_addr_t buffer_addr, const bm_size_t buffer_size,
tpu_tensor_t* out_tensors, const int out_num);
int absadd_local(const void* param, const int param_size,
const tpu_tensor_t* in_tensors, const int in_num,
const bm_addr_t buffer_addr, const bm_size_t buffer_size,
tpu_tensor_t* out_tensors, const int out_num);
int absadd_reshape(const void* param, const int param_size,
const tpu_tensor_t* in_tensors, const int in_num,
tpu_tensor_t* out_tensors, const int out_num);

bmcompiler::tpu_layer_def_t* register_tpu_layers(int* num){
   static ABSADDDescriptor inst_ABSADDDescriptor;
   static bmcompiler::tpu_layer_def_t layer_list[] = {
      {TPU_ABSADD, absadd_global, absadd_local, absadd_reshape, &inst_ABSADDDescriptor},
   };
   if(num){
      *num = sizeof(layer_list)/sizeof(layer_list[0]);
   }
   return layer_list;
}
}

第三步，添加 plugin/host_absadd.cpp，为空文件。

BMLang调用

在BMLang中，通过 custom_tpu_op 调用 abs_add 算子实现对tensor的计算。测试示例compile_case1.py的代码如下：

import bmlang
import numpy as np
import ctypes

# 0 使用ctypes封装算子参数，结构与device/device_absadd.c中的tpu_absadd_param_t保持一致
class ABSADDParam(ctypes.Structure):
   _fields_ = [
      ("b_val", ctypes.c_float)]

dtype='float32'
bmlang.init('bm1684x', assist=True)
param = ABSADDParam()
param.b_val = 1.23

# 1 定义输入和输出Tensor
in_shape = [1,18,14,14]
in_data = np.random.random(size=in_shape).astype('float32')
t_input = bmlang.Tensor('in', dtype=dtype, shape=in_shape, data=in_data)
t_output = bmlang.Tensor(name='out', dtype=dtype)

# 2 添加模型的计算层
y = bmlang.custom_tpu_op([t_input], [t_output], 0, param)
z = bmlang.relu(y[0])
o = bmlang.concat([y[0], z],axis=0)

# 3 编译模型
bmlang.compile('absadd_concat_demo', inputs=[t_input])

# 4 deinit bmlang
bmlang.deinit()

# 5 保存ref_data (可选)
in_data.tofile("compilation/input_ref_data.dat")
y_ref = np.abs(in_data).astype('float32') + param.b_val
z_ref = y_ref * (y_ref > 0)
ref_data = np.concatenate((y_ref, z_ref), axis=0)
ref_data.tofile("compilation/output_ref_data.dat")