知乎链接：https://zhuanlan.zhihu.com/p/611606630

若是初学者，建议先看前面的，尤其是其中的关于MagicMind的模块部分。

建议结合实测试和同步查看，这样更容易理解各步骤的意义。

1、自定义算子概述

1）为什么需要自定义算子？

MagicMind提供了丰富的算子库，能够满足绝大多数场景的使用需求。但出于以下原因，可能需要定制化算子的实现，从而满足特定需求：

• 某些运算逻辑没有对应算子，或无法组合基础 API 实现

• 通过算子融合追求更高的计算性能

2）流程图

2、自定义算子实现步骤

1 自定义算子定义

先看第一部分：

自定义算子定义就是根据算子的特性，将算子对应的输入、输出，输入和输出数据类型、形状推导函数、属性信息通过调用 IPluginDefBuilder相关接口注册到 Magicmind算子库中； PLUGIN_REGISTER_OP 宏是自定义算子注册的入口。

以 C++ API 为例， PluginAdd的注册示例

PLUGIN_REGISTER_OP("PluginAdd")
   .Input("input1").TypeConstraint("T")
   .Input("input2").TypeConstraint("T")
   .Output("output").TypeConstraint("T")
   .Param("T").Type("type")
   .Default(DataType::FLOAT32)
   .ShapeFn(DoShapeInfer);

接下来我们介绍下注册算子名字，输入输出及输入输出数据类型，形状推导的语法

形状推导是根据输入形状推导输出形状，是必须实现并注册的；形状推导接口的输入参数 IShapeInferResource 以键值对方式（name, value），存储了算子编译和运行期所需要的各种参数，包括输入输出形状、输入输出地址、算子参数等，用户可以通过 name 获取到所需要的参数；下面示例代码为 PluginAdd 的形状推导函数实现；该算子输出形状和输入形状一致，所以只需要获取输入形状并设置给输出形状即可。

Status DoShapeInfer(IShapeInferResource *context) {
    Status ret;
    std::vector<int64_t> input_shape;
    // 获取输入'input1'的形状，获取成功返回Status::OK()，获取失败返回其它错误    ret = context->GetShape("input1", &input_shape);
    if (ret != Status::OK()) {
        return ret;
    }
        // 设置输入'input1'的形状，设置成功返回Status::OK()，设置失败返回其它错误        ret = context->SetShape("output", input_shape);
    return ret;}

除了上述介绍的注册接口，还有其他属性的注册接口，介绍如下

复杂算子注册示例

PLUGIN_REGISTER_OP("PluginFusedMatMul") // 自定义算子名    .Input(“a”).TypeConstraint(“T”) // 注册输入 Input a，并对其进行属性设置    .Input(“b”).TypeConstraint(“T”) // 注册输入 Input b，并对其进行属性设置    .Input(“args”).TypeConstraint(“T”).NumberConstraint(“num_args”) // 注册输入 args，并对其进行属性设置    .Output(“product”).TypeConstraint(“T”)  // 注册输出 product，并对其进行属性设置    .Param(“transpose_a”).Type(“bool”).Default(false) // 注册 transpose_a 属性，并对其进行属性设置    .Param(“transpose_b”).Type(“bool”).Default(false) // 注册 transpose_b 属性，并对其进行属性设置    // 注册自定义算子的 Param 范型描述 T，并对其进行属性设置    .Param("T").Type("type").Allowed({DataType::FLOAT16,DataType::FLOAT32})
    .Param(“num_args”).Type(“int”).Minimum(0)  // 注册 num_args 属性，并对其进行属性设置    .Param(“fused_ops”).TypeList(“string”).Default({“op”})// 注册 fused_ops属性，并对其进行属性设置    .Param(“epsilon”).Type(“float”).Default(0.0001f) // 注册 epsilon 属性，并对其进行属性设置    // 注册 leakyrelu_alpha 属性，并对其进行属性设置    .Param("leakyrelu_alpha").Type("float").Default(0.2f) 
    .ShapeFn([](IShapeInferResource* rsc)->Status { // 注册自定义算子的形状推导函数        std::vector<int64_t> input_shape;
        rsc->GetShape("input1", &input_shape);
        rsc->SetShape("output", input_shape);
        return Status::OK();
    });

2 实现 MagicMind 算子 Kernel

MagicMind算子Kernel的实现，这一步主要实现算子的执行行为

Kernel定义了算子的执行行为，会在运行期被调用。所有自定义算子的 Kernel 都派生自 IPluginKernel 基类，并要求实现 SetLocalVar、GetWorkspaceSize、Enqueue 接口。下面示例为派生自 IPluginKernel 的 PluginAddKernel 类。

class PluginAddKernel : public IPluginKernel {
 public:
  // 完成参数检查和用户自定义成员变量的初始化  Status SetLocalVar(INodeResource *context) override;
  // 获取加法操作运行时所需要的额外内存（如果有的话）  size_t GetWorkspaceSize(INodeResource *context) override;
  // 执行加法运算  Status Enqueue(INodeResource *context) override;
  ~PluginAddKernel();
 private:
  uint32_t input_count;};

SetLocalVar、GetWorkspaceSize、Enqueue 接口的作用说明如下：

SetLocalVar、GetWorkspaceSize、Enqueue 接口代码如下：

#  SetLocalVarStatus PluginAddKernel::SetLocalVar(INodeResource *context) {
  std::vector<int64_t> input_shape;
  context->GetTensorShape("input1", &input_shape);
  if (input_shape.size() < 1) {
    return Status(magicmind::error:Code::INVALID_ARGUMENT, "FAILED");   }
  input_count = GetInputCount(input_shape);
  return Status::OK();}# GetWorkspaceSizesize_t PluginAddKernel::GetWorkspaceSize(INodeResource *context) {
  size_t workspace_size = 0;
  return workspace_size;}# EnqueueStatus PluginAddKernel::Enqueue(INodeResource *context) {
  float *input1_addr = nullptr;
  float *input2_addr = nullptr;
  float *output_addr = nullptr;
  context->GetTensorDataPtr("input1", &input1_addr);
  context->GetTensorDataPtr("input2", &input2_addr);
  context->GetTensorDataPtr("output", &output_addr);
  for (int i = 0; i < input_count; i++) {
    output_addr[i] = input1_addr[i] + input2_addr[i];
  }
  return Status::OK();}

PluginAdd在MLU上运行，需要调用支持MLU运行的CNNL算子或BANGC实现，这里我们用Bangc实现了在MLU执行加法的VectorAdd接口, 在获取到对应输入输出的设备指针后，调用VectorAdd接口完成计算。

Status PluginAddKernel::Enqueue(INodeResource *context) {
  float *input1_addr = nullptr;
  float *input2_addr = nullptr;
  float *output_addr = nullptr;
  cnrtQueue_t queue;
  context->GetTensorDataPtr("input1", static_cast<void **>(&input1_addr));
  context->GetTensorDataPtr("input2", static_cast<void **>(&input2_addr));
  context->GetTensorDataPtr("output", static_cast<void **>(&output_addr));
  context->GetQueue(&queue); // 获取所需要的输入输出地址和Queue  cnrtDim3_t dim{4,1,1};
  cnrtFunctionType_t type = cnrtFuncTypeUnion1;
  VectorAdd<<<dim, type, queue>>>(input1_addr, input2_addr, output_addr, input_count);
  // 调用BangC 实现好的VectorAdd接口完成计算  return Status::OK();}// VectorAdd__mlu_global__ void VectorAdd(const float *input1,  const float *input2,
                              float *output, const uint32_t count) {
  __nram__ float src1_nram[BANG_NRAM_SIZE/8]; // 在nram上分配空间  __nram__ float src2_nram[BANG_NRAM_SIZE/8];
  __nram__ float dst_nram[BANG_NRAM_SIZE/8];
  uint32_t num_per_core = count / taskDim;
  // 将ddr内存上的数据拷贝到nram上  __memcpy(src1_nram, input1 + taskId * num_per_core, num_per_core * sizeof(float), GDRAM2NRAM);
  __memcpy(src2_nram, input2 + taskId * num_per_core, num_per_core * sizeof(float), GDRAM2NRAM);
  __bang_add(dst_nram, src1_nram, src2_nram, num_per_core); // 在mlu上执行加法操作  // 将数据拷出到ddr内存上  __memcpy(output + taskId * num_per_core, dst_nram, num_per_core * sizeof(float), NRAM2GDRAM);}

3 注册 MagicMind 算子 Kernel

MagicMind算子 Kernel 的注册，这一步主要是实现 KernelFactory 以及注册执行设备，以便调用该算子。

在实现完算子 Kernel 之后，还需要对算子 Kernel 进行注册，使得 Runtime 运行时能够识别并调用它。需要注册的信息包括算子 Name、设备类型以及 KernelFactory， 注册Kernel 使用 PLUGIN_REGISTER_KERNEL 进行注册。PluginAdd示例如下

class PluginAddKernelFactory : public IPluginKernelFactory {
 public:
  IPluginKernel* Create() override {  // 重写Create接口    return new PluginAddKernel();
  }
  ~PluginAddKernelFactory() {}};PLUGIN_REGISTER_KERNEL(CreatePluginKernelDefBuilder("PluginAdd").DeviceType("MLU"), PluginAddKernelFactory);

注意：注册 Kernel 的 Name 要和前面注册算子的 Name 保持一致

4 编译 PluginOp动态库

算子 Kernel 注册完之后，还需要编译出一个动态库，供后面使用

如果涉及 BangC Kernel，需要使用 CNCC 将对应的 MLU 文件编译成.o文件，如下所示

cncc --bang-arch=compute_30 -c vectoradd_mlu_kernel.mlu -O3 -o vectoradd_mlu_kernel.o -fPIC

接着将实现好的 Plugin 文件和编译好的.o(如有）编译成动态库；如下所示

gcc --shared -fPIC plugin_add.cc vectoradd_mlu_kernel.o  \
-I /usr/local/neuware/include -o libplugin.so

注：上面 BangC Kernel 的编译不是必须的，取决于 Enqueue接口是调用 BangC实现还是其他（如CNNL算子）

3、自定义算子的使用

自定义算子的使用有两种方式，可以通过API使用，也可以通过parser使用

3.1 通过 API 使用自定义算子

MagicMind INetwork 类提供了 AddIPluginNode 接口，它用来往网络中添加 Plugin 算子，并返回一个 Plugin 算子对象引用，C++示例如下，（调用PluginOp时需要在函数一开始调用dlopen打开对应plugin.so)；如果调用PluginOp和PluginOp的文件同时编译，则不需要用dlopen

//调用PluginOp时需要调用dlopen打开对应plugin.soauto kernel_lib = dlopen("libplugin.so", RTLD_LAZY); //添加网络输入INetwork *network = CreateINetwork();ITensor *input1 = network->AddInput(magicmind::DataType::FLOAT16, Dims({1, 2, 2, 4}));ITensor *input2 = network->AddInput(magicmind::DataType::FLOAT16, Dims({1, 2, 2, 4}));TensorMap plugin_inputs;DataTypeMap plugin_outputs_dtype;//准备PluginOp的输入和输出类型std::vector<ITensor *> VInput1{input1};std::vector<ITensor *> VInput2{input2};plugin_inputs["input1"] = VInput1;plugin_inputs["input2"] = VInput2;plugin_outputs_dtype["output"] = std::vector<DataType>{magicmind::DataType::FLOAT16};// 添加PluginAdd节点IPluginNode *plugin_add = network->AddIPluginNode("PluginAdd", plugin_inputs, plugin_outputs_dtype);// 添加其它算子IMaxPool2DNode *pool = network->AddIMaxPool2DNode(plugin_add->GetOutput(0), true);// …

3.2 通过 Parser 使用自定义算子及限制

通过Parser使用自定义算子，分别支持Pytorch/TF/ONNX，不支持Caffe.

1）使用PyTorch Parser

加载模型之前还需要注册自定义算子接口到 libtorch 中，可通过调用 LoadRegistedPlugin 接口来实现 libtorch 的算子注册，如下示例；

const char* op_name = "MyPlugin::custom_add";  // 算子名称static auto TorchLoadRegistedPlugin = LoadRegistedPlugin<ModelKind::kPytorch>;TorchLoadRegistedPlugin(
    [op_name](const std::string& name) { return name == op_name; }, nullptr);

2）OnnxParser

直接加载模型即可

3）使用 TensorFlow Parser 加载模型

TensorFlow Parser 加载模型之前还需要注册自定义算子到 libtf中，目前TensorFlow Parser 支持调用 LoadRegistedPlugin接口来实现libtf的算子注册。

const char* op_name = "TF_custom_add_op";  // 算子名称static auto TFLoadRegistedPlugin = LoadRegistedPlugin<ModelKind::kTensorFlow>;TFLoadRegistedPlugin(
    [](const std::string& name) { return name == op_name; },
    [](const std::string& name) { return op_name; });

*注册自定义算子名字首字母要大写，后面的字符要属于字母、数字、下划线，否则注册到libtf中会失败。

4、使用自定义算子加速案例

案例：在 MagicMind上测试 Conformer 网络性能时发现该网络时延较大，将网络结构中的 ATTN/FFN 结构分别替换为大算子，时延得到显著的减少，下图为 FFN 结构替换示例

该模型在 SeqLen=2000，BatchSize=1,16,32的时延对比表如下(测试环境：MLU370-X4, 测试精度为 Float32)：

SyntaxHighlighter.all();