CANN Samples（十三）：Ascend C 算子开发入门

在之前的文章中，我们一直在使用 CANN 为我们提供的各种高级接口和工具。但你是否想过，那些构成神经网络基本单元的"算子"（Operator），比如卷积、加法、激活函数等，它们在底层是如何实现的？如果遇到一个框架不支持的、非常新颖的算子，我们又该如何自己动手去创造它？

深入到 CANN 的最底层，学习如何使用 Ascend C 语言来开发自定义算子。这绝对是硬核的一章，但别担心，我会像之前一样，用最通俗易懂的方式，带你一步步揭开算子开发的神秘面纱。

1. 算子开发：两种主流模式

在昇腾（Ascend）平台上，开发一个算子通常有两种主流的模式：Kernel Launch 模式和 Framework Launch 模式。

Kernel Launch 模式 ：可以理解为"直接调用模式"。这种方式非常直接，我们像写普通的 C++ 程序一样，在主机（CPU）侧的代码里，通过一个特殊的调用语法 <<<...>>>，直接启动在昇腾 AI 处理器（我们称之为设备或 NPU）上运行的函数（我们称之为"核函数"或 Kernel）。这种模式非常灵活，代码写起来也更自由，特别适合快速验证算法原型或实现一些高度定制化的逻辑。
Framework Launch 模式：可以理解为"框架集成模式"。这种方式更工程化、更规范。我们把自定义算子按照一套标准的规范进行开发和打包，最终生成一个标准的算子包。这个算子包可以被 CANN 的运行时框架（比如 ACL）加载，也可以被 TensorFlow、PyTorch 等深度学习框架识别和调用，就像使用官方提供的原生算子一样。

简单打个比方，Kernel Launch 模式就像是你自己组装了一台性能强劲的赛车，然后自己下场去开，怎么开、何时开都由你说了算。而 Framework Launch 模式则像是你设计并制造了一款符合 F1 标准的引擎，然后把它卖给各大车队，让他们把你的引擎集成到自己的赛车里去参加比赛。

为了让你对这两种模式都有一个直观的认识，我们将通过 samples/operator/ascendc/ 目录下的官方示例，从易到难，一步步掌握它们。

1.1. Hello World：迈出第一步 (`Kernel Launch`)

任何编程语言的学习都始于一句"Hello, World!"。算子开发也不例外。让我们首先聚焦于 operator/ascendc/0_introduction/0_helloworld 这个最简单的示例，看看 Kernel Launch 模式是如何工作的。

这个示例的目录结构非常简单：

复制代码

.
├── CMakeLists.txt          // 编译脚本
├── hello_world.cpp         // 设备端（NPU）代码，定义核函数
├── main.cpp                // 主机端（CPU）代码，调用核函数
└── run.sh                  // 一键运行脚本

1.1.1. 设备端代码：`hello_world.cpp`

我们先来看 hello_world.cpp，这里的代码将运行在 NPU 上。

cpp 复制代码

#include "kernel_operator.h"

extern "C" __global__ __aicore__ void hello_world()
{
    AscendC::printf("Hello World!!!\n");
}

void hello_world_do(uint32_t blockDim, void *stream)
{
    hello_world<<<blockDim, nullptr, stream>>>();
}

这段代码虽然短小，但信息量巨大：

__global__ 和 __aicore__：这是两个"魔法"般的关键字。__global__ 告诉编译器，这个函数（hello_world）是一个"核函数"，它应该被编译成能在设备（NPU）上运行的代码。__aicore__ 则进一步指定，这个核函数将在 NPU 的 AI Core 上执行。
AscendC::printf：这是 Ascend C 提供的一个内置函数，功能类似标准 C++ 的 printf，但它的特殊之处在于，它能在 NPU 的设备上打印信息，这对于调试核函数非常有帮助。
hello_world_do 函数和 <<<...>>>：hello_world_do 是一个普通的 C++ 函数，它封装了对核函数的调用。最关键的就是 hello_world<<<blockDim, nullptr, stream>>>() 这行代码。这正是 Kernel Launch 模式的核心，这个特殊的语法用于从主机端启动一个核函数。
- blockDim：指定要用多少个 AI Core 来并行执行这个核函数。
- stream：指定在哪个"执行流"上运行。你可以把 stream 理解成一个独立的任务队列。

1.1.2. 主机端代码：`main.cpp`

看完了设备端的"士兵"，我们再来看主机端的"指挥官" main.cpp。

cpp 复制代码

#include "acl/acl.h"
extern void hello_world_do(uint32_t coreDim, void *stream);

int32_t main(int argc, char const *argv[])
{
    // 1. 初始化 AscendCL
    aclInit(nullptr);
    int32_t deviceId = 0;
    aclrtSetDevice(deviceId);
    
    // 2. 创建一个执行流
    aclrtStream stream = nullptr;
    aclrtCreateStream(&stream);
    
    // 3. 启动核函数
    constexpr uint32_t blockDim = 8; // 在 8 个 AI Core 上执行
    hello_world_do(blockDim, stream);
    
    // 4. 等待任务执行完成
    aclrtSynchronizeStream(stream);

    // 5. 释放资源
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}

主机端的代码逻辑非常清晰，遵循了标准的 AscendCL 应用开发流程：

初始化 ：调用 aclInit 和 aclrtSetDevice 来初始化 AscendCL 库并指定在哪块 NPU 设备上工作。
创建资源 ：调用 aclrtCreateStream 创建一个用于执行任务的 stream。
执行任务 ：调用我们之前在 hello_world.cpp 中定义的封装函数 hello_world_do，通过 <<<...>>> 语法将 hello_world 核函数下发到 NPU 上执行。
同步等待 ：aclrtSynchronizeStream 是一个阻塞函数，它会一直等待，直到这个 stream 上的所有任务（包括我们刚刚下发的核函数）都执行完毕。
清理收尾：释放所有申请的资源。

通过这个最简单的例子，我们理清了 Kernel Launch 模式的基本思想：主机端通过 AscendCL API 负责资源管理和任务下发，设备端通过带有 __global__ 标记的核函数负责具体的计算任务。

1.2. 自定义加法算子：处理真实数据

HelloWorld 只是让我们"听个响"，一个真正的算子需要处理输入数据并产生输出。现在，我们把目光投向一个更实用的例子：operator/ascendc/tutorials/AddCustomSample。这个例子实现了一个自定义的向量加法算子 z = x + y。

我们将再次聚焦于 KernelLaunch 模式，看看它是如何处理数据的。相关代码位于 KernelLaunch/AddKernelInvocationNeo 目录下。

1.2.1. 设备端：`add_custom.cpp` 的流水线艺术

与 HelloWorld 不同，add_custom.cpp 的代码引入了算子性能优化的核心思想：流水线（Pipeline）。

在 NPU 中，数据存储在"全局内存"（Global Memory，通常是 DDR），而计算则发生在速度极快但容量很小的"片上内存"（Local Memory）中。因此，一个高效的算子必须精心设计数据在两种内存之间的搬运过程。

add_custom.cpp 将整个计算过程拆分成了三个阶段，形成了一条高效的流水线：

CopyIn ：从 Global Memory 中拷贝一小块输入数据（x 和 y）到 Local Memory。
Compute：在 Local Memory 中执行加法计算。
CopyOut ：将计算结果（z）从 Local Memory 拷贝回 Global Memory。

为了让这条流水线能够不间断地运行，代码中使用了 TPipe 和 TQue 这两个 Ascend C 提供的工具，构建了"双缓冲（Double Buffering）"机制。你可以想象一下，当计算单元正在处理第 N 块数据时，数据搬运单元可以同时去 Global Memory 预取第 N+1 块数据，并将第 N-1 块的计算结果写回。这样，计算和数据传输可以高度重叠，极大地提升了硬件利用率。

让我们看看关键的核函数代码：

cpp 复制代码

// operator/ascendc/tutorials/AddCustomSample/KernelLaunch/AddKernelInvocationNeo/add_custom.cpp

class KernelAdd {
public:
    // ... 初始化 ...
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z) {
        // ... 设置 Global Memory 和 Local Memory 的映射关系 ...
        // ... 初始化用于流水线的 Queue ...
    }

    __aicore__ inline void Process() {
        // 将数据切分成 TILE_NUM * BUFFER_NUM 个小块，循环处理
        int32_t loopCount = TILE_NUM * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);  // 预取下一块数据
            Compute(i); // 计算当前数据
            CopyOut(i); // 写回上一块结果
        }
    }

private:
    // 从 Global Memory 拷贝到 Local Memory
    __aicore__ inline void CopyIn(int32_t progress) { /* ... */ }
    // 在 Local Memory 中执行计算
    __aicore__ inline void Compute(int32_t progress) { /* ... */ }
    // 从 Local Memory 拷贝回 Global Memory
    __aicore__ inline void CopyOut(int32_t progress) { /* ... */ }

private:
    AscendC::TPipe pipe; // 流水线工具
    AscendC::TQue<...> inQueueX, inQueueY; // 输入数据的队列
    AscendC::TQue<...> outQueueZ; // 输出数据的队列
    AscendC::GlobalTensor<half> xGm, yGm, zGm; // Global Memory 上的张量
};

// 核函数入口
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    KernelAdd op;
    op.Init(x, y, z);
    op.Process();
}

1.2.2. 主机端：`main.cpp` 的数据交互

主机端的 main.cpp 现在不仅要负责启动核函数，还要负责准备输入数据，并将计算结果从 NPU 拷贝回来。

cpp 复制代码

// operator/ascendc/tutorials/AddCustomSample/KernelLaunch/AddKernelInvocationNeo/main.cpp

int32_t main(int32_t argc, char *argv[])
{
    // ... 初始化 ACL ...

    // 1. 在 Host 和 Device 上都分配内存
    uint8_t *xHost, *yHost, *zHost; // Host 侧内存
    uint8_t *xDevice, *yDevice, *zDevice; // Device 侧内存
    aclrtMallocHost((void **)(&xHost), inputByteSize);
    // ... (为 yHost, zHost 分配)
    aclrtMalloc((void **)&xDevice, inputByteSize, ...);
    // ... (为 yDevice, zDevice 分配)

    // 2. 准备输入数据
    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);
    ReadFile("./input/input_y.bin", inputByteSize, yHost, inputByteSize);

    // 3. 将输入数据从 Host 拷贝到 Device
    aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(yDevice, inputByteSize, yHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE);

    // 4. 启动核函数，注意这次传入了 Device 上的内存地址
    add_custom_do(blockDim, stream, xDevice, yDevice, zDevice);
    aclrtSynchronizeStream(stream);

    // 5. 将计算结果从 Device 拷贝回 Host
    aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST);
    
    // 6. 保存结果并验证
    WriteFile("./output/output_z.bin", zHost, outputByteSize);

    // ... 释放所有内存和资源 ...
    return 0;
}

这个流程非常清晰地展示了主机与设备之间的数据交互：Host 准备数据 -> 拷贝到 Device -> Device 执行计算 -> 结果拷贝回 Host -> Host 验证结果。

1.3. 算子工程化：`Framework Launch` 模式

Kernel Launch 模式虽然灵活，但它更像是一个"个人作坊"。如果想让我们的算子被 AI 框架这样的"正规军"使用，就需要采用 Framework Launch 模式，对算子进行标准化、工程化的封装。

我们回到 AddCustomSample 示例，这次看 FrameworkLaunch/AddCustom 目录。

1.3.1. 角色分工：主机端与设备端的职责分离

Framework Launch 模式下，代码结构发生了显著变化，职责划分更加清晰：

op_kernel/add_custom.cpp：设备端实现 。这部分代码和 Kernel Launch 模式下的核函数几乎一样，专注于在 NPU 上执行高性能的计算流水线。
op_host/add_custom.cpp：主机端实现 。这不再是一个 main 函数，而是定义了算子与框架交互的"规则"。
AddCustom.json：算子定义文件。这是一个 JSON 文件，像一张"身份证"，向外界声明了算子的接口信息。

1.3.2. 算子的"身份证"：`AddCustom.json`

json 复制代码

[
    {
        "op": "AddCustom",
        "language": "cpp",
        "input_desc": [
            { "name": "x", "param_type": "required", "format": ["ND"], "type": ["float16"] },
            { "name": "y", "param_type": "required", "format": ["ND"], "type": ["float16"] }
        ],
        "output_desc": [
            { "name": "z", "param_type": "required", "format": ["ND"], "type": ["float16"] }
        ]
    }
]

这个文件清晰地定义了 AddCustom 算子需要两个名为 x 和 y 的 float16 类型的输入，并会产生一个名为 z 的 float16 类型的输出。框架通过读取这个文件，就能了解算子的基本信息。

1.3.3. 主机端的"大脑"：`op_host/add_custom.cpp`

主机端的代码是 Framework Launch 模式的精髓所在。它定义了三件重要的事情：

算子注册 (ops::AddCustom)：通过一个 C++ 类，将算子的接口信息（输入、输出、数据类型等）以编程的方式注册到 CANN 框架中。
形状推导 (InferShape)：提供一个函数，告诉框架如何根据输入的形状来推断输出的形状。这对于 AI 框架进行计算图优化至关重要。
Tiling 策略 (TilingFunc)：这是最核心的部分。它是一个回调函数，在算子执行前被框架调用。它会根据当前任务的实际输入数据大小，动态地计算出最佳的数据切分和并行方案（比如用几个 Core，每个 Core 分多少数据），然后将这个方案打包，传递给设备端的核函数去执行。

cpp 复制代码

// operator/ascendc/tutorials/AddCustomSample/FrameworkLaunch/AddCustom/op_host/add_custom.cpp

namespace optiling {
// Tiling 函数：动态计算并行和切分策略
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{
    // 1. 获取输入数据的总长度
    uint32_t totalLength = context->GetInputShape(0)->GetOriginShape().GetShapeSize();
    // 2. 设置并行执行的 Core 的数量
    context->SetBlockDim(BLOCK_DIM);
    // 3. 将 totalLength 和 tileNum 等信息打包到 tiling data 中
    tiling.SaveToBuffer(...);
    // 4. 将打包好的 tiling data 交给框架
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
// 算子定义和注册
class AddCustom : public OpDef {
public:
    explicit AddCustom(const char *name) : OpDef(name)
    {
        // 定义输入、输出
        this->Input("x")...;
        this->Input("y")...;
        this->Output("z")...;

        // 注册形状推导函数和 Tiling 函数
        this->SetInferShape(ge::InferShape);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend910"); // 声明支持的芯片型号
    }
};
OP_ADD(AddCustom); // 将算子注册到系统中
}

通过这种方式，算子的计算逻辑（Kernel）和调度策略（Host）被完全解耦。主机端的代码像一个聪明的"大脑"，为设备端的"士兵"制定作战计划，使得算子能够被框架智能地调度，并高效地处理各种不同尺寸的数据。

总结

今天，我们踏上了一段深入 CANN 底层的旅程，学习了如何使用 Ascend C 开发自定义算子。

我们从 Kernel Launch 模式的 HelloWorld 开始，理解了主机端与设备端协同工作的基本模型。
接着，通过 AddCustom 示例，我们学习了如何设计带有数据输入输出的算子，并利用流水线和双缓冲技术来优化性能。
最后，我们探索了 Framework Launch 模式，了解了如何将算子进行标准化、工程化的封装，通过定义算子接口、形状推导和动态 Tiling 策略，让算子能够无缝地融入主流 AI 框架。

掌握算子开发，意味着你拥有了突破框架限制、实现极致性能优化的能力。虽然这部分内容比之前的章节更具挑战性，但它也为你打开了一扇通往高性能计算新世界的大门。希望这篇文章能成为你探索这个新世界的起点。

CANN Samples（十三）：Ascend C 算子开发入门

1. 算子开发：两种主流模式

1.1. Hello World：迈出第一步 (Kernel Launch)

1.1.1. 设备端代码：hello_world.cpp

1.1.2. 主机端代码：main.cpp

1.2. 自定义加法算子：处理真实数据

1.2.1. 设备端：add_custom.cpp 的流水线艺术

1.2.2. 主机端：main.cpp 的数据交互

1.3. 算子工程化：Framework Launch 模式

1.3.1. 角色分工：主机端与设备端的职责分离

1.3.2. 算子的"身份证"：AddCustom.json

1.3.3. 主机端的"大脑"：op_host/add_custom.cpp

总结

1.1. Hello World：迈出第一步 (`Kernel Launch`)

1.1.1. 设备端代码：`hello_world.cpp`

1.1.2. 主机端代码：`main.cpp`

1.2.1. 设备端：`add_custom.cpp` 的流水线艺术

1.2.2. 主机端：`main.cpp` 的数据交互

1.3. 算子工程化：`Framework Launch` 模式

1.3.2. 算子的"身份证"：`AddCustom.json`

1.3.3. 主机端的"大脑"：`op_host/add_custom.cpp`