核函数

SPMD模型可以得知,使用Ascend C进行编程时,我们编写一份算子实现代码,算子被调用时,将启动N个运行示例,在N个核上运行。本节将介绍算子实现的入口函数。

核函数(Kernel Function)是Ascend C算子设备侧实现的入口。在核函数中,需要为在一个核上执行的代码规定要进行的数据访问和计算操作,当核函数被调用时,多个核都执行相同的核函数代码,具有相同的参数,并行执行。

Ascend C允许用户使用核函数这种C/C++函数的语法扩展来管理设备端的运行代码,用户在核函数中进行算子类对象的创建和其成员函数的调用,由此实现该算子的所有功能。核函数是主机端和设备端连接的桥梁,本章将具体介绍核函数的用法。

核函数定义

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

以上是一个核函数声明的例子,编写核函数时需要遵循以下规则。

核函数调用

核函数的调用语句是C/C++函数调用语句的一种扩展。

常见的函数调用方式是如下的形式:

function_name(argument list);

核函数使用内核调用符<<<...>>>这种语法形式,来规定核函数的执行配置:

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

内核调用符仅可在NPU侧编译时调用,CPU侧编译无法识别该符号。

执行配置由3个参数决定:
  • blockDim,规定了核函数将会在几个核上执行。每个执行该核函数的核会被分配一个逻辑ID,即block_idx,可以在核函数的实现中调用GetBlockIdx来获取block_idx;
  • l2ctrl,保留参数,暂时设置为固定值nullptr,开发者无需关注;
  • stream,类型为aclrtStream,stream用于维护一些异步操作的执行顺序,确保按照应用程序中的代码调用顺序在Device上执行。stream创建等管理接口请参考Stream管理AscendCL应用开发指南(C&C++)中的"AscendCL API参考 > 运行时管理 > Stream管理"章节。

如下名为add_custom的核函数,实现两个矢量的相加,调用示例如下:

// blockDim设置为8表示在8个核上调用了add_custom核函数,每个核都会独立且并行地执行该核函数,该核函数的参数列表为x,y,z。
add_custom<<<8, nullptr, stream>>>(x, y, z);

核函数的调用是异步的,核函数的调用结束后,控制权立刻返回给主机端,可以调用以下aclrtSynchronizeStream函数来强制主机端程序等待所有核函数执行完毕。

aclError aclrtSynchronizeStream(aclrtStream stream);

核函数示例

下面提供核函数实现和调用的样例代码片段,完整样例请参考Add算子示例

// 实现核函数
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();
}

// 调用核函数
void add_custom_do(uint32_t blockDim, void* l2ctrl, void* stream, uint8_t* x, uint8_t* y, uint8_t* z)
{
    add_custom<<<blockDim, l2ctrl, stream>>>(x, y, z);
}