核函数(Kernel)是在设备侧并行执行的C/C++函数。它由主机端代码调用,调用时被实例化为多份并行副本,分别调度到多个AI Core上运行(每个AI Core被抽象为一个Block)。这些Block利用AI Core的计算能力,协同完成某一算子的计算。
Kirin9050/Kirin9020/Kirin9030/KirinX90系列处理器暂不支持通过<<<>>>调用核函数。
extern "C" __global__ __aicore__ void add_custom(__gm__ uint8_t* x, __gm__ uint8_t* y, __gm__ uint8_t* z)
以上是一个核函数声明的例子,编写核函数时需要遵循以下规则。
使用函数类型限定符
除了需要按照C/C++函数声明的方式定义核函数之外,还要为核函数加上额外的函数类型限定符,包含__global__和__aicore__。
使用__global__函数类型限定符来标识它是一个核函数,可以被<<<>>>调用;使用__aicore__函数类型限定符来标识该核函数在设备端AI Core上执行:__global__ __aicore__ void kernel_name(argument list)。
编程中使用到的函数可以分为三类:核函数(device侧执行)、host侧执行函数、device侧执行函数(除核函数之外的)。三者的调用关系如下图所示:
图1 核函数(device侧执行)、host侧执行函数、device侧执行函数(除核函数之外的)调用关系

使用变量类型限定符
指针入参变量需要增加变量类型限定符__gm__,表明该指针变量指向Global Memory上某处内存地址。
其他规则或建议
规则1:核函数必须具有void返回类型。
规则2:仅支持入参为指针或C/C++内置数据类型(Primitive data types),如:half* s0、float* s1、int32_t c。
建议:为了统一表达,建议使用GM_ADDR宏来修饰入参,GM_ADDR宏定义如下。
#define GM_ADDR __gm__ uint8_t*
使用GM_ADDR修饰入参的样例如下。
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
这里统一使用uint8_t类型的指针,在后续的使用中需要将其转化为实际的指针类型。
核函数的调用语句是C/C++函数调用语句的一种扩展。本节仅描述最基础的调用方式,实际进行算子开发时,开发者可以通过该方式调用,也可以通过接入AscendC框架的方式调用,更多细节请参考工程化算子开发。
常见的函数调用方式是如下的形式:
function_name(argument list);
核函数使用内核调用符<<<...>>>这种语法形式,来规定核函数的执行配置:
kernel_name<<<blockDim, l2ctrl, stream>>>(argument list);
内核调用符仅可在NPU侧编译时调用,CPU侧编译无法识别该符号。
执行配置由以下blockDim, l2ctrl, stream参数决定:
规定了核函数将会在几个核上执行。每个执行该核函数的核会被分配一个逻辑ID,即block_idx,可以在核函数的实现中调用GetBlockIdx来获取block_idx。
blockDim是逻辑核的概念,取值范围为[1, 65535]。为了充分利用硬件资源,一般设置为物理核的核数或其倍数。
如下名为add_custom的核函数,实现两个矢量的相加,调用示例如下。
// blockDim设置为8表示在8个核上调用了add_custom核函数,每个核都会独立且并行地执行该核函数,该核函数的参数列表为x,y,z。 add_custom<<<8, nullptr, stream>>>(x, y, z);
核函数的调用是异步的,核函数的调用结束后,控制权立刻返回给host端,可以调用以下aclrtSynchronizeStream函数来强制主机端程序等待所有核函数执行完毕。
aclError aclrtSynchronizeStream(aclrtStream stream);
下面提供核函数实现的样例代码片段,完整样例请参考矢量编程。
// 实现核函数
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();
}