CANN ops-nn Softshrink 算子實戰(zhàn)解析數(shù)學原理、參數(shù)約束與 aclnn 兩段式調(diào)用【免費下載鏈接】ops-nn本項目是CANN提供的神經(jīng)網(wǎng)絡類計算算子庫實現(xiàn)網(wǎng)絡在NPU上加速計算。項目地址: https://gitcode.com/cann/ops-nnSoftshrink軟收縮是一種帶閾值的逐元素激活函數(shù)在神經(jīng)網(wǎng)絡中被用于對張量進行平滑稀疏化與噪聲抑制。本文以 CANN ops-nn 倉庫中 activation/softshrink 模塊為對象完整講解其數(shù)學定義、產(chǎn)品支持情況、參數(shù)規(guī)范、aclnn 兩段式調(diào)用方式并結(jié)合op_api、op_host、op_kernel源碼與單元測試深入剖析其實現(xiàn)原理。讀完本文你將能夠理解 Softshrink 算子在全流程中的定位掌握通過aclnnSoftshrink接口在 NPU 上正確調(diào)用該算子并校驗結(jié)果的完整方法。一、算子功能與數(shù)學原理Softshrink 算子對輸入張量x逐元素施加“軟收縮”變換給定閾值λ將落在[-λ, λ]區(qū)間內(nèi)的元素收縮為 0對區(qū)間外的元素則向 0 方向收縮一個λ的偏移量。計算公式如下見 activation/softshrink/README.md 與 aclnnSoftshrink 接口文檔$$ Softshrink(x)\begin{cases} x-λ, \text{if } x λ \ xλ, \text{if } x -λ \ 0, \text{otherwise} \end{cases} $$從直觀上看Softshrink 與 Hard Shrink硬收縮x -λ取xλx λ取x-λ否則為 0的區(qū)別在于Hard Shrink 在閾值處是跳躍不連續(xù)的而 Softshrink 在閾值處是連續(xù)過渡的因此梯度行為更平滑。該算子常用于去噪、特征稀疏化以及需要對小幅度信號進行抑制的場景。二、產(chǎn)品支持情況根據(jù) activation/softshrink/README.md 中的產(chǎn)品支持矩陣Softshrink 算子在各系列產(chǎn)品上的支持情況如下產(chǎn)品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 訓練系列產(chǎn)品 / Atlas A3 推理系列產(chǎn)品√Atlas A2 訓練系列產(chǎn)品 / Atlas A2 推理系列產(chǎn)品√Atlas 200I/500 A2 推理產(chǎn)品√Atlas 推理系列產(chǎn)品√Atlas 訓練系列產(chǎn)品√說明算子 README 的產(chǎn)品支持矩陣為整體能力聲明具體到aclnnSoftshrink接口層面aclnnSoftshrink.md 中列出了 Atlas 200I/500 A2 推理產(chǎn)品“不支持”的接口級說明并在 Atlas 推理/訓練系列產(chǎn)品上僅支持 FLOAT、FLOAT16 兩種數(shù)據(jù)類型。實際開發(fā)時應以目標版本與具體接口文檔為準。三、參數(shù)說明Softshrink 算子共包含兩個輸入和一個輸出參數(shù)定義如下摘自 activation/softshrink/README.md參數(shù)名輸入/輸出/屬性描述數(shù)據(jù)類型數(shù)據(jù)格式x輸入待進行 Softshrink 計算的輸入張量公式中的 xfp16、fp32、bf16NDlambd輸入Softshrink 的閾值參數(shù)默認值為 0.5fp321y輸出Softshrink 計算的輸出張量fp16、fp32、bf16ND從算子定義源碼 op_host/softshrink_def.cpp 可以印證以上約束輸入x與輸出y的數(shù)據(jù)類型均限定為DT_FLOAT16、DT_FLOAT、DT_BF16格式為ND且開啟AutoContiguous()處理lambd以算子屬性Attr而非輸入張量的形式注冊類型為浮點默認值 0.5與 PyTorch 語義保持一致算子配置為動態(tài) shape、動態(tài) rank 支持DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)能夠處理編譯期未知的張量形狀。在 aclnn 接口層面見下文lambd作為aclScalar*傳入數(shù)據(jù)類型為 FLOAT數(shù)值要求大于等于 0。四、aclnn 兩段式接口詳解CANN 的 aclnn 算子 API 采用兩段式接口設計必須先調(diào)用aclnnSoftshrinkGetWorkspaceSize獲取計算所需 workspace 大小以及包含算子計算流程的執(zhí)行器再調(diào)用aclnnSoftshrink執(zhí)行計算。兩個接口的函數(shù)原型如下aclnnStatus aclnnSoftshrinkGetWorkspaceSize( const aclTensor* self, const aclScalar* lambd, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)aclnnStatus aclnnSoftshrink( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)4.1 第一段接口aclnnSoftshrinkGetWorkspaceSize該接口完成入?yún)⑿ｒ?、?gòu)建算子執(zhí)行器并返回執(zhí)行計算所需的 workspace 大小。參數(shù)說明如下參數(shù)名輸入/輸出描述使用說明數(shù)據(jù)類型數(shù)據(jù)格式維度(shape)非連續(xù) TensorselfaclTensor*輸入輸入的張量公式中的 x支持空 TensorFLOAT、FLOAT16、BFLOAT16ND0-8√lambdaclScalar*輸入輸入的標量公式中的輸入 λ數(shù)值要求大于等于 0FLOATND0-8√outaclTensor*輸出Softshrink 計算的出參支持空 Tensorshape 需要與 self 一致FLOAT、FLOAT16、BFLOAT16ND0-8√workspaceSizeuint64_t*輸出返回需要在 Device 側(cè)申請的 workspace 大小-----executoraclOpExecutor**輸出返回 op 執(zhí)行器包含了算子計算流程-----需要說明的是在 Atlas 推理系列產(chǎn)品、Atlas 訓練系列產(chǎn)品上數(shù)據(jù)類型僅支持 FLOAT、FLOAT16見 aclnnSoftshrink.md。返回值與錯誤碼兩個接口的返回值均為aclnnStatus狀態(tài)碼具體參見 aclnn 返回碼說明。第一段接口會完成入?yún)⑿ｒ炓韵聢鼍皶箦e返回碼錯誤碼描述ACLNN_ERR_PARAM_NULLPTR161001傳入的 self、lambd 或 out 是空指針ACLNN_ERR_PARAM_INVALID161002self、lambd 或 out 的數(shù)據(jù)類型不在支持的范圍之內(nèi)ACLNN_ERR_PARAM_INVALID161002self 和 out 的 shape 不一致ACLNN_ERR_PARAM_INVALID161002self 或 out 的維數(shù)大于 8ACLNN_ERR_PARAM_INVALID161002lambd 0以上校驗邏輯在源碼 op_api/aclnn_softshrink.cpp 的CheckParams流程中一一對應實現(xiàn)CheckNotNull檢查self、lambd、out是否為空指針失敗返回ACLNN_ERR_PARAM_NULLPTRCheckDtypeValid檢查輸入輸出數(shù)據(jù)類型是否在{DT_FLOAT, DT_FLOAT16, DT_BF16}支持列表內(nèi)并根據(jù)當前 SoC 版本判斷是否支持 BF16僅ASCEND910B~ASCEND910E區(qū)間內(nèi)的 SoC 支持否則返回ACLNN_ERR_PARAM_INVALIDCheckShape檢查self與out的 shape 一致且維數(shù)不超過MAX_SUPPORT_DIMS_NUMS8 維CheckLambdValue檢查lambd-ToFloat() 0CheckFormat在 Regbase 架構(gòu)下拒絕私有格式Private format的輸入輸出。當self為空 Tensor 時第一段接口直接返回workspaceSize 0并釋放執(zhí)行器屬于合法路徑。4.2 第二段接口aclnnSoftshrink該接口在指定 Stream 上執(zhí)行算子計算。參數(shù)說明如下參數(shù)名輸入/輸出描述workspace輸入在 Device 側(cè)申請的 workspace 內(nèi)存地址workspaceSize輸入在 Device 側(cè)申請的 workspace 大小由第一段接口aclnnSoftshrinkGetWorkspaceSize獲取executor輸入op 執(zhí)行器包含了算子計算流程stream輸入指定執(zhí)行任務的 Stream從源碼看第二段接口內(nèi)部調(diào)用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成真正的算子下發(fā)執(zhí)行屬于框架固定寫法。五、調(diào)用示例與完整調(diào)用流程以下示例代碼摘自 examples/test_aclnn_softshrink.cpp與接口文檔 aclnnSoftshrink.md 中的示例一致演示了在 NPU 上調(diào)用 Softshrink 算子的完整流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_softshrink.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) { // 固定寫法資源初始化 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); // 調(diào)用aclrtMalloc申請device側(cè)內(nèi)存 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); // 調(diào)用aclrtMemcpy將host側(cè)數(shù)據(jù)拷貝到device側(cè)內(nèi)存上 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); // 計算連續(xù)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]; } // 調(diào)用aclCreateTensor接口創(chuàng)建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手冊 // 根據(jù)自己的實際device填寫deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // check根據(jù)自己的需要處理 CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 構(gòu)造輸入與輸出需要根據(jù)API的接口自定義構(gòu)造 std::vectorint64_t selfShape {4, 2}; std::vectorint64_t outShape {4, 2}; void* selfDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclScalar* lambd nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}; std::vectorfloat outHostData {0, 0, 0, 0, 0, 0, 0, 0}; float lambdValue 0.5f; // 創(chuàng)建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 創(chuàng)建lambd aclTensor lambd aclCreateScalar(lambdValue, aclDataType::ACL_FLOAT); CHECK_RET(lambd ! nullptr, return ret); // 創(chuàng)建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 調(diào)用CANN算子庫API需要修改為具體的API名稱 uint64_t workspaceSize 0; aclOpExecutor* executor; // 調(diào)用aclnnSoftshrink第一段接口 ret aclnnSoftshrinkGetWorkspaceSize(self, lambd, out, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSoftshrinkGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根據(jù)第一段接口計算出的workspaceSize申請device內(nèi)存 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); } // 調(diào)用aclnnSoftshrink第二段接口 ret aclnnSoftshrink(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnSoftshrink failed. ERROR: %d\n, ret); return ret); // 4. 固定寫法同步等待任務執(zhí)行結(jié)束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 獲取輸出的值將device側(cè)內(nèi)存上的結(jié)果拷貝至host側(cè)需要根據(jù)具體API的接口定義修改 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需要根據(jù)具體API的接口定義修改 aclDestroyTensor(self); aclDestroyScalar(lambd); aclDestroyTensor(out); // 7. 釋放device資源需要根據(jù)具體API的接口定義修改 aclrtFree(selfDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }該示例按固定套路分為七個階段可歸納為“三步核心調(diào)用 四步固定收尾”資源初始化aclInit→aclrtSetDevice→aclrtCreateStream完成 ACL 運行環(huán)境與 Stream 創(chuàng)建構(gòu)造輸入輸出使用aclrtMalloc申請 Device 側(cè)內(nèi)存aclrtMemcpy將 Host 數(shù)據(jù)拷貝至 Device再通過aclCreateTensor/aclCreateScalar創(chuàng)建aclTensor與aclScalar。注意示例中l(wèi)ambdValue 0.5f與默認閾值一致輸入數(shù)據(jù){0.1, ..., 0.8}均落在[-0.5, 0.5]內(nèi)因此預期輸出全為 0兩段式調(diào)用先調(diào)用aclnnSoftshrinkGetWorkspaceSize獲取workspaceSize與executor按需用aclrtMalloc申請 workspace再調(diào)用aclnnSoftshrink執(zhí)行同步等待aclrtSynchronizeStream等待任務執(zhí)行結(jié)束結(jié)果回拷aclrtMemcpy將 Device 側(cè)結(jié)果拷貝回 Host 并打印釋放張量對象aclDestroyTensor、aclDestroyScalar釋放資源aclrtFree釋放 Device 內(nèi)存含 workspace、aclrtDestroyStream、aclrtResetDevice、aclFinalize。關于示例的具體編譯與執(zhí)行過程請參考倉庫中的 編譯與運行樣例 文檔在arch35目錄下還提供了面向 Ascend 950 的 test_aclnn_softshrink.cpp 版本。六、源碼級實現(xiàn)原理Softshrink 算子完整覆蓋了 CANN 算子開發(fā)的“API 層 → L0 算子層 → 算子定義層 → Kernel 層”全鏈路以下按調(diào)用鏈自頂向下分析。6.1 API 層入?yún)⑿ｒ炁c算子圖構(gòu)建op_api/aclnn_softshrink.cpp 實現(xiàn)了兩段式接口的對外封裝。第一段接口內(nèi)部除了前文所述的 5 步參數(shù)校驗外還完成了算子計算圖的構(gòu)建對輸入執(zhí)行l(wèi)0op::Contiguous將非連續(xù) Tensor 規(guī)整為連續(xù)內(nèi)存布局調(diào)用 L0 算子l0op::SoftShrink(selfContiguous, lambd-ToFloat(), executor)對結(jié)果執(zhí)行l(wèi)0op::Cast轉(zhuǎn)換到輸出要求的數(shù)據(jù)類型通過l0op::ViewCopy將結(jié)果拷貝到輸出out上從而支持非連續(xù)輸出 Tensor匯總得到workspaceSize并把uniqueExecutor的所有權轉(zhuǎn)移給調(diào)用方。這也解釋了接口文檔中“支持非連續(xù) Tensor√”的能力來源。6.2 L0 算子層AICore 路徑分發(fā)op_api/softshrink.cpp 定義了 L0 算子SoftShrink通過OP_TYPE_REGISTER(SoftShrink)注冊算子類型IsAiCoreSupport依據(jù)數(shù)據(jù)類型{DT_FLOAT, DT_FLOAT16, DT_BF16}判斷是否走 AICore 路徑支持時通過ADD_TO_LAUNCHER_LIST_AICORE(SoftShrink, ...)宏將算子加入任務隊列l(wèi)ambd作為算子屬性OP_ATTR隨任務下發(fā)輸出張量由executor-AllocTensor(input-GetViewShape(), input-GetDataType())自動分配shape 與數(shù)據(jù)類型與輸入保持一致。6.3 算子定義與形狀推導op_host/softshrink_def.cpp 使用OpDef框架注冊算子聲明輸入x、輸出yfp16/fp32/bf16ND 格式AutoContiguous以及可選屬性lambd默認 0.5f并配置ascend950的 AICore 配置動態(tài) shape/rank 支持、精度提升標志PrecisionReduceFlag(true)。形狀推導分為兩個層面op_host/softshrink_infershape.cpp輸出 shape 與輸入 shape 完全一致逐元素算子復用InferShape4Elewise通用實現(xiàn)op_graph/softshrink_graph_infer.cpp圖模式下輸出 dtype 與輸入 dtype 保持一致InferDataType直接透傳輸入數(shù)據(jù)類型。6.4 Kernel 層schMode 模板分發(fā)與 fp16/bf16 升精op_kernel/softshrink.cpp 是 AICore Kernel 的入口采用單 schMode 模板按數(shù)據(jù)類型分發(fā)schMode數(shù)據(jù)類型模板實例化說明0FP32Softshrinkfloat, 1, 0直通計算NEED_UPCAST01FP16Softshrinkhalf, 1, 1升精計算NEED_UPCAST12BF16Softshrinkbfloat16_t, 1, 1升精計算NEED_UPCAST1其中BUFFER_MODE固定為 1雙緩沖。源碼注釋指出一個重要實現(xiàn)細節(jié)fp16/bf16 路徑均先 Cast 到 fp32 計算再 Cast 回原精度目的是對齊 PyTorch CPU 的 fp16/bf16 路徑避免在λ ∈ {0.1, 0.3 ...}等無法被 fp16 精確表示的場景下因邊界判斷與減法誤差累積導致與 golden 數(shù)據(jù)偏差。與之配套的 Tiling 實現(xiàn) op_host/arch35/softshrink_tiling.cpp 按芯片核數(shù)AIV Core 數(shù)與 UB 內(nèi)存大小切分任務并根據(jù)數(shù)據(jù)類型計算每元素所需的 UB 緩沖區(qū)占用fp32 路徑為 36 字節(jié)/元素9 個 fp32 緩沖區(qū)fp16/bf16 升精路徑為 32 字節(jié)/元素4 * sizeof(half/bf16) 6 * sizeof(float)。lambd屬性在 Tiling 階段從context-GetAttrs()-GetFloat(0)讀取默認值為 0.5f。6.5 單元測試驗證activation/softshrink/tests/ut/op_host/op_api/test_aclnn_softshrink.cpp 覆蓋了多種數(shù)據(jù)類型組合的調(diào)用場景{5, 6, 7}形狀的 FLOAT16 輸入 FLOAT16 輸出TestGetWorkspaceSize返回ACL_SUCCESS同樣形狀的 FLOAT 輸入 FLOAT 輸出返回ACL_SUCCESSBF16 輸入 BF16 輸出調(diào)用成功預期在支持 BF16 的平臺通過DOUBLE 輸入預期返回ACLNN_ERR_PARAM_INVALID驗證了 dtype 校驗邏輯FLOAT16 輸入 FLOAT 輸出輸入輸出 dtype 不一致返回ACL_SUCCESS印證了內(nèi)部Cast環(huán)節(jié)對輸出數(shù)據(jù)類型的適配能力。測試通過OP_API_UT(aclnnSoftshrink, INPUT(tensor_desc, lambd), OUTPUT(out_tensor_desc))框架封裝了張量與標量描述含取值區(qū)間與精度閾值開發(fā)者可在接入環(huán)境后按需放開ut.TestPrecision()進行數(shù)值精度比對。七、約束與確定性計算約束說明算子 README 中明確指出 Softshrink 算子無額外約束activation/softshrink/README.md。實操中的限制均來自接口層面的入?yún)⑿ｒ灁?shù)據(jù)類型需為 fp16/fp32/bf16、維數(shù)不超過 8、lambd ≥ 0、輸出 shape 與輸入一致等詳見第四節(jié)錯誤碼表。確定性計算aclnnSoftshrink默認采用確定性實現(xiàn)見 aclnnSoftshrink.md即相同輸入在多次執(zhí)行下結(jié)果可復現(xiàn)這對調(diào)試與結(jié)果比對十分友好。關于確定性計算的更完整約定可參考倉庫文檔 determinism_compute.md。八、總結(jié)Softshrink 是 CANN ops-nn 中結(jié)構(gòu)完整、極具代表性的逐元素激活算子對外提供符合兩段式約定的aclnnSoftshrink接口內(nèi)部則貫通了參數(shù)校驗、Contiguous/Cast/ViewCopy 算子圖構(gòu)建、AICore 模板分發(fā)、Tiling 切分與雙緩沖 Kernel 的完整實現(xiàn)鏈路并通過 fp16/bf16 升精計算保證了與 PyTorch golden 的一致性。開發(fā)者無論是直接使用該算子還是將其作為學習 CANN 自定義算子開發(fā)范式的參考模板都可以從 activation/softshrink 目錄下的 README、接口文檔、示例與測試用例中獲得完整的閉環(huán)指引?！久赓M下載鏈接】ops-nn本項目是CANN提供的神經(jīng)網(wǎng)絡類計算算子庫實現(xiàn)網(wǎng)絡在NPU上加速計算。項目地址: https://gitcode.com/cann/ops-nn創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考