ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

CANN ops-nn 算子实战:aclnnSituGlu 两段式接口详解与 SiTU 门控线性单元 NPU 实现

CANN ops-nn 算子实战:aclnnSituGlu 两段式接口详解与 SiTU 门控线性单元 NPU 实现 CANN ops-nn 算子实战aclnnSituGlu 两段式接口详解与 SiTU 门控线性单元 NPU 实现【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn本篇技术指南围绕 CANN ops-nn 仓库中 SiTU 门控线性单元SituGlu算子的官方 aclnn 接口文档展开系统讲解aclnnSituGluGetWorkspaceSize/aclnnSituGlu两段式接口的函数原型、全部入参与返回值、错误码语义以及可复制的完整调用示例并结合 activation/situ_glu 目录下的算子定义、Tiling 与 Kernel 源码深入解读其在 NPU 上的计算流程与实现细节。读完本文你可以直接在自己的 NPU 工程中接入该算子并理解其在 AscendC 内核层面的工作原理。产品支持情况aclnnSituGlu 作为 SituGlu 算子的 C 语言 aclnn 接口在不同 NPU 产品上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持这一支持范围与算子注册的 AICore 配置一一对应。在 算子定义文件 中SituGlu通过AICore().AddConfig(ascend910b)、AddConfig(ascend910_93)、AddConfig(ascend950)注册了三个硬件平台Tiling 侧的IsCapable()校验见 situ_glu_tiling.cpp同样只接受ASCEND910_93、ASCEND910B、ASCEND950三种 SoC 版本与上表结论一致。功能说明接口功能SiTU 门控线性单元SiTU Gated Linear Unit激活函数。对输入张量x沿指定维度切分为门控gate与上路径up两半按 SiTU 公式计算输出。对给定的输入张量x其维度为[a, b, c, d, e, f, g, ...]aclnnSituGlu 执行以下计算流程合轴将x基于输入参数dim进行合轴合轴后维度为[pre, cut]。其中cut轴为需要切分为两个张量的轴cut必须为偶数令h cut // 2。前后切分根据输入参数activateLeft对x进行切分activateLeft为true时$$ gate x[:, :h], \quad up x[:, h:] $$activateLeft为false时$$ gate x[:, h:], \quad up x[:, :h] $$SiTU 计算根据输入参数beta、linearBeta计算$$ situ_a \beta \cdot \tanh\left(\frac{gate}{\beta}\right) \cdot \text{sigmoid}(gate) $$当linearBeta 0时对 up 路径施加有界化线性 tanh 变换$$ up linearBeta \cdot \tanh\left(\frac{up}{linearBeta}\right) $$输出$$ y situ_a \cdot up $$重塑输出输出张量y的维度数量与合轴前的x一致dim轴上的大小为x的一半其他维度与x相同。从数学形式上看situ_a是 tanh有界非线性与 sigmoid软门控的逐元素乘积再与 up 路径相乘融合了门控线性单元GLU的通道选择特性与 tanh 的有界激活特性linearBeta 0时还会把 up 路径限制在(-linearBeta, linearBeta)区间内linearBeta 0时 up 直接透传。函数原型与两段式接口每个算子分为两段式接口详见 两段式接口说明必须先调用aclnnSituGluGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnSituGlu接口执行计算。aclnnStatus aclnnSituGluGetWorkspaceSize( const aclTensor *x, int64_t dim, double beta, double linearBeta, bool activateLeft, const aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnSituGlu( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)接口头文件为aclnnop/aclnn_situ_glu.h调用示例见 examples/test_aclnn_situ_glu.cpp。aclnnSituGluGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorxaclTensor*输入公式中的输入 x。支持空 Tensor输出为对应空 Tensorshape 必须在入参 dim 对应的维度上为偶数。FLOAT、FLOAT16、BFLOAT16ND1-8√dimint64_t输入公式中的输入 dim表示对 x 进行合轴以及切分的维度序号。取值范围为[-x.dim(), x.dim()-1]。----betadouble输入公式中的输入 betaSiTU 门控部分的缩放系数。建议为非零值。----linearBetadouble输入公式中的输入 linearBetaup 路径线性 tanh 的缩放系数。小于等于 0 时 up 路径不进行变换。-----activateLeftbool输入公式中的输入 activateLeft表示切分 x 时 gate 是否为前半部分。true 表示 gate 为前半、up 为后半false 表示 gate 为后半、up 为前半。----outaclTensor*输出公式中的输出 y。支持空 Tensorshape 在入参 dim 对应的维度上为 x 的一半其他维度与 x 一致数据类型与 x 一致。FLOAT、FLOAT16、BFLOAT16ND1-8×workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----几个值得注意的细节非连续 Tensor输入x支持非连续 Tensor表中为 √而输出out要求连续表中为 ×。这与算子定义中x输入声明了.AutoContiguous()的处理方式相呼应。空 Tensorx与out均支持空 Tensor此时输出为对应的空 Tensor适合动态 shape 场景下的占位。dim 的负索引语义dim取值区间为[-x.dim(), x.dim()-1]负数表示从后往前数。在 Tiling 实现中负数会被统一转换为正数cutDim_ cutDim_ 0 ? (cutDim_ xDims_) : cutDim_见 situ_glu_tiling.cpp。beta 建议非零由于公式中存在gate / beta与up / linearBetabeta 为 0 会导致除零文档建议 beta 取非零值linearBeta 0时 up 路径直接透传不会触发除法。返回值与错误码aclnnStatus返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001x、out 存在空指针。ACLNN_ERR_PARAM_INVALID161002x 或 out 的数据类型不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002x 或 out 的 shape 维度不在支持的范围之内。ACLNN_ERR_PARAM_INVALID161002dim 不在指定的取值范围内或 x 在 dim 对应维度上不能被 2 整除。ACLNN_ERR_PARAM_INVALID161002out 的维度数与 x 不一致或 out 在 dim 对应维度上的大小不等于 x 的一半。ACLNN_ERR_PARAM_INVALID161002out 的数据类型与 x 不匹配。这些校验规则与底层 Tiling/Infershape 中的检查逻辑一一对应CheckAndGetXAndAttrs中校验dim范围、cutDimSize_ % 2 1奇数报错以及数据类型仅支持 float32/float16/bfloat16situ_glu_tiling.cppCheckY中校验输出维度数、输出 dim 大小等于输入一半以及 dtype 一致situ_glu_tiling.cpp。Shape 推导侧的校验见 situ_glu_infershape.cpp。aclnnSituGlu 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnSituGluGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值aclnnStatus返回状态码具体参见 aclnn 返回码。关于 workspace 大小Tiling 阶段通过GetWorkspaceSize()返回固定值16 * 1024 * 102416MB见 situ_glu_tiling.cpp 与 PostTiling 中的workspaceSize_ WORKSPACE_SIZE并在PostTiling中通过context_-SetBlockDim(coreNumAll_)设置使用全部可用 AIV Core 并行计算。Kernel 侧situ_glu入口对 workspace 做了空指针保护并通过GetUserWorkspace获取用户工作区后再进行搬运与计算见 situ_glu.cpp。约束说明确定性说明aclnnSituGlu 默认确定性实现即相同输入与参数下多次执行结果一致。这也与 PyTorch 封装侧文档默认支持确定性计算的说明一致见 torchapi_situ_glu.md。输入支持 1-8 维、ND 格式dim 对应维度必须为偶数。调用示例示例代码如下完整可编译版本见 examples/test_aclnn_situ_glu.cpp具体编译和执行过程请参考 编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_situ_glu.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法acl初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API手册 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出 std::vectorint64_t xShape {2, 32}; std::vectorint64_t outShape {2, 16}; void* xDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* x nullptr; aclTensor* out nullptr; std::vectorfloat xHostData {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; int64_t dim -1; double beta 1.0; double linear_beta 0.0; bool activate_left true; // 创建x aclTensor ret CreateAclTensor(xHostData, xShape, xDeviceAddr, aclDataType::ACL_FLOAT, x); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnSituGlu第一段接口 ret aclnnSituGluGetWorkspaceSize(x, dim, beta, linear_beta, activate_left, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSituGluGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnSituGlu第二段接口 ret aclnnSituGlu(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSituGlu failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor aclDestroyTensor(x); aclDestroyTensor(out); // 7. 释放device资源 aclrtFree(xDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例运行流程拆解初始化aclInit→aclrtSetDevice→aclrtCreateStream这是所有 aclnn 算子调用前的固定环境准备。构造张量通过aclrtMalloc在 Device 侧申请内存aclrtMemcpy将 Host 数据拷入再计算连续 Tensor 的 strides 并调用aclCreateTensor创建x与out两个aclTensor。示例中x的 shape 为[2, 32]out为[2, 16]即dim-1最后一维上从 32 切半到 16。两段式调用先调用aclnnSituGluGetWorkspaceSize获取 workspace 大小与 executor若workspaceSize 0则申请对应 Device 内存再调用aclnnSituGlu在指定 stream 上执行。同步与取数aclrtSynchronizeStream等待异步任务完成后用aclrtMemcpy把结果从 Device 拷回 Host 并打印。资源释放依次释放aclTensor、Device 内存、workspace、stream并aclrtResetDeviceaclFinalize收尾。源码级实现原理算子定义OP Def算子定义文件 中通过OpDef注册了SituGlu算子一个必选输入x支持 FLOAT/FLOAT16/BF16、ND 格式、未知 shape一个必选输出y以及四个可选属性dim默认 -1、beta默认 1.0、linear_beta默认 0.0、activate_left默认 true。与 GE 算子原型 中REG_OP(SituGlu)的声明完全一致是 aclnn 接口参数默认值的底层来源。Shape 推导InferShapeInferShape 实现 的逻辑为复制x的 shape 到y然后将dim负数先转正对应维度除以 2SPLIT_NUM若该维为动态维度-1则直接返回若为奇数或负数则报错。数据类型推导则直接将输入 dtype 赋给输出InferDataTypeForSituGlu。Tiling 与多核切分Tiling 侧的核心思想是把任意维度x合轴为二维[dimBatchSize, dim2H]situ_glu_tiling.cppdimBatchSizedim之前所有维度的乘积dim2Hdim及之后所有维度的乘积切分前总长切分后半行长度dimH dim2H / 2。在此基础上根据 UBUnified Buffer容量计算单次最多可处理的 gateup 元素对 数量ubMaxPair_并区分两条搬运路径Short-H 路径isLongH_ 0半行满足 32B 对齐时按行 stride 拷贝blockCount 1并在 Core 间、Core 内两层切分 batchLong-H 路径isLongH_ 1半行不满足 32B 对齐或单行过大时回退为连续单块拷贝由DataCopyPad处理 sub-block 对齐功能正确但性能略降situ_glu_tiling.cpp。PyTorch 封装文档中dim 非 last 维时若尾部维度组合不满足半行 32B 对齐会回退到 Long-H 路径的说明正源于此。Kernel 计算内核Kernel 侧是一个模板类SituGluBaseTsitu_glu.hpp覆盖三种 dtypeSituGluBasefloatfp32 原生计算无类型转换SituGluBasehalf/SituGluBasebfloat16_t先Cast到 float32 计算再转回原类型bf16 使用CAST_RINT舍入模式避免非 A2 平台上产生垃圾值。核心计算在SituCore中展开situ_glu.hpp与文档公式严格一致Muls(gate, 1/beta)→Tanh→Muls(beta)得到beta·tanh(gate/beta)Sigmoid(gate)与其相乘得到situ_a当linearBeta 0时对 up 做同样的linearBeta·tanh(up/linearBeta)变换后与situ_a相乘否则直接相乘。activateLeft通过选择 gate/up 在 xQueue 缓冲中的偏移前半[0:pairNum_]还是后半[halfOff_:halfOff_pairNum_]来决定前后切分方向situ_glu.hpp。Kernel 入口situ_glu根据编译期宏ORIG_DTYPE_X分派到对应模板实例situ_glu.cpp。配套调用方式与测试验证除 aclnn C 接口外SituGlu 还提供 PyTorch 封装cann_ops_nn.situ_glu其函数原型、参数默认值与 aclnn 接口完全对应示例见 torchapi_situ_glu.md底层通过 situ_glu.py 中注册的situ_glu(Tensor x, *, int dim-1, float beta1.0, float linear_beta0.0, bool activate_leftTrue) - Tensorschema 封装调用。模块级说明可参考 situ_glu 模块 README。仓库同时提供了 ATK 算子测试用例 atk_aclnnSituGlu.json覆盖 bf16/fp16/fp32 三种 dtype、dim取负值与正值、beta ∈ [0.1, 10.0]、linearBeta ∈ [0.0, 50.0]、activateLeft取 true/false以及 1-8 维各种 shape含不规则维度如[7, 9, 8, 9, 11, 6, 13, 12]、超长单维如[15432]、多维组合如[14870, 658]可作为自行验证算子正确性与性能边界的参考用例集。使用建议与注意事项使用前请确认目标设备属于文档列出的支持产品Ascend 950PR/DT、Atlas A3、Atlas A2否则接口将校验失败输入x在dim维必须为偶数且out的 shape 需预先正确分配dim 维为 x 的一半否则第一段接口返回ACLNN_ERR_PARAM_INVALID161002beta建议取非零值不需要 up 路径变换时保持linearBeta 0默认 0.0即可workspace 必须按第一段接口返回的大小在 Device 侧申请第二段接口执行后按需释放默认确定性实现适合对可复现性有要求的训练场景。【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-nn创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表