1. 简介

AscendC算子在经过编译后，会生成一个aclnn接口的算子包。本文以一个custom_silu的自定义激活函数算子为例，讲述如何将Pytorch的Aten算子层与aclnn接口绑定，并在Python侧接口成功调用AscendC算子。

2. 安装前置

本文以pytorch的v2.1.0版本为例子。pip安装命令如下。

pip install torch==2.1.0

然后下载华为官方的算子注册仓库，本文是以在op-plugin上修改的方式支持算子。

git clone git@gitee.com:ascend/op-plugin.git

3. 注册算子代码

3.1 算子yaml注册

找到算子注册的yaml目录[op_plugin_functions.yaml]
写入这样的配置。

- func: custom_silu(Tensor self, float k) -> Tensor
    op_api: all_version

func是函数的开头标签，后面写入Pytorch的算子表达式。
表达式格式：函数名字(函数参数) -> 返回类型。
op_api是表示你希望算子放置的命名空间。op_plugin一共定义了两种命名空间：aclop和op_api。op_api使用EXEC_NPU_CMD调用aclnn算子。aclop使用OpCommand调用算子，使用aclopCompileAndExecute执行算子。
自定义算子包，都是基于aclnn封装的，所以使用op_api命名空间。

3.2 cpp代码接口实现

yaml编写之后，torchgen模块会生成算子接口定义的头文件。

namespace op_api {
    at::Tensor custom_silu(const at::Tensor &self, float k);
}

在op_plugin/ops/opapi文件夹下创建一个cpp新文件，例如Custo
SiluKernel.cpp。

#include "op_plugin/AclOpsInterface.h"
#include "op_plugin/OpApiInterface.h"
#include "op_plugin/utils/op_api_common.h"

namespace op_api {
using npu_preparation = at_npu::native::OpPreparation;

at::Tensor custom_silu(const at::Tensor &x, float k)
{
    auto options = x.options();
    at::Tensor result = npu_preparation::apply_tensor_without_format(x.sizes(), options); // 创建输出内存

    // calculate the output result of the NPU
    EXEC_NPU_CMD(aclnnCustomSilu, x, y, result);
    return result;
}
} // namespace op_api

npu_preparation::apply_tensor_without_format是一种快速创建Tensor的封装函数。输出Tensor与输入Tensor的大小和格式保持一致。

aclnnCustomSilu是本文已经实现的aclnn算子包。EXEC_NPU_CMD是封装了动态调用aclnnCustomSiluGetWorkSize的宏，其实现在op-plugin/op_plugin/utils/op_api_common.h。

EXEC_NPU_CMD宏还会自动把输入的Pytorch at::Tensor转换成昇腾的aclTensor，极大地简化了代码。

4. Python接口测试

代码完成后，执行一键编译命令和安装。

bash ci/build.sh --python=3.9 --pytorch=v2.1.0
pip install dist/*.whl

测试你的自定义算子接口。

import torch
import torch_npu

a = torch.rand([1, 100])
b = torch_npu.custom_silu(a, 0.5)

print(b)