cann/cann-launch-camp: FastGelu自定义算子实现

cann/cann-launch-camp: FastGelu自定义算子实现 FastGeluCustom Ascend C 自定义算子【免费下载链接】cann-launch-camp用户可借助此项目规范管理 CANN 开源社区启航营高校活动的课程作业、课设、毕设等实践作品其核心功能是统一标准化目录层级与提交规范保障作品提交规整、可追溯、可评审。项目地址: https://gitcode.com/cann/cann-launch-camp1. 项目简介本项目在昇腾 NPU 上使用 Ascend C 实现自定义FastGelu算子面向float16和float32输入输出 shape 与 dtype 均与输入保持一致。算子采用逐元素计算方式支持 ND 格式输入可覆盖 1D、2D、3D、4D 及更高维张量。当前实现重点兼顾输出正确性float16精度中大规模张量吞吐小规模张量启动开销控制代码结构简洁便于提交和维护。2. 算子信息项目内容算子名称FastGelu输入x输出y输入 dtypefloat16,float32输出 dtype与输入一致输入 formatND输出 formatND目标硬件ascend910bworkspace03. 数学公式题目给出的 FastGelu 形式为y x * exp(0.851 * (x - abs(x))) / (1 exp(-1.702 * abs(x)))本实现采用等价化简形式y x / (1 exp(-1.702 * x))该形式可以减少算子内部的向量计算链核心计算只需要tmp -1.702 * x tmp exp(tmp) tmp tmp 1 y x / tmp与稳定原式相比该版本少用了Abs、Sub、第二次Exp以及额外Mul因此计算链更短适合追求运行速率的测试场景。4. 目录结构FastGeluCustom/ ├── CMakeLists.txt ├── README.md ├── op_host/ │ ├── CMakeLists.txt │ └── fast_gelu.cpp └── op_kernel/ ├── CMakeLists.txt ├── fast_gelu.cpp ├── fast_gelu_tiling.h └── tiling_key_fast_gelu.h各文件作用文件作用CMakeLists.txt工程总入口设置ascend910b添加 Host 与 Kernel 子目录op_host/fast_gelu.cppHost 侧算子注册、shape 推导、dtype 推导、tiling 计算op_kernel/fast_gelu.cppKernel 侧 Ascend C 计算逻辑op_kernel/fast_gelu_tiling.hHost 与 Kernel 共享的 tiling 数据结构op_kernel/tiling_key_fast_gelu.hdtype 模板 tiling key生成float16/float32两套 Kernel5. 实现特点5.1 模板 TilingKey 区分 dtypeHost 侧通过ASCENDC_TPL_SEL_PARAM(context, DT_X);Kernel 侧通过template typename DT_X __global__ __aicore__ void fast_gelu(...)让平台在编译期生成float16和float32两个 Kernel 实例。这样可以避免运行期 dtype 分支使if constexpr在编译期裁剪无关路径。5.2 half 输入使用 float 中间计算float16分支先将输入 cast 到float32Cast(xf, xLocal, RoundMode::CAST_NONE, len);之后在 float 中完成Muls、Exp、Adds、Div最后 cast 回 half。这样比直接 half 计算更稳尤其对隐藏点中的小数值、负数和精度边界更有利。5.3 按最小元素数控制启动核数Host 侧使用MIN_ELEMENTS_PER_CORE 64当张量较小时不会盲目启动所有 AIV 核减少小规模输入下的调度和流水初始化开销。5.4 blockLength 按 32B 对齐Host 侧根据 dtype 计算alignNum 32 / typeSize blockLength AlignUp(ceil(length / blockDim), alignNum)这样每个核处理的起始位置更规整有利于 GM 到 UB 的搬运并且兼容非 32B 对齐的尾块。5.5 Double Buffer 与预取下一 tileKernel 使用BUFFER_NUM 2并在处理当前 tile 前预取下一 tileCopyIn(nextOffset, nextLen);整体流程近似为预取第 0 块 循环预取下一块 - 计算当前块 - 写回当前块这种方式可以降低 DMA 搬运等待对 Vector 计算的影响。5.6 不同 dtype 使用不同 TILE_LENGTHKernel 根据 dtype 在编译期选择 tile 大小TILE_LENGTH IsSameTypeT, half::value ? 10240 : 6144;half每个元素 2 字节可以使用更大的 tilefloat每个元素 4 字节tile 略小避免 UB 压力过高。6. 编译方法进入工程根目录后执行source /usr/local/Ascend/ascend-toolkit/set_env.sh rm -rf build mkdir build cd build cmake .. make -j8如果平台 CANN 路径不同先查找环境脚本find /usr/local/Ascend -name set_env.sh然后替换source路径。7. 编译成功检查编译完成后可检查生成文件find . -name *.so -o -name *.run -o -name *.json比赛平台通常只需要工程编译通过然后直接提交评测。8. 常见问题8.1Could not find ASCCANN 环境变量没有加载。执行source /usr/local/Ascend/ascend-toolkit/set_env.sh然后重新清理 build 编译。8.2fast_gelu_x_tilingkey does not name a template说明 Kernel 入口和tiling_key_fast_gelu.h的模板机制不匹配。当前版本要求保留template typename DT_X __global__ __aicore__ void fast_gelu(...)同时保留tiling_key_fast_gelu.h中的ASCENDC_TPL_ARGS_DECL和ASCENDC_TPL_SEL。8.3 输出结果异常优先检查是否清理了旧的build是否使用了当前版本所有文件fast_gelu_tiling.h中的结构体是否与 Host / Kernel 完全一致op_host/fast_gelu.cpp是否调用了ASCENDC_TPL_SEL_PARAM(context, DT_X)。建议每次替换文件后执行rm -rf build mkdir build cd build cmake .. make -j89. 版本备注该版本偏向性能优化使用等价化简公式减少计算链并对float16使用 float 中间计算保证精度。若测试数据包含极大负数且平台严格检查Inf/NaN行为可考虑将 Kernel 的Compute部分切换回题目给出的稳定原式。【免费下载链接】cann-launch-camp用户可借助此项目规范管理 CANN 开源社区启航营高校活动的课程作业、课设、毕设等实践作品其核心功能是统一标准化目录层级与提交规范保障作品提交规整、可追溯、可评审。项目地址: https://gitcode.com/cann/cann-launch-camp创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考