CANN ops-nn 中 SwishGrad 算子解析原理、實(shí)現(xiàn)與 aclnnSwishBackward 調(diào)用指南【免費(fèi)下載鏈接】ops-nn本項(xiàng)目是CANN提供的神經(jīng)網(wǎng)絡(luò)類計算算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計算。項(xiàng)目地址: https://gitcode.com/cann/ops-nn導(dǎo)讀SwishGrad 是 CANN 神經(jīng)網(wǎng)絡(luò)算子庫ops-nn中用于計算 Swish 激活函數(shù)梯度的反向算子實(shí)現(xiàn)在倉庫的 experimental/activation/swish_grad 目錄下面向 Atlas A2 訓(xùn)練系列產(chǎn)品/Atlas 800I A2 推理產(chǎn)品提供 fp16、fp32、bf16 三種精度的梯度計算能力。本文以 swish_grad/README.md 為核心骨架結(jié)合倉庫中算子定義、Host 側(cè) tiling、Kernel 側(cè)昇騰 AI Core 實(shí)現(xiàn)與單元測試完整講解 SwishGrad 的數(shù)學(xué)原理、參數(shù)約束、兩段式 aclnn 接口調(diào)用流程與源碼級實(shí)現(xiàn)機(jī)制幫助開發(fā)者快速上手在 NPU 上完成 Swish 反向傳播的接入與調(diào)試。一、算子功能與數(shù)學(xué)原理1.1 功能定位SwishGrad 算子是 Swish 激活函數(shù)activation/swish的反向傳播算子求 Swish 函數(shù)的梯度。它接收反向傳播傳入的梯度grad、正向輸入x以及可調(diào)節(jié)參數(shù)scale輸出對輸入x的梯度grad_x。從算子庫的接口關(guān)系來看上層對外暴露的 API 名為aclnnSwishBackward其功能說明文檔 docs/aclnnSwishBackward.md 明確指出它是 aclnnSwish倉庫實(shí)際路徑為 activation/swish/docs/aclnnSwish.md激活函數(shù)的反向傳播用于計算 Swish 激活函數(shù)的梯度。這與 torch 生態(tài)中的swish_backward語義一一對應(yīng)。1.2 計算公式根據(jù) swish_grad/README.md 的功能說明SwishGrad 的計算公式如下$$ y sigmoid(scalex) xsigmoid(scale*x) $$$$ sigmoid sigmoid*(1 - sigmoid) $$其中scale為可調(diào)節(jié)的斜率參數(shù)。更精確的推導(dǎo)形式見 docs/aclnnSwishBackward.mdSwish 正向函數(shù)$$ s(x) x*\sigma(\beta x) $$其導(dǎo)數(shù)$$ s^\prime(x) \beta s(x)\sigma(\beta x)(1-\beta s(x)) \sigma(\beta x)*(1\beta x(1-\sigma(\beta x))) $$最終輸出梯度$$ gradInput gradOutput * s^\prime(x) $$其中 Sigmoid 函數(shù)定義為$$ \sigma(x) {\frac{1} {1{e}^{-x}}} $$可見 README 中的兩條公式本質(zhì)是先算 sigmoid 值再算 1 - sigmoid即導(dǎo)數(shù)因子最后與梯度逐元素相乘這一計算鏈的簡寫。對照 op_kernel/swish_grad.h 的Compute實(shí)現(xiàn)可以逐行印證該公式σ(βx) 1 / (1 exp(-βx)) σ(βx) σ(βx) * (1 - σ(βx)) s(x) σ(βx) βx * σ(βx) // 即 σ(βx) * (1 βx(1 - σ(βx))) grad_x grad * s(x)1.3 參數(shù)說明README 中對算子原語的參數(shù)定義如下參數(shù)名輸入/輸出/屬性描述數(shù)據(jù)類型數(shù)據(jù)格式grad輸入待進(jìn)行 SwishGrad 計算的入?yún)⒐街械?gradfp16、fp32、bf16ND, FRACTAL_NZ, NC1HWC0x輸入待進(jìn)行 SwishGrad 計算的入?yún)⒐街械?xfp16、fp32、bf16ND, FRACTAL_NZ, NC1HWC0scale輸入待進(jìn)行 SwishGrad 計算的入?yún)⒐街械?scalefp321標(biāo)量y輸入待進(jìn)行 SwishGrad 計算的入?yún)簳r不參與運(yùn)算fp16、fp32、bf16ND, FRACTAL_NZ, NC1HWC0grad_x輸出待進(jìn)行 SwishGrad 計算的出參公式中的輸出fp16、fp32、bf16ND, FRACTAL_NZ, NC1HWC0約束說明無。需要指出兩點(diǎn)值得注意的細(xì)節(jié)算子定義中的y輸入暫時不參與運(yùn)算。這在 op_kernel/swish_grad.cpp 的 kernel 入口中可以看到swish_grad內(nèi)核函數(shù)雖然接收GM_ADDR y但 op_kernel/swish_grad.h 的Init只為grad、x、grad_x建立了 GlobalBuffery并未被讀取屬于為后續(xù)版本預(yù)留的入?yún)?。上述表格描述的是底層算子的參?shù)形態(tài)對上層用戶而言真正接觸的是aclnnSwishBackward接口的參數(shù)見第三節(jié)。二、算子源碼結(jié)構(gòu)SwishGrad 在倉庫中采用標(biāo)準(zhǔn)的 CANN 算子工程布局experimental/activation/swish_grad/ ├── docs/aclnnSwishBackward.md # aclnn 接口使用文檔 ├── examples/test_aclnn_swish_grad.cpp # 可運(yùn)行調(diào)用樣例 ├── op_host/ │ ├── op_api/ # l0 算子接口封裝l0op::SwishGrad │ │ ├── aclnn_swish_backward.cpp/h # 對外 aclnn 兩段式接口 │ │ └── swish_grad.cpp/h # l0 算子封裝與 launcher │ ├── swish_grad_def.cpp # 算子原語注冊O(shè)pDef │ ├── swish_grad_infershape.cpp # 輸出 shape 推導(dǎo) │ └── swish_grad_tiling.cpp # Host 側(cè) tiling 計算 ├── op_kernel/ │ ├── swish_grad.cpp # AI Core 內(nèi)核入口__global__ __aicore__ │ ├── swish_grad.h # Kernel 類實(shí)現(xiàn)雙緩沖流水 │ ├── swish_grad_tiling_data.h # tiling 參數(shù)結(jié)構(gòu)體 │ └── swish_grad_tiling_key.h # tiling key └── tests/ut/ # op_api / op_host / op_kernel 三級單測整個調(diào)用鏈為aclnnSwishBackwardGetWorkspaceSize→l0op::SwishGradHost 側(cè)拼裝算子→ADD_TO_LAUNCHER_LIST_AICORE下發(fā) → AI Core 內(nèi)核swish_grad執(zhí)行。下面逐層展開。2.1 算子定義注冊op_defswish_grad_def.cpp 中通過OP_TYPE_REGISTER/OP_ADD注冊名為SwishGrad的算子四個數(shù)據(jù)對象輸入grad、x、y輸出grad_x數(shù)據(jù)類型支持DT_FLOAT16、DT_FLOAT、DT_BF16格式支持ND、NC1HWC0、FRACTAL_NZ三種屬性sconfAttrType(OPTIONAL).Float(1.0)即可選的 float 型斜率參數(shù)默認(rèn)值為 1.0對應(yīng)公式中的scale/βAICore 配置DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)、PrecisionReduceFlag(true)注冊到ascend910bSoC對應(yīng) Atlas A2 系列并通過ExtendCfgInfo(opFile.value, swish_grad)指定 kernel 入口文件名為swish_grad.cpp。2.2 輸出 shape 推導(dǎo)infershapeswish_grad_infershape.cpp 的實(shí)現(xiàn)非常直接讀取輸入x的 shape并將輸出grad_x的 shape 賦值為與之完全相同*yShape *xShape即逐元素反向算子輸出與輸入同 shape、同 dtype。這與 README 中g(shù)rad、x 與 grad_x 的 shape/數(shù)據(jù)類型一致的約束一致。2.3 Host 側(cè) tilingswish_grad_tiling.cppswish_grad_tiling.cpp 負(fù)責(zé)在 Host 側(cè)根據(jù)輸入規(guī)模與硬件資源計算切分方案GetPlatformInfo通過platform_ascendc::PlatformAscendC獲取 UBUnified Buffer大小、AI Core 核數(shù)以及塊對齊粒度GetUbBlockSizeGetWorkspaceSize查詢系統(tǒng)庫 API 所需 workspace 大小并回填給框架GetShapeAttrsInfo根據(jù)輸入存儲 shape 與數(shù)據(jù)類型字節(jié)數(shù)計算總數(shù)據(jù)量結(jié)合 UB 容量與雙緩沖需求BUFFER_NUM 2確定每個核負(fù)責(zé)的數(shù)據(jù)量bigCoreDataNum/smallCoreDataNum、tile 內(nèi)數(shù)據(jù)量tileDataNum以及尾部數(shù)據(jù)tailDataNum。最終產(chǎn)出的 tiling 數(shù)據(jù)結(jié)構(gòu)定義在 swish_grad_tiling_data.hstruct SwishGradTilingData { uint64_t smallCoreDataNum; uint64_t bigCoreDataNum; uint64_t finalBigTileNum; uint64_t finalSmallTileNum; uint64_t tileDataNum; uint64_t smallTailDataNum; uint64_t bigTailDataNum; uint64_t tailBlockNum; float sconf; };其核心思想是多核負(fù)載均衡 尾部數(shù)據(jù)處理把總數(shù)據(jù)量按核數(shù)均分若不能整除則前tailBlockNum個核多承擔(dān)bigCoreDataNum個數(shù)據(jù)其余核承擔(dān)smallCoreDataNum個數(shù)據(jù)每個核內(nèi)部再按tileDataNum切成多個 tile 循環(huán)處理最后的余量作為尾 tile 單獨(dú)處理。2.4 AI Core 內(nèi)核實(shí)現(xiàn)op_kernelKernel 入口 swish_grad.cpp 通過REGISTER_TILING_DEFAULT讀取 tiling 數(shù)據(jù)實(shí)例化NsSwishGrad::KernelSwishGradDTYPE_X并調(diào)用Init與Process。核心實(shí)現(xiàn)位于 swish_grad.h其設(shè)計要點(diǎn)如下1雙緩沖流水線TPipe上建立了 VECIN 隊(duì)列inQueueGrad、inQueueX、inQueueY與 VECOUT 隊(duì)列outQueueGrad緩沖區(qū)個數(shù)均為BUFFER_NUM 2實(shí)現(xiàn) CopyIn → Compute → CopyOut 三階段重疊執(zhí)行。2fp16/bf16 與 fp32 兩條計算路徑對float16_t/bfloat16_t先把x通過Cast提升到 float 計算使用臨時緩沖區(qū)tmpQueue0/1/2經(jīng)歷Muls乘 β→Muls(-1.0)→Exp→Adds(1.0)→Div得到 σ(βx)→Sub得到 1-σ(βx)→Mul得到 βx·σ(βx)→Add(1.0)得到 1βx(1-σ(βx))→Mul與 σ(βx) 相乘 → 再與 grad 相乘最后以CAST_RINT舍入模式轉(zhuǎn)回半精度輸出對float直接在浮點(diǎn)路徑上完成同樣序列的運(yùn)算全程使用 float 精度輸出不降精度。3多核數(shù)據(jù)切片Init中根據(jù)coreId與tailBlockNum的關(guān)系選擇大核/小核的數(shù)據(jù)區(qū)間與 GlobalBuffer 偏移每個核只處理自己負(fù)責(zé)的連續(xù)數(shù)據(jù)段。內(nèi)核從sconf即 scale/β讀取斜率參數(shù)其值由 Host 側(cè)在aclnnSwishBackwardGetWorkspaceSize中從betaOptional標(biāo)量轉(zhuǎn)換而來缺省為 1.0f。2.5 測試用例佐證倉庫為算子提供了三級單元測試可用于驗(yàn)證實(shí)現(xiàn)與本文結(jié)論op_api 層tests/ut/op_api/test_aclnn_swish_backward.cpp覆蓋 float32、float16、bf16 等場景構(gòu)造{2,5}的 gradOutput、self值域 [-1,1]與 beta如 1.1f、0.01f斷言aclnnSwishBackwardGetWorkspaceSize返回ACLNN_SUCCESS并設(shè)置了精度容差op_host 層tests/ut/op_host/test_swish_grad_tiling.cpp驗(yàn)證 tiling 計算產(chǎn)物各核數(shù)據(jù)量、tile 數(shù)等符合預(yù)期op_kernel 層tests/ut/op_kernel/test_swish_grad.cpp直接對 AI Core 內(nèi)核做數(shù)值驗(yàn)證。三、aclnnSwishBackward 接口詳解3.1 產(chǎn)品支持情況產(chǎn)品是否支持Atlas A2 訓(xùn)練系列產(chǎn)品/Atlas 800I A2 推理產(chǎn)品√對應(yīng)到 SoC 層面算子原語注冊于ascend910b平臺見 swish_grad_def.cpp且 aclnn_swish_backward.cpp 中的CheckSocVersionIsSupportBf16顯示 bf16 支持范圍限定在ASCEND910B至ASCEND910E之間的 SoC 版本。3.2 兩段式接口模型與 CANN 其他 aclnn 算子一致aclnnSwishBackward采用兩段式接口詳見 docs/zh/context/two_phase_api.md先調(diào)用GetWorkspaceSize接口完成入?yún)⑿ｒ?yàn)、算子拼裝并獲取所需 workspace 大小與執(zhí)行器再調(diào)用執(zhí)行接口真正下發(fā)計算。第一段接口原型aclnnStatus aclnnSwishBackwardGetWorkspaceSize( const aclTensor* gradOutput, // 正向輸出梯度公式中的 gradOutput const aclTensor* self, // Swish 激活函數(shù)輸入公式中的 x const aclScalar* betaOptional, // 可調(diào)節(jié)參數(shù) β標(biāo)量空指針時按 1.0 計算 aclTensor* gradInput, // 輸出對輸入的梯度 uint64_t* workspaceSize, aclOpExecutor** executor)第二段接口原型aclnnStatus aclnnSwishBackward( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)3.3 第一段接口參數(shù)說明參數(shù)名輸入/輸出描述使用說明數(shù)據(jù)類型數(shù)據(jù)格式維度(shape)非連續(xù) TensorgradOutput輸入Swish 激活函數(shù)正向輸出的梯度公式中的 gradOutput支持空 TensorgradOutput、self 與 gradInput 的 shape 一致三者的數(shù)據(jù)類型一致BFLOAT16、FLOAT16、FLOATND0-8√self輸入用于計算激活函數(shù)的張量公式中的 x支持空 Tensor同上BFLOAT16、FLOAT16、FLOATND0-8√betaOptional輸入可調(diào)節(jié)參數(shù)控制 Swish 函數(shù)的形狀和斜率公式中的 β數(shù)據(jù)類型需可轉(zhuǎn)換為 FLOAT參見 docs/zh/context/deduction_relationship.md 的互推導(dǎo)關(guān)系空指針時以 1.0 計算----gradInput輸出backward 計算的輸出Swish 正向輸入的梯度值支持空 Tensorshape 與數(shù)據(jù)類型同前兩者BFLOAT16、FLOAT16、FLOATND0-8√workspaceSize輸出返回需要在 Device 側(cè)申請的 workspace 大小-----executor輸出返回 op 執(zhí)行器包含算子計算流程-----返回值aclnnStatus具體返回碼參見 docs/zh/context/aclnn_return_code.md。第一段接口會完成入?yún)⑿ｒ?yàn)以下場景會報錯返回碼錯誤碼描述ACLNN_ERR_PARAM_NULLPTR161001傳入的 gradOutput、self 或 gradInput 是空指針ACLNN_ERR_PARAM_INVALID161002gradOutput、self、betaOptional 或 gradInput 的數(shù)據(jù)類型不在支持范圍之內(nèi)ACLNN_ERR_PARAM_INVALID161002gradOutput、self 和 gradInput 的數(shù)據(jù)類型不同ACLNN_ERR_PARAM_INVALID161002gradOutput、self 和 gradInput 的 shape 不同3.4 第二段接口參數(shù)說明參數(shù)名輸入/輸出描述workspace輸入在 Device 側(cè)申請的 workspace 內(nèi)存地址workspaceSize輸入在 Device 側(cè)申請的 workspace 大小由第一段接口aclnnSwishBackwardGetWorkspaceSize獲取executor輸入op 執(zhí)行器包含算子計算流程stream輸入指定執(zhí)行任務(wù)的 Stream返回值同樣為aclnnStatus。約束說明無。3.5 接口內(nèi)部實(shí)現(xiàn)要點(diǎn)aclnn_swish_backward.cpp 的實(shí)現(xiàn)清晰展示了接口層的處理流程可作為閱讀模板參數(shù)校驗(yàn)CheckParams依次執(zhí)行空指針檢查CheckNotNull、數(shù)據(jù)類型檢查CheckDtypeValid支持列表為 FLOAT16/FLOAT/BF16且校驗(yàn) betaOptional 可轉(zhuǎn) FLOAT、三個張量 dtype 一致、bf16 受 SoC 版本約束與 shape 檢查CheckShapeValid維度上限為MAX_SUPPORT_DIMS_NUMS三個張量 shape 一致空 Tensor 短路當(dāng)gradOutput或self為空時workspaceSize直接置 0 返回不執(zhí)行計算連續(xù)性處理對非連續(xù)輸入調(diào)用l0op::Contiguous轉(zhuǎn)為連續(xù) Tensor對應(yīng)參數(shù)表中非連續(xù) Tensor √的支持能力高維 reshape當(dāng)輸入維度超過 8 時將輸入展平為 1 維長 Tensor 交給 kernel計算完后再Reshape回原維度β 轉(zhuǎn)換betaOptional通過ToFloat()轉(zhuǎn)為 float空指針時取默認(rèn)值1.0f下發(fā)與收尾調(diào)用l0op::SwishGrad完成算子拼裝通過l0op::ViewCopy將結(jié)果拷貝到可能非連續(xù)的輸出gradInput上最終以GetWorkspaceSize()匯總整體 workspace 需求第二段接口則直接調(diào)用框架通用執(zhí)行函數(shù)CommonOpExecutorRun完成異步下發(fā)。四、完整調(diào)用示例與運(yùn)行指引4.1 可運(yùn)行樣例倉庫提供了可直接參考的完整樣例 examples/test_aclnn_swish_grad.cpp編譯與運(yùn)行方法可參照 docs/zh/context/compile_and_run_sample.md。樣例的核心流程如下為便于閱讀做了精簡完整代碼以倉庫文件為準(zhǔn)#include acl/acl.h #include aclnn_swish_backward.h int main() { // 1. 固定寫法device/stream 初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // aclInit aclrtSetDevice aclrtCreateStream // 2. 構(gòu)造輸入與輸出示例 shape 均為 {2, 3} std::vectorint64_t gradOutputShape {2, 3}; std::vectorint64_t selfShape {2, 3}; std::vectorint64_t gradInputShape {2, 3}; std::vectorfloat gradOutHostData {1, 1, 1, 1, 1, 1}; std::vectorfloat selfHostData {1, 2, 3, 4, 5, 6}; float betaValue 1.1f; // 用 aclrtMalloc/aclrtMemcpy 把 host 數(shù)據(jù)搬到 device并 aclCreateTensor 創(chuàng)建 aclTensor aclTensor* gradOut CreateAclTensor(gradOutHostData, gradOutputShape, ..., ACL_FLOAT, ...); aclTensor* self CreateAclTensor(selfHostData, selfShape, ..., ACL_FLOAT, ...); aclTensor* gradInput CreateAclTensor(/* 輸出占位 */, gradInputShape, ..., ACL_FLOAT, ...); aclScalar* betaOptional aclCreateScalar(betaValue, aclDataType::ACL_FLOAT); // 3. 兩段式調(diào)用 uint64_t workspaceSize 0; aclOpExecutor* executor; ret aclnnSwishBackwardGetWorkspaceSize(gradOut, self, betaOptional, gradInput, workspaceSize, executor); void* workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } ret aclnnSwishBackward(workspaceAddr, workspaceSize, executor, stream); // 4. 同步等待執(zhí)行結(jié)束 aclrtSynchronizeStream(stream); // 5. 將 device 側(cè)結(jié)果拷回 host 并打印本例輸出為 gradInput grad * s(x) std::vectorfloat outData(6, 0); aclrtMemcpy(outData.data(), ..., gradInputDeviceAddr, ..., ACL_MEMCPY_DEVICE_TO_HOST); // 6-7. 釋放 aclTensor / aclScalar / device 內(nèi)存 / stream并 aclFinalize return 0; }樣例中g(shù)radOutput {1,1,1,1,1,1}、self {1,2,3,4,5,6}、beta 1.1執(zhí)行后打印的 6 個輸出值即為對應(yīng)位置的grad * s(x)數(shù)值可直接對照第一節(jié)的公式手工驗(yàn)算。4.2 調(diào)用步驟要點(diǎn)初始化aclInit(nullptr)→aclrtSetDevice(deviceId)→aclrtCreateStream(stream)構(gòu)造 Tensorshape/stride 一致、數(shù)據(jù)類型為 FLOAT16/FLOAT/BF16、格式為 ND支持非連續(xù) Tensor接口內(nèi)部會自動Contiguous支持空 Tensor兩段式調(diào)用第一段拿到workspaceSize與executor若workspaceSize 0則在 device 側(cè)申請內(nèi)存然后調(diào)用第二段aclnnSwishBackward同步與取數(shù)aclrtSynchronizeStream后通過aclrtMemcpyDEVICE_TO_HOST讀取輸出資源釋放依次釋放aclTensor/aclScalar、device 內(nèi)存、workspace、stream最后aclrtResetDeviceaclFinalize。4.3 注意事項(xiàng)三個張量gradOutput、self、gradInput的shape 與數(shù)據(jù)類型必須完全一致否則第一段接口返回ACLNN_ERR_PARAM_INVALID161002betaOptional傳nullptr時按β 1.0計算此時算子退化為 PyTorch 中默認(rèn)swish的梯度形式傳入時其數(shù)據(jù)類型需可互推導(dǎo)為 FLOATbf16 輸入僅在支持 bf16 的 SoC 版本ASCEND910B~ASCEND910E上可用該算子底層為逐元素計算輸出與輸入同 shapey入?yún)轭A(yù)留項(xiàng)當(dāng)前版本不參與運(yùn)算。五、總結(jié)SwishGrad 算子是 ops-nn 倉庫中一個典型的逐元素反向算子樣例完整展示了 CANN 算子的標(biāo)準(zhǔn)五件套OpDef 注冊swish_grad_def.cpp、InferShapeswish_grad_infershape.cpp、Host 側(cè) tilingswish_grad_tiling.cpp、AI Core 內(nèi)核op_kernel/swish_grad.h與 aclnn 兩段式對外接口aclnn_swish_backward.cpp。理解其數(shù)學(xué)推導(dǎo)σ(βx)·(1βx(1-σ(βx)))、多核均衡 tiling 策略、半精度轉(zhuǎn)浮點(diǎn)計算的精度處理以及兩段式接口的調(diào)用范式即可舉一反三地接入 ops-nn 中其他同構(gòu)的激活函數(shù)反向算子如 SiluGrad。實(shí)際開發(fā)中若需驗(yàn)證算子行為可參考倉庫 tests/ut 下的 op_api、op_host、op_kernel 三級測試用例或直接編譯運(yùn)行 examples/test_aclnn_swish_grad.cpp 樣例?！久赓M(fèi)下載鏈接】ops-nn本項(xiàng)目是CANN提供的神經(jīng)網(wǎng)絡(luò)類計算算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計算。項(xiàng)目地址: https://gitcode.com/cann/ops-nn創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考