深入解析华为Ascend CANN算子开发：核函数的定义与调用

深入解析华为Ascend CANN算子开发：核函数的定义与调用

随着人工智能计算的不断发展，算子（Operator）作为深度学习模型底层计算单元的核心组成部分，其性能直接影响模型推理效率。在华为Ascend平台上，CANN（Compute Architecture for Neural Networks）提供了完整的算子开发框架，使开发者能够充分利用Ascend AI处理器的硬件优势。在算子开发中，核函数（Kernel Function）是实现设备端计算逻辑的关键入口。本文将从核函数的定义、调用规则、执行配置以及模板核函数使用等方面进行系统解析，帮助开发者快速掌握Ascend算子开发技巧。

训练营简介

2025年昇腾CANN训练营第二季，基于CANN开源开放全场景，推出0基础入门系列、码力全开特辑、开发者案例等专题课程，助力不同阶段开发者快速提升算子开发技能。获得Ascend C算子中级认证，即可领取精美证书，完成社区任务更有机会赢取华为手机，平板、开发板等大奖。

报名链接：https://www.hiascend.com/developer/activities/cann20252#cann-camp-2502-intro

一、核函数概述

核函数是CANN算子在设备端执行的核心入口。与普通C/C++函数不同，核函数在调用时会被多个计算核（Core）并行执行，每个核执行相同的代码，并拥有相同的输入参数。通过核函数，开发者可以直接操作Global Memory，实现数据搬运和计算操作，从而完成算子的全部功能。

核函数的定义必须遵循以下关键要求：

1. 函数类型限定符：必须同时使用 __global__ 和 __aicore__，分别标识该函数为核函数并在AI Core上执行。
2. 入参限定符：指针类型的参数必须使用 __gm__ 修饰，表示指向全局内存（Global Memory）。
3. 返回类型：核函数必须为 void 类型。
4. 参数类型：支持C/C++原生数据类型或指针类型，不支持自定义复杂数据类型。

以Add算子为例，一个典型的核函数实现如下：

extern "C" __global__ __aicore__ void add_custom(__gm__ uint8_t* x, __gm__ uint8_t* y, __gm__ uint8_t* z) {
    KernelAdd op;
    op.Init(x, y, z);
    op.Process();
}

这里的 KernelAdd 类封装了算子的初始化和核心计算逻辑，而核函数本身则负责将其绑定到设备端的多个核并行执行。

二、核函数调用机制

在CANN中，核函数的调用方式是对C/C++函数调用的扩展，采用内核调用符 <<<...>>> 指定执行配置。基本语法如下：

kernel_name<<<blockDim, l2ctrl, stream>>>(argument list);

1. 执行参数说明

• blockDim：指定启动多少个逻辑核执行核函数，每个逻辑核对应一个独立的执行实例。实际应用中，blockDim通常设置为物理核数或其倍数，以充分利用硬件资源。
• l2ctrl：保留参数，目前无需关注，通常设置为 nullptr。
• stream：类型为 aclrtStream，用于控制核函数异步执行的顺序，确保在设备端按主机端调用顺序执行。

调用示例：

add_custom<<<8, nullptr, stream>>>(x, y, z);
aclError aclrtSynchronizeStream(stream); // 等待所有核函数执行完成

核函数调用是异步的，调用完成后控制权立即返回主机端。开发者可通过 aclrtSynchronizeStream 强制等待所有核函数完成计算。

三、变量类型限定符与内存访问

核函数的入参指针必须使用 __gm__ 修饰，表明其指向全局内存。为了统一表达，CANN推荐使用宏 GM_ADDR 来修饰指针：

#define GM_ADDR __gm__ uint8_t*
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z);

在实际开发中，指针类型通常为 uint8_t*，在使用时再根据具体数据类型进行类型转换。这样的设计有助于算子在不同数据类型下保持通用性。

四、核函数执行配置细节

核函数的执行配置主要受blockDim控制，其设置规则因计算模式而异：

1. 耦合模式：Vector和Cube单元集成在同一个AI Core上，blockDim直接指定启动核实例数，不区分Vector或Cube。

2. 分离模式：

   • Vector算子：blockDim指定启动的Vector核实例数量。
   • Cube算子：blockDim指定启动的Cube核实例数量。
   • 融合算子：同时包含Vector和Cube计算，blockDim设置组合数，例如一个组合为2个Vector核和1个Cube核。

此外，开发者可通过 GetCoreNumAic、GetCoreNumAiv 接口获取AI Core核数量，以避免超过硬件资源。

五、模板核函数的使用

CANN支持模板核函数，可定义非类型模板参数和类型模板参数，实现算子逻辑的高度复用。例如：

template<int a, typename T>
__global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z) {
    AscendC::printf("Print Template a: %d\n", a);
    xGm.SetGlobalBuffer((__gm__T*)x + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
    yGm.SetGlobalBuffer((__gm__T*)y + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
    zGm.SetGlobalBuffer((__gm__T*)z + BLOCK_LENGTH * AscendC::GetBlockIdx(), BLOCK_LENGTH);
}

调用方式如下：

add_custom<20, float><<<blockDim, nullptr, stream>>>(x, y, z);

注意事项：

• 模板核函数仅支持 <<<...>>> 调用方式。
• 暂不支持自定义复杂数据类型。

模板核函数的使用可以显著提升算子开发的灵活性和可复用性，使得算子在不同数据类型和参数下无需重复实现。

六、总结

核函数是华为Ascend CANN算子开发中最核心的部分，它承载了算子在设备端的计算逻辑。理解核函数的定义规则、调用方式、执行配置以及模板使用方法，是高效开发算子的基础。通过核函数，开发者能够充分利用Ascend AI处理器的并行计算能力，实现高性能的算子计算。

未来，在掌握核函数基础上，开发者还可以结合Kernel直调算子和工程化算子开发方式，进一步优化算子性能，实现复杂模型在Ascend平台上的高效部署。