5.Swish 激活函数算子
Swish 激活函数算子
大致代码参考官方教程即可快速入门-自定义算子开发-AscendC算子开发-CANN - 华为HarmonyOS开发者
一、创建工程
在~/mywork目录创建json文件：swish_custom.json
[
{
"op": "SwishCustom",
"input_desc": [
{
"name": "x",
"param_type": "required",
"format": [
"ND",
"ND",
"ND"
],
"type": [
"fp16",
"float",
"int32"
]
}
],
"output_desc": [
{
"name": "z",
"param_type": "required",
"format": [
"ND",
"ND",
"ND"
],
"type": [
"fp16",
"float",
"int32"
]
}
]
}
]
使用msopgen命令创建工程：
msopgen gen -i ~/mywork/swish_custom.json -c ai_core-kirin9020 -out ~/mywork/SwishCustom
二、算子实现
代码同官方教程中的AddCustom算子基本一样，除了名字之外，只需要修改kernel侧代码（在官方教程的基础上修改变量、compute函数）：
#include "kernel_operator.h"
constexpr int32_t BUFFER_NUM = 2;
class KernelSwish {
public:
__aicore__ inline KernelSwish() {}
// 初始化函数，完成内存初始化相关操作
__aicore__ inline void Init(GM_ADDR x, GM_ADDR z, uint32_t totalLength, uint32_t tileNum)
{
// 使用获取到的TilingData计算得到singleCoreSize(每个核上总计算数据大小)、tileNum（每个核上分块个数）、singleTileLength（每个分块大小）等变量
this->blockLength = totalLength / AscendC::GetBlockNum();
this->tileNum = tileNum;
this->tileLength = this->blockLength / tileNum / BUFFER_NUM;
// 获取当前核的起始索引
xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
zGm.SetGlobalBuffer((__gm__ DTYPE_Z*)z + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
// 通过Pipe内存管理对象为输入输出Queue分配内存
pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Z));
}
// 核心处理函数，实现算子逻辑，调用私有成员函数CopyIn、Compute、CopyOut完成矢量算子的三级流水操作
__aicore__ inline void Process()
{
int32_t loopCount = this->tileNum * BUFFER_NUM;
for (int32_t i = 0; i < loopCount; i++) {
CopyIn(i);
Compute(i);
CopyOut(i);
}
}
private:
// 搬入函数，完成CopyIn阶段的处理，被核心Process函数调用
__aicore__ inline void CopyIn(int32_t progress)
{
// 从Queue中分配输入Tensor
AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
// 将GlobalTensor数据拷贝到LocalTensor
AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
// 将LocalTensor放入VECIN（代表矢量编程中搬入数据的逻辑存放位置）的Queue中
inQueueX.EnQue(xLocal);
}
// 计算函数，完成Compute阶段的处理，被核心Process函数调用
__aicore__ inline void Compute(int32_t progress)
{
// 将Tensor从队列中取出，用于后续计算
AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
// 从输出队列分配临时 Tensor
AscendC::LocalTensor<DTYPE_Z> expNegX = outQueueZ.AllocTensor<DTYPE_Z>();
// 1. -x
AscendC::Muls(expNegX, xLocal, (DTYPE_X)-1.0, this->tileLength);
// 2. exp(-x)
AscendC::Exp(expNegX, expNegX, this->tileLength);
// 3. 1 + exp(-x) 使用 Adds 直接加标量 1
AscendC::Adds(expNegX, expNegX, (DTYPE_Z)1.0, this->tileLength);
// 4. 倒数得到 sigmoid(x)
AscendC::Reciprocal(expNegX, expNegX, this->tileLength);
// 5. z = x * sigmoid(x)
AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();
AscendC::Mul(zLocal, xLocal, expNegX, this->tileLength);
outQueueZ.EnQue(zLocal);
outQueueZ.FreeTensor(expNegX);
inQueueX.FreeTensor(xLocal);
}
// 搬出函数，完成CopyOut阶段的处理，被核心Process函数调用
__aicore__ inline void CopyOut(int32_t progress)
{
// 从VecOut的Queue中取出输出Tensor
AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
// 将输出Tensor拷贝到GlobalTensor中
AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
// 将不再使用的LocalTensor释放
outQueueZ.FreeTensor(zLocal);
}
private:
// Pipe内存管理对象
AscendC::TPipe pipe;
// 输入数据Queue队列管理对象，QuePosition为VECIN
AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX;
// 输出数据Queue队列管理对象，QuePosition为VECOUT
AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
// 管理输入输出Global Memory内存地址的对象，其中xGm, 为输入，zGm为输出
AscendC::GlobalTensor<DTYPE_X> xGm;
AscendC::GlobalTensor<DTYPE_Z> zGm;
// 每个核上总计算数据大小
uint32_t blockLength;
// 每个核上总计算数据分块个数
uint32_t tileNum;
// 每个分块大小
uint32_t tileLength;
};
extern "C" __global__ __aicore__ void swish_custom(GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
GET_TILING_DATA(tiling_data, tiling);
KernelSwish op;
op.Init(x, z, tiling_data.totalLength, tiling_data.tileNum);
op.Process();
}
然后编译
./build.sh
三、运行测试
生成测试数据：
import numpy as np
# 生成随机输入，例如形状 (256,)，数据类型 float16
x = np.random.randn(256).astype(np.float16)
x.tofile('swish_x.bin')
# 计算 golden
def swish(x):
return x / (1 + np.exp(-x))
golden = swish(x).astype(np.float16)
golden.tofile('swish_golden.bin')
print("数据已生成并保存到 swish_x.bin 和 swish_golden.bin")
创建json文件：swish_config.json
{
"op_type": "SwishCustom",
"gen_data": false,
"inputs": [
{
"name": "x",
"dtype": "float16",
"format": "ND",
"shape": [256],
"param_type": "required",
"data_file": "/home/dj/mywork/SwishCustom/swish_x.bin"
}
],
"outputs": [
{
"name": "z",
"dtype": "float16",
"format": "ND",
"shape": [256],
"param_type": "required",
"data_file": "/home/dj/mywork/SwishCustom/swish_golden.bin"
}
]
}
CPU测试命令：
ascendebug kernel \
--backend cpu \
--chip-version kirin9020 \
--repo-type customize \
--json-file ./swish_config.json \
--core-type AiCore \
--work-dir ./debug_workspace
仿真测试命令：
ascendebug kernel --backend simulator --repo-type customize --json-file ./swish_config.json --core-type AiCore --chip-version kirin9020 --work-dir ./debug_workspace --block-num 1 --timeout 1200
浙公网安备 33010602011771号