FastGelu 自定义算子【免费下载链接】cann-launch-camp用户可借助此项目规范管理 CANN 开源社区启航营高校活动的课程作业、课设、毕设等实践作品其核心功能是统一标准化目录层级与提交规范保障作品提交规整、可追溯、可评审。项目地址: https://gitcode.com/cann/cann-launch-camp1. 算子概述FastGelu是基于 Ascend C 实现的自定义 AICore 算子用于计算 Fast GELU 激活函数。该算子支持float16和float32输入输出张量与输入张量保持相同的 shape 和数据类型。Fast GELU 是 GELU 激活函数的一种高效近似形式常用于 Transformer、BERT 等神经网络模型中可在保持较好非线性表达能力的同时降低计算复杂度。2. 计算公式算子对输入张量x中的每个元素执行如下计算y x * sigmoid(1.702 * x)其中sigmoid(t) 1 / (1 exp(-t))3. 算子接口3.1 输入参数名类型ShapeFormat说明xfloat16/float32任意 shapeND输入张量3.2 输出参数名类型ShapeFormat说明y与x相同与x相同NDFast GELU 计算结果3.3 支持约束支持数据类型float16、float32支持数据格式ND输入和输出 shape 相同输入和输出数据类型相同当前注册平台配置ascend910b4. 文件结构. ├── CMakeLists.txt ├── op_host │ ├── CMakeLists.txt │ └── fast_gelu.cpp └── op_kernel ├── CMakeLists.txt ├── fast_gelu.cpp ├── fast_gelu_tiling.h └── tiling_key_fast_gelu.h5. Host 侧实现说明Host 侧实现位于op_host/fast_gelu.cpp主要负责算子注册、shape 推导、数据类型推导和 Tiling 参数计算。5.1 算子注册算子名为FastGelu通过OP_ADD(FastGelu)注册。算子定义包含1 个必选输入x1 个必选输出y输入输出均支持float16和float32输入输出 format 均为NDAICore 配置目标为ascend910b5.2 Shape 推导输出张量y的 shape 直接继承输入张量x的 shapey.shape x.shape5.3 数据类型推导输出张量y的数据类型与输入张量x保持一致y.dtype x.dtype5.4 Tiling 逻辑Host 侧根据输入长度、数据类型大小、AIV 核数以及 UB 大小计算 Kernel 执行所需的 Tiling 参数。Tiling 数据结构定义在op_kernel/fast_gelu_tiling.h中字段类型说明lengthuint32_t输入张量元素总数blockLengthuint32_t每个 Core 负责处理的元素数量上限tileLengthuint32_t每次搬运和计算的 Tile 元素数量Tiling 过程主要包括获取输入张量元素总数length_x。检查输入数据类型仅允许float16和float32。根据数据类型大小计算 32 字节 block 对应的元素个数。根据输入规模和 AIV Core 数确定block_dim。小规模张量长度不超过8192时最多使用8个 Core减少调度开销。根据 UB 大小计算tileLength并按 32 字节对齐对应的元素个数向下对齐。设置 workspace 大小为0。6. Kernel 侧实现说明Kernel 侧实现位于op_kernel/fast_gelu.cpp核心类为KernelFastGeluDT_X。Kernel 采用按 Core 切分、按 Tile 循环处理的方式执行计算。每个 Core 只处理自己负责的连续数据片段。6.1 初始化流程Init函数完成如下工作保存 Host 侧传入的length、blockLength和tileLength。根据当前 Core ID 计算该 Core 的全局偏移coreOffset。计算当前 Core 实际需要处理的元素数量coreLength。设置输入、输出 Global Memory Buffer。初始化输入队列、输出队列和临时计算 buffer。计算完整 Tile 数量fullTileNum和尾块元素数量tailDataNum。判断完整 Tile 是否满足 32 字节对齐以选择普通搬运或 Pad 搬运路径。6.2 数据处理流程Process函数为 Kernel 主流程处理完整 Tile。根据对齐情况选择ProcessFullTilesAligned对齐场景使用DataCopy。ProcessFullTilesPad非对齐场景使用DataCopyPad。处理尾块数据ProcessTail。尾块同样根据地址和长度是否 32 字节对齐选择对齐或非对齐搬运方式。6.3 计算流程每个 Tile 的计算逻辑在Compute函数中完成scale x * 1.702 y sigmoid(scale) y y * x对应使用的 Ascend C 向量指令包括Muls将输入乘以常数1.702Sigmoid计算 sigmoidMul将 sigmoid 结果与原输入相乘6.4 数据搬运策略Kernel 中根据数据是否满足 32 字节对齐分别使用两种搬运策略场景搬入搬出对齐DataCopyDataCopy非对齐或尾块DataCopyPadDataCopyPad这样可以同时支持对齐数据和非对齐尾块数据避免因 shape 或切分导致的非 32 字节对齐问题。7. 模板与数据类型选择模板选择相关定义位于op_kernel/tiling_key_fast_gelu.h。该文件通过ASCENDC_TPL_ARGS_DECL和ASCENDC_TPL_SEL定义DT_X模板参数使 Kernel 能够根据输入x的数据类型分别生成DT_X float16DT_X float32Host 侧通过ASCENDC_TPL_SEL_PARAM(context, DT_X)根据输入数据类型选择对应的 Kernel 模板实例。8. 编译说明项目根目录CMakeLists.txt中通过npu_op_package定义自定义算子包并分别添加 Host 和 Kernel 子目录add_subdirectory(op_host) add_subdirectory(op_kernel)通常可在支持 Ascend C 开发环境的机器上执行如下流程进行构建mkdir -p build cd build cmake .. make生成的算子包安装路径由根目录CMakeLists.txt中的配置决定INSTALL_PATH ${CMAKE_BINARY_DIR}9. 注意事项输入数据类型必须为float16或float32否则 Host 侧 Tiling 会返回失败。当前算子 workspace 大小为0计算过程中仅使用 UB 中的输入、输出和临时 buffer。算子按一维连续内存处理输入张量因此支持任意 shape 的ND张量。对于小张量Host 侧会限制最多使用8个 Core以减少多核调度带来的额外开销。tileLength会根据 UB 大小和 32 字节对齐要求动态计算。10. 参考实现位置Host 侧实现op_host/fast_gelu.cppKernel 侧实现op_kernel/fast_gelu.cppTiling 数据结构op_kernel/fast_gelu_tiling.h模板数据类型选择op_kernel/tiling_key_fast_gelu.h【免费下载链接】cann-launch-camp用户可借助此项目规范管理 CANN 开源社区启航营高校活动的课程作业、课设、毕设等实践作品其核心功能是统一标准化目录层级与提交规范保障作品提交规整、可追溯、可评审。项目地址: https://gitcode.com/cann/cann-launch-camp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考