昇腾社区首页
中文
注册
开发者
下载
支持
积分兑换 NEW

算子NPU实现（Ascend C）

自定义PyTorch算子的NPU实现一般采用Ascend C编程语言实现，推荐“工程化方式”实现算子，其开发的NPU算子简称为“Ascend C算子”。

工程化开发是标准的算子开发流程，其简化了NPU适配过程，同时会自动注册Ascend C算子对应的Ascend IR，以保证PyTorch算子能与TorchAir max-autotune模式（Ascend IR）配合工作。

使用TorchAir reduce-overhead模式时，您可以使用任意方式开发算子NPU实现，只需要保证Torch算子能在Eager模式下正常工作。

使用TorchAir max-autotune模式时，需要确保算子NPU实现是以Ascend C算子工程化方式开发的aclnnXxx接口，以Kernel直调方式开发的NPU实现不会生成Ascend IR注册逻辑，没有对应的Ascend IR，无法完成PyTorch算子到Ascend IR的转换。

本章仅提供基于Ascend C工程化开发算子的关键步骤说明，详细的操作请参考《CANN Ascend C算子开发指南》中的“算子实现>工程化算子开发”章节，例如算子原型json文件的参数含义、msOpGen工具的命令说明等。

  1. 创建自定义算子工程与原型
  2. 实现Kernel与Tiling
  3. 实现InferShape与InferDataType（可选）
  4. 自定义算子包编译部署

创建自定义算子工程与原型

  1. 编写算子的原型定义json文件，用于生成算子开发工程。

    算子定义时，注意名称为Torch算子名的大驼峰格式，同时入参顺序与类型应当与PyTorch算子原型完全一致。

    例如MyOp算子的原型json文件名为my_op.json，文件内容如下：

     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
    [
    {
        "op": "MyOp",
        "input_desc": [
            {
                "name": "x",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float"]
            },
            {
                "name": "y",
                "param_type": "optional",
                "format": ["ND"],
                "type": ["float"]
            },
            {
                "name": "z",
                "param_type": "dynamic",
                "format": ["ND"],
                "type": ["float"]
            }
        ],
        "attr": [
            {
                "name": "attr1",
                "param_type": "required",
                "type": "float"
            },
            {
                "name": "attr2",
                "param_type": "required",
                "type": "int"
            }
        ],
        "output_desc": [
            {
                "name": "out",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float"]
            }
        ]
    }
]
  2. 使用msOpGen工具生成算子的开发工程。
    ${INSTALL_DIR}/python/site-packages/bin/msopgen gen -i my_op.json -c ai_core-<soc_version> -f pytorch -lan cpp -out ./MyOp
    • ${INSTALL_DIR}为CANN软件安装后文件存储路径，请根据实际环境进行替换。
    • -i：指定算子原型定义的json文件所在路径，请根据实际情况修改。
    • -c：ai_core-<soc_version>代表算子在AI Core上执行，<soc_version>为昇腾AI处理器的型号，请与实际环境保持一致。
    • -lan：参数cpp代表算子基于Ascend C编程框架，使用C/C++编程语言开发。
    • -out：生成文件所在路径，可配置为绝对路径或者相对路径，并且工具执行用户对路径具有可读写权限。若不配置，则默认生成在执行命令的当前路径。
  3. 生成的算子核心工程目录结构如下：
    MyOp
├── build.sh                   // 算子包编译脚本
├── ......
├── op_host                    // Host侧实现
│  ├── my_op.cpp              // 算子定义、Tiling、InferShape、InferDataType实现文件
│  └── my_op_tiling.h         // 算子Tiling定义文件
└── op_kernel                  // Kernel侧实现
    └── my_op.cpp              // Kernel代码实现文件

实现Kernel与Tiling

为方便演示，本样例直接使用默认生成的Kernel和Tiling空实现，不影响后续的编译与执行。

实际业务场景下，您可以参考《CANN Ascend C算子开发指南》中的“算子实现>工程化算子开发”章节下“Kernel侧算子实现”和“Host侧Tiling实现”，进行核心代码开发。

实现InferShape与InferDataType（可选）

本步骤为可选操作，仅当需要使用TorchAir max-autotune模式（Ascend IR）提供的能力时，才需要实现InferShape与InferDataType。

前文msOpGen生成的自定义算子工程会在“op_host/my_op.cpp”中生成一份简单但通常合适的默认实现。默认实现假设输出的shape和data type与第一个输入完全相同，与MyOp的定义一致，无需重新实现。

实际业务场景下，您可以参考《CANN Ascend C算子开发指南》中“算子入图（GE图）开发”章节下“开发流程”完成datatype推导与shape推导。

自定义算子包编译部署

  1. 在自定义算子工程目录下执行如下命令，进行算子工程编译。
    cd ./MyOp
bash build.sh

    编译完成后，会在当前目录下创建build_out目录，并在build_out目录下生成自定义算子安装包custom_opp_<target os>_<target architecture>.run。

    注意，自定义算子包的默认vendor名为customize，相同vendor名称的算子包会互相覆盖。

  2. 自定义算子包安装。
    bash build_out/custom_opp_<target os>_<target architecture>.run
cd ..