現(xiàn)、圖模式調(diào)用與源碼級(jí)原理)
CANN ops-math AssignSub 算子深度解析原址減法實(shí)現(xiàn)、圖模式調(diào)用與源碼級(jí)原理【免費(fèi)下載鏈接】ops-math本項(xiàng)目是CANN提供的數(shù)學(xué)類基礎(chǔ)計(jì)算算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 數(shù)學(xué)算子庫中的 AssignSub 算子為講解主體圍繞其功能定義、參數(shù)約束、產(chǎn)品適配與圖模式調(diào)用展開并結(jié)合倉庫內(nèi) op_graph / op_host / op_kernel / framework 四個(gè)層面的源碼實(shí)現(xiàn)幫助讀者既能在實(shí)際工程中正確調(diào)用 AssignSub又能理解該算子從算子原型、Tiling 計(jì)算到 NPU Kernel 執(zhí)行的完整鏈路。一、算子概述與核心公式AssignSub 是 CANN ops-math 提供的原址in-place減法算子在已有張量var的基礎(chǔ)上直接完成減法計(jì)算結(jié)果寫回var所在的內(nèi)存地址不需要額外分配輸出張量。這一點(diǎn)與普通的 elementwise 減法算子如sub會(huì)生成新的輸出張量有本質(zhì)區(qū)別因此它天然適合用作訓(xùn)練/推理圖中的變量Variable更新節(jié)點(diǎn)。其數(shù)學(xué)定義非常簡(jiǎn)單文檔給出的計(jì)算公式為out_i var_i - value_i其中var_i是被更新的原張量元素value_i是要減去的值out_i與var_i共享同一塊地址。這一行為在算子原型注釋中也有明確描述Updates var by subtracting value from it. This operation outputs var after the update is done且明確指出該算子與 TensorFlow 的 AssignSub 算子兼容見 math/assign_sub/op_graph/assign_sub_proto.h。二、產(chǎn)品支持情況當(dāng)前倉庫中的 AssignSub 實(shí)現(xiàn)適配了 Ascend 主流訓(xùn)練/推理產(chǎn)品線支持情況如下表產(chǎn)品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 訓(xùn)練系列產(chǎn)品 / Atlas A3 推理系列產(chǎn)品√Atlas A2 訓(xùn)練系列產(chǎn)品 / Atlas A2 推理系列產(chǎn)品√Atlas 200I/500 A2 推理產(chǎn)品×Atlas 推理系列產(chǎn)品√Atlas 訓(xùn)練系列產(chǎn)品√從源碼角度可以印證上述適配范圍算子的 AICore 配置在 math/assign_sub/op_host/assign_sub_def.cpp 中通過this-AICore().AddConfig(ascend950, aicoreConfig)和AddConfig(ascend350, aicoreConfig)注冊(cè)對(duì)應(yīng)支持 asceng 950 與 350 系列平臺(tái)Tiling 與 Kernel 實(shí)現(xiàn)則位于arch35目錄math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp、math/assign_sub/op_kernel/arch35/assign_sub_dag.h。需要說明的是產(chǎn)品支持矩陣以當(dāng)前倉庫實(shí)現(xiàn)為準(zhǔn)具體到目標(biāo)環(huán)境時(shí)仍建議以對(duì)應(yīng)版本的《CANN 算子支持矩陣》做最終確認(rèn)。三、參數(shù)說明AssignSub 算子共包含兩個(gè)輸入、一個(gè)輸出和一個(gè)可選屬性完整參數(shù)定義如下表數(shù)據(jù)范圍與算子原型REG_OP(AssignSub)的聲明嚴(yán)格一致見 math/assign_sub/op_graph/assign_sub_proto.h參數(shù)名輸入/輸出/屬性描述數(shù)據(jù)類型數(shù)據(jù)格式var輸入輸入張量公式中的var_iBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvalue輸入輸入張量表示要從var_i上減去的值公式中的value_iBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvar輸出輸出張量公式中的out_i與var_i同地址BFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDuse_locking屬性可選是否使用鎖來保護(hù)更新操作bool—在使用時(shí)需要注意以下幾點(diǎn)數(shù)據(jù)類型一致性var、value與輸出var必須為同一數(shù)據(jù)類型。這一約束不僅在算子原型中體現(xiàn)還會(huì)在 Tiling 階段被顯式校驗(yàn)AssignSubTiling::CheckDtype()會(huì)逐一取出三個(gè)張量的 dtype 并比較若不一致會(huì)記錄The dtypes of var, value and output must be the same錯(cuò)誤并返回失敗見 math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp。Shape 一致性var、value、輸出三者的 shape 也必須完全一致CheckShape()中通過varShape ! valueShape || varShape ! outputShape進(jìn)行校驗(yàn)同一文件 L74-L96。這與該算子使用逐元素elewise推斷規(guī)則是一致的。use_locking 屬性可選 bool默認(rèn)值為False。原型注釋說明當(dāng)其為True時(shí)減法更新操作將由鎖保護(hù)避免多線程并發(fā)下的競(jìng)爭(zhēng)為False時(shí)行為未定義但可以減少鎖競(jìng)爭(zhēng)帶來的開銷。當(dāng)前倉庫的 Tiling/Kernel 實(shí)現(xiàn)中未對(duì)鎖語義做額外處理屬性主要隨圖下發(fā)給框架側(cè)。約束說明原文檔明確約束說明無即在已支持的 dtype / ND 格式范圍內(nèi)沒有額外使用約束。四、圖模式調(diào)用單算子構(gòu)圖完整樣例AssignSub 支持圖模式調(diào)用即通過算子 IR 構(gòu)圖GE Graph 接口的方式在計(jì)算圖中插入 AssignSub 節(jié)點(diǎn)后運(yùn)行。倉庫在 math/assign_sub/examples/test_geir_assign_sub.cpp 中提供了完整的可運(yùn)行樣例下面拆解其關(guān)鍵流程。4.1 定義算子 IR 與構(gòu)圖首先通過ge::op::AssignSub創(chuàng)建算子實(shí)例并為兩個(gè)輸入分別創(chuàng)建Data占位節(jié)點(diǎn)、設(shè)置輸入/輸出 TensorDescauto assignSub1 op::AssignSub(assignSub1); std::vectorint64_t xShape {1, 16}; std::vectorint64_t yShape {1, 16}; ADD_INPUT(1, var, inDtype, yShape); // 輸入 varshape 為 {1, 16} ADD_INPUT(2, value, inDtype, xShape); // 輸入 valueshape 為 {1, 16} ADD_INPUT_ATTR(use_locking, false); // 設(shè)置屬性 use_locking false ADD_OUTPUT(1, var, inDtype, yShape); // 輸出 var與輸入 var 同 shape樣例中的宏展開后會(huì)為每個(gè)輸入調(diào)用assignSub1.set_input_var(...)/assignSub1.set_input_value(...)建立數(shù)據(jù)流邊并通過assignSub1.update_output_desc_var(...)聲明輸出描述。這里示例使用DT_FLOAT16FP16作為輸入數(shù)據(jù)類型shape 取{1, 16}。4.2 初始化 GE 與會(huì)話運(yùn)行構(gòu)圖完成后樣例依次完成 GE 初始化、建圖、執(zhí)行三步// 1. 初始化 GE指定 deviceId 與圖運(yùn)行模式 std::mapAscendString, AscendString global_options {{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}}; Status ret ge::GEInitialize(global_options); // 2. 創(chuàng)建 Session 并添加計(jì)算圖 ge::Session* session new Session(build_options); uint32_t graph_id 0; ret session-AddGraph(graph_id, graph, graph_options); // 3. 運(yùn)行圖輸入/輸出均為 ge::Tensor std::vectorge::Tensor output; ret session-RunGraph(graph_id, input, output);其中g(shù)e.graphRunMode1表示按圖運(yùn)行模式執(zhí)行運(yùn)行結(jié)束后樣例還會(huì)通過aclgrphDumpGraph(graph, ./dump, ...)將圖結(jié)構(gòu) dump 到本地便于調(diào)試查看構(gòu)圖是否正確。4.3 數(shù)據(jù)生成與結(jié)果落盤樣例內(nèi)置了GenOnesData輔助函數(shù)用于構(gòu)造輸入數(shù)據(jù)根據(jù) shape 計(jì)算元素個(gè)數(shù)結(jié)合GetDataTypeSize換算字節(jié)數(shù)將數(shù)據(jù)全部初始化為固定值后封裝為ge::Tensor。運(yùn)行完成后輸入與輸出張量會(huì)分別寫入tc_ge_irrun_test_0008_npu_input_{i}.bin與tc_ge_irrun_test_0008_npu_output_{i}.bin并在終端逐元素打印輸出結(jié)果方便與預(yù)期var - value結(jié)果比對(duì)。從該樣例可以看到圖模式調(diào)用的標(biāo)準(zhǔn)姿勢(shì)op::AssignSub構(gòu)圖 →GEInitialize→Session::AddGraph→RunGraph→GEFinalize這套流程同樣適用于 ops-math 中其它算子。五、源碼級(jí)原理從原型注冊(cè)到 NPU Kernel 的完整鏈路了解了調(diào)用方式后再深入一層看 AssignSub 在倉庫內(nèi)部是如何被定義、分派與執(zhí)行的。整條鏈路可以概括為算子原型 → Host 側(cè)定義/推斷 → Tiling 計(jì)算 → Kernel 執(zhí)行四個(gè)階段。5.1 算子原型op_graphmath/assign_sub/op_graph/assign_sub_proto.h 使用REG_OP宏聲明算子原型REG_OP(AssignSub) .INPUT(var, TensorType::NumberType()) .INPUT(value, TensorType::NumberType()) .OUTPUT(var, TensorType::NumberType()) .ATTR(use_locking, Bool, false) .OP_END_FACTORY_REG(AssignSub)注意輸入與輸出同名var這正是原址更新語義在原型層面的體現(xiàn)輸出與輸入var指向同一塊數(shù)據(jù)。5.2 Host 側(cè)算子定義op_hostmath/assign_sub/op_host/assign_sub_def.cpp 通過OpDef注冊(cè)了算子的運(yùn)行時(shí)定義包括輸入/輸出的數(shù)據(jù)類型白名單DT_BF16, DT_FLOAT16, DT_FLOAT, DT_INT8, DT_INT32, DT_INT64, DT_UINT8數(shù)據(jù)格式均為FORMAT_ND屬性u(píng)se_lockingOPTIONAL類型默認(rèn)falseAICore 配置開啟動(dòng)態(tài)編譯DynamicCompileStaticFlag(true)、動(dòng)態(tài) rank 與動(dòng)態(tài) shape 支持DynamicRankSupportFlag(true)/DynamicShapeSupportFlag(true)并將 Kernel 文件指向assign_sub_apt。Shape 推斷方面math/assign_sub/op_host/assign_sub_infershape.cpp 直接復(fù)用逐元素通用推斷邏輯Ops::Base::InferShape4Elewise——這再次印證了 AssignSub 在 shape 語義上與 elementwise 算子一致輸入輸出 shape 相同。對(duì)應(yīng)的單測(cè) math/assign_sub/tests/ut/op_host/test_assign_sub_infershape.cpp 驗(yàn)證了 FP16、ND 格式下輸入輸出 shape 推斷的結(jié)果。5.3 Tiling 計(jì)算op_host/arch35在 NPU 上執(zhí)行前Host 側(cè)需要先完成 Tiling數(shù)據(jù)分片計(jì)算。math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp 實(shí)現(xiàn)了AssignSubTiling::RunTiling()其執(zhí)行順序?yàn)閐type 校驗(yàn)三個(gè)張量 dtype 必須一致CheckDtypeshape 校驗(yàn)三個(gè)張量 shape 必須一致CheckShape按 dtype 選擇分片策略通過Ops::Base::ElewiseBaseTiling::DoTilingAssignSubOpT::OpDag()生成基礎(chǔ) Tiling 數(shù)據(jù)baseTiling設(shè)置 Tiling KeySetTilingData()依據(jù)輸出 dtype 寫入對(duì)應(yīng)的 Tiling KeyFP16/BF16/FP32/INT8/INT32/INT64/UINT8 分別映射為101/102/103/104/105/106/107申請(qǐng) Workspace設(shè)置系統(tǒng) Workspace 大小SYS_WORKSPACE_SIZE 16MB供 Kernel 側(cè)使用。Tiling 階段還通過TilingPrepareForAssignSub從平臺(tái)信息中讀取 AIV 核數(shù)與 UB 內(nèi)存大小GetCoreNumAiv()/GetCoreMemSize(CoreMemType::UB)用于決定分塊粒度與核數(shù)分配。5.4 Kernel 執(zhí)行op_kernelNPU 側(cè)入口在 math/assign_sub/op_kernel/assign_sub_apt.cpp其assign_sub內(nèi)核函數(shù)以GM_ADDR x1/x2/y分別接收var、value與輸出地址核心邏輯是一個(gè)按 Tiling Key 分派的 elementwise 調(diào)度器if (TILING_KEY_IS(101UL)) { ElementwiseSch0UL, AssignSubOphalf::OpDag sch((tilingData.baseTiling), pipe); sch.Init(x1, x2, y); sch.Process(); } else if (TILING_KEY_IS(103UL)) { ElementwiseSch0UL, AssignSubOpfloat::OpDag sch(...); ... }真正執(zhí)行減法運(yùn)算的計(jì)算圖在 math/assign_sub/op_kernel/arch35/assign_sub_dag.h 中定義通過CopyIn0 / CopyIn1將兩個(gè)輸入從 GM 搬入 UBOpCopyIn自定義向量減法節(jié)點(diǎn)SubCustom使用寄存器指令Reg::Sub(vregOutput, vregInput1, vregInput2, mask)完成a - b其中向量寬度按數(shù)據(jù)類型VECTOR_REG_WIDTH / sizeof(T)計(jì)算并通過CeilDivision(count, vl)分循環(huán)處理剩余元素CopyOut將結(jié)果寫回 GM 輸出地址最終由DAGSchOutputs組成可調(diào)度的計(jì)算 DAG。Tiling 數(shù)據(jù)本身在 math/assign_sub/op_kernel/arch35/assign_sub_tiling_struct.h 中定義為AssignSubTilingData { EleBaseTilingData baseTiling; }即直接復(fù)用 elementwise 基礎(chǔ) Tiling 結(jié)構(gòu)這也是 AssignSub 與倉庫中其他 elementwise 算子共享同一套調(diào)度框架atvoss/elewise的體現(xiàn)。六、框架適配TensorFlow 算子映射除了圖模式直接構(gòu)圖AssignSub 還通過框架插件完成了與第三方框架的對(duì)接。math/assign_sub/framework/assign_sub_tf_plugin.cpp 中使用REGISTER_CUSTOM_OP將自定義算子與 TensorFlow 側(cè)算子建立映射REGISTER_CUSTOM_OP(AssignSub) .FrameworkType(TENSORFLOW) .OriginOpType(std::vectorge::AscendString{ge::AscendString(AssignSub), ge::AscendString(AssignSubVariableOp)}) .ParseParamsByOperatorFn(AutoMappingByOpFn) .ImplyType(ImplyType::TVM);這里同時(shí)映射了 TensorFlow 的AssignSub與AssignSubVariableOp兩個(gè)原始算子類型后者對(duì)應(yīng)tf.Variable的assign_sub方法參數(shù)通過AutoMappingByOpFn自動(dòng)映射無需手寫逐字段轉(zhuǎn)換邏輯。這意味著在 PyTorch/TensorFlow 生態(tài)中以var.assign_sub(value)這類語義發(fā)起的減法更新可以無縫落到 NPU 上的 AssignSub 內(nèi)核。七、測(cè)試與驗(yàn)證倉庫為 AssignSub 提供了 Host 側(cè)單測(cè)與 ST 用例可作為功能正確性的驗(yàn)證參考Infershape 單測(cè)math/assign_sub/tests/ut/op_host/test_assign_sub_infershape.cpp 構(gòu)造{ -1 }動(dòng)態(tài) shape 的 FP16 輸入斷言推斷輸出 shape 與預(yù)期一致且返回GRAPH_SUCCESSTiling 單測(cè)math/assign_sub/tests/ut/op_host/arch35/test_assign_sub_tiling.cpp 覆蓋 arch35 平臺(tái)上的 Tiling 計(jì)算ST 用例math/assign_sub/tests/st/arch35/ttk_kernel_assign_sub_st.csv 定義了 Kernel 級(jí) ST 用例矩陣配合 math/assign_sub/tests/assets/golden.py 的 golden 數(shù)據(jù)生成腳本可批量校驗(yàn)內(nèi)核輸出。八、小結(jié)AssignSub 雖然只是一個(gè)單輸入對(duì)的原址減法算子但其倉庫實(shí)現(xiàn)完整覆蓋了從算子原型op_graph、Host 定義與推斷op_host、Tiling 計(jì)算arch35、NPU Kernelop_kernel到框架適配framework的全棧鏈路是理解 ops-math 中 elementwise 類算子實(shí)現(xiàn)范式的極佳樣本。實(shí)際使用時(shí)只需記住三個(gè)要點(diǎn)輸入輸出 dtype 與 shape 必須一致、數(shù)據(jù)格式為 ND、通過圖模式GE IR或 TensorFlow 插件路徑發(fā)起調(diào)用即可在支持的 Ascend 產(chǎn)品上正確完成原址減法更新。【免費(fèi)下載鏈接】ops-math本項(xiàng)目是CANN提供的數(shù)學(xué)類基礎(chǔ)計(jì)算算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-math創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考 CANN ops-math AssignSub 算子深度解析：原址減法實(shí)現(xiàn)、圖模式調(diào)用與源碼級(jí)原理
CANN ops-math AssignSub 算子深度解析原址減法實(shí)現(xiàn)、圖模式調(diào)用與源碼級(jí)原理【免費(fèi)下載鏈接】ops-math本項(xiàng)目是CANN提供的數(shù)學(xué)類基礎(chǔ)計(jì)算算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 數(shù)學(xué)算子庫中的 AssignSub 算子為講解主體圍繞其功能定義、參數(shù)約束、產(chǎn)品適配與圖模式調(diào)用展開并結(jié)合倉庫內(nèi) op_graph / op_host / op_kernel / framework 四個(gè)層面的源碼實(shí)現(xiàn)幫助讀者既能在實(shí)際工程中正確調(diào)用 AssignSub又能理解該算子從算子原型、Tiling 計(jì)算到 NPU Kernel 執(zhí)行的完整鏈路。一、算子概述與核心公式AssignSub 是 CANN ops-math 提供的原址in-place減法算子在已有張量var的基礎(chǔ)上直接完成減法計(jì)算結(jié)果寫回var所在的內(nèi)存地址不需要額外分配輸出張量。這一點(diǎn)與普通的 elementwise 減法算子如sub會(huì)生成新的輸出張量有本質(zhì)區(qū)別因此它天然適合用作訓(xùn)練/推理圖中的變量Variable更新節(jié)點(diǎn)。其數(shù)學(xué)定義非常簡(jiǎn)單文檔給出的計(jì)算公式為out_i var_i - value_i其中var_i是被更新的原張量元素value_i是要減去的值out_i與var_i共享同一塊地址。這一行為在算子原型注釋中也有明確描述Updates var by subtracting value from it. This operation outputs var after the update is done且明確指出該算子與 TensorFlow 的 AssignSub 算子兼容見 math/assign_sub/op_graph/assign_sub_proto.h。二、產(chǎn)品支持情況當(dāng)前倉庫中的 AssignSub 實(shí)現(xiàn)適配了 Ascend 主流訓(xùn)練/推理產(chǎn)品線支持情況如下表產(chǎn)品是否支持Ascend 950PR / Ascend 950DT√Atlas A3 訓(xùn)練系列產(chǎn)品 / Atlas A3 推理系列產(chǎn)品√Atlas A2 訓(xùn)練系列產(chǎn)品 / Atlas A2 推理系列產(chǎn)品√Atlas 200I/500 A2 推理產(chǎn)品×Atlas 推理系列產(chǎn)品√Atlas 訓(xùn)練系列產(chǎn)品√從源碼角度可以印證上述適配范圍算子的 AICore 配置在 math/assign_sub/op_host/assign_sub_def.cpp 中通過this-AICore().AddConfig(ascend950, aicoreConfig)和AddConfig(ascend350, aicoreConfig)注冊(cè)對(duì)應(yīng)支持 asceng 950 與 350 系列平臺(tái)Tiling 與 Kernel 實(shí)現(xiàn)則位于arch35目錄math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp、math/assign_sub/op_kernel/arch35/assign_sub_dag.h。需要說明的是產(chǎn)品支持矩陣以當(dāng)前倉庫實(shí)現(xiàn)為準(zhǔn)具體到目標(biāo)環(huán)境時(shí)仍建議以對(duì)應(yīng)版本的《CANN 算子支持矩陣》做最終確認(rèn)。三、參數(shù)說明AssignSub 算子共包含兩個(gè)輸入、一個(gè)輸出和一個(gè)可選屬性完整參數(shù)定義如下表數(shù)據(jù)范圍與算子原型REG_OP(AssignSub)的聲明嚴(yán)格一致見 math/assign_sub/op_graph/assign_sub_proto.h參數(shù)名輸入/輸出/屬性描述數(shù)據(jù)類型數(shù)據(jù)格式var輸入輸入張量公式中的var_iBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvalue輸入輸入張量表示要從var_i上減去的值公式中的value_iBFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDvar輸出輸出張量公式中的out_i與var_i同地址BFLOAT16、FLOAT16、FLOAT、INT8、INT32、INT64、UINT8NDuse_locking屬性可選是否使用鎖來保護(hù)更新操作bool—在使用時(shí)需要注意以下幾點(diǎn)數(shù)據(jù)類型一致性var、value與輸出var必須為同一數(shù)據(jù)類型。這一約束不僅在算子原型中體現(xiàn)還會(huì)在 Tiling 階段被顯式校驗(yàn)AssignSubTiling::CheckDtype()會(huì)逐一取出三個(gè)張量的 dtype 并比較若不一致會(huì)記錄The dtypes of var, value and output must be the same錯(cuò)誤并返回失敗見 math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp。Shape 一致性var、value、輸出三者的 shape 也必須完全一致CheckShape()中通過varShape ! valueShape || varShape ! outputShape進(jìn)行校驗(yàn)同一文件 L74-L96。這與該算子使用逐元素elewise推斷規(guī)則是一致的。use_locking 屬性可選 bool默認(rèn)值為False。原型注釋說明當(dāng)其為True時(shí)減法更新操作將由鎖保護(hù)避免多線程并發(fā)下的競(jìng)爭(zhēng)為False時(shí)行為未定義但可以減少鎖競(jìng)爭(zhēng)帶來的開銷。當(dāng)前倉庫的 Tiling/Kernel 實(shí)現(xiàn)中未對(duì)鎖語義做額外處理屬性主要隨圖下發(fā)給框架側(cè)。約束說明原文檔明確約束說明無即在已支持的 dtype / ND 格式范圍內(nèi)沒有額外使用約束。四、圖模式調(diào)用單算子構(gòu)圖完整樣例AssignSub 支持圖模式調(diào)用即通過算子 IR 構(gòu)圖GE Graph 接口的方式在計(jì)算圖中插入 AssignSub 節(jié)點(diǎn)后運(yùn)行。倉庫在 math/assign_sub/examples/test_geir_assign_sub.cpp 中提供了完整的可運(yùn)行樣例下面拆解其關(guān)鍵流程。4.1 定義算子 IR 與構(gòu)圖首先通過ge::op::AssignSub創(chuàng)建算子實(shí)例并為兩個(gè)輸入分別創(chuàng)建Data占位節(jié)點(diǎn)、設(shè)置輸入/輸出 TensorDescauto assignSub1 op::AssignSub(assignSub1); std::vectorint64_t xShape {1, 16}; std::vectorint64_t yShape {1, 16}; ADD_INPUT(1, var, inDtype, yShape); // 輸入 varshape 為 {1, 16} ADD_INPUT(2, value, inDtype, xShape); // 輸入 valueshape 為 {1, 16} ADD_INPUT_ATTR(use_locking, false); // 設(shè)置屬性 use_locking false ADD_OUTPUT(1, var, inDtype, yShape); // 輸出 var與輸入 var 同 shape樣例中的宏展開后會(huì)為每個(gè)輸入調(diào)用assignSub1.set_input_var(...)/assignSub1.set_input_value(...)建立數(shù)據(jù)流邊并通過assignSub1.update_output_desc_var(...)聲明輸出描述。這里示例使用DT_FLOAT16FP16作為輸入數(shù)據(jù)類型shape 取{1, 16}。4.2 初始化 GE 與會(huì)話運(yùn)行構(gòu)圖完成后樣例依次完成 GE 初始化、建圖、執(zhí)行三步// 1. 初始化 GE指定 deviceId 與圖運(yùn)行模式 std::mapAscendString, AscendString global_options {{ge.exec.deviceId, 0}, {ge.graphRunMode, 1}}; Status ret ge::GEInitialize(global_options); // 2. 創(chuàng)建 Session 并添加計(jì)算圖 ge::Session* session new Session(build_options); uint32_t graph_id 0; ret session-AddGraph(graph_id, graph, graph_options); // 3. 運(yùn)行圖輸入/輸出均為 ge::Tensor std::vectorge::Tensor output; ret session-RunGraph(graph_id, input, output);其中g(shù)e.graphRunMode1表示按圖運(yùn)行模式執(zhí)行運(yùn)行結(jié)束后樣例還會(huì)通過aclgrphDumpGraph(graph, ./dump, ...)將圖結(jié)構(gòu) dump 到本地便于調(diào)試查看構(gòu)圖是否正確。4.3 數(shù)據(jù)生成與結(jié)果落盤樣例內(nèi)置了GenOnesData輔助函數(shù)用于構(gòu)造輸入數(shù)據(jù)根據(jù) shape 計(jì)算元素個(gè)數(shù)結(jié)合GetDataTypeSize換算字節(jié)數(shù)將數(shù)據(jù)全部初始化為固定值后封裝為ge::Tensor。運(yùn)行完成后輸入與輸出張量會(huì)分別寫入tc_ge_irrun_test_0008_npu_input_{i}.bin與tc_ge_irrun_test_0008_npu_output_{i}.bin并在終端逐元素打印輸出結(jié)果方便與預(yù)期var - value結(jié)果比對(duì)。從該樣例可以看到圖模式調(diào)用的標(biāo)準(zhǔn)姿勢(shì)op::AssignSub構(gòu)圖 →GEInitialize→Session::AddGraph→RunGraph→GEFinalize這套流程同樣適用于 ops-math 中其它算子。五、源碼級(jí)原理從原型注冊(cè)到 NPU Kernel 的完整鏈路了解了調(diào)用方式后再深入一層看 AssignSub 在倉庫內(nèi)部是如何被定義、分派與執(zhí)行的。整條鏈路可以概括為算子原型 → Host 側(cè)定義/推斷 → Tiling 計(jì)算 → Kernel 執(zhí)行四個(gè)階段。5.1 算子原型op_graphmath/assign_sub/op_graph/assign_sub_proto.h 使用REG_OP宏聲明算子原型REG_OP(AssignSub) .INPUT(var, TensorType::NumberType()) .INPUT(value, TensorType::NumberType()) .OUTPUT(var, TensorType::NumberType()) .ATTR(use_locking, Bool, false) .OP_END_FACTORY_REG(AssignSub)注意輸入與輸出同名var這正是原址更新語義在原型層面的體現(xiàn)輸出與輸入var指向同一塊數(shù)據(jù)。5.2 Host 側(cè)算子定義op_hostmath/assign_sub/op_host/assign_sub_def.cpp 通過OpDef注冊(cè)了算子的運(yùn)行時(shí)定義包括輸入/輸出的數(shù)據(jù)類型白名單DT_BF16, DT_FLOAT16, DT_FLOAT, DT_INT8, DT_INT32, DT_INT64, DT_UINT8數(shù)據(jù)格式均為FORMAT_ND屬性u(píng)se_lockingOPTIONAL類型默認(rèn)falseAICore 配置開啟動(dòng)態(tài)編譯DynamicCompileStaticFlag(true)、動(dòng)態(tài) rank 與動(dòng)態(tài) shape 支持DynamicRankSupportFlag(true)/DynamicShapeSupportFlag(true)并將 Kernel 文件指向assign_sub_apt。Shape 推斷方面math/assign_sub/op_host/assign_sub_infershape.cpp 直接復(fù)用逐元素通用推斷邏輯Ops::Base::InferShape4Elewise——這再次印證了 AssignSub 在 shape 語義上與 elementwise 算子一致輸入輸出 shape 相同。對(duì)應(yīng)的單測(cè) math/assign_sub/tests/ut/op_host/test_assign_sub_infershape.cpp 驗(yàn)證了 FP16、ND 格式下輸入輸出 shape 推斷的結(jié)果。5.3 Tiling 計(jì)算op_host/arch35在 NPU 上執(zhí)行前Host 側(cè)需要先完成 Tiling數(shù)據(jù)分片計(jì)算。math/assign_sub/op_host/arch35/assign_sub_tiling_arch35.cpp 實(shí)現(xiàn)了AssignSubTiling::RunTiling()其執(zhí)行順序?yàn)閐type 校驗(yàn)三個(gè)張量 dtype 必須一致CheckDtypeshape 校驗(yàn)三個(gè)張量 shape 必須一致CheckShape按 dtype 選擇分片策略通過Ops::Base::ElewiseBaseTiling::DoTilingAssignSubOpT::OpDag()生成基礎(chǔ) Tiling 數(shù)據(jù)baseTiling設(shè)置 Tiling KeySetTilingData()依據(jù)輸出 dtype 寫入對(duì)應(yīng)的 Tiling KeyFP16/BF16/FP32/INT8/INT32/INT64/UINT8 分別映射為101/102/103/104/105/106/107申請(qǐng) Workspace設(shè)置系統(tǒng) Workspace 大小SYS_WORKSPACE_SIZE 16MB供 Kernel 側(cè)使用。Tiling 階段還通過TilingPrepareForAssignSub從平臺(tái)信息中讀取 AIV 核數(shù)與 UB 內(nèi)存大小GetCoreNumAiv()/GetCoreMemSize(CoreMemType::UB)用于決定分塊粒度與核數(shù)分配。5.4 Kernel 執(zhí)行op_kernelNPU 側(cè)入口在 math/assign_sub/op_kernel/assign_sub_apt.cpp其assign_sub內(nèi)核函數(shù)以GM_ADDR x1/x2/y分別接收var、value與輸出地址核心邏輯是一個(gè)按 Tiling Key 分派的 elementwise 調(diào)度器if (TILING_KEY_IS(101UL)) { ElementwiseSch0UL, AssignSubOphalf::OpDag sch((tilingData.baseTiling), pipe); sch.Init(x1, x2, y); sch.Process(); } else if (TILING_KEY_IS(103UL)) { ElementwiseSch0UL, AssignSubOpfloat::OpDag sch(...); ... }真正執(zhí)行減法運(yùn)算的計(jì)算圖在 math/assign_sub/op_kernel/arch35/assign_sub_dag.h 中定義通過CopyIn0 / CopyIn1將兩個(gè)輸入從 GM 搬入 UBOpCopyIn自定義向量減法節(jié)點(diǎn)SubCustom使用寄存器指令Reg::Sub(vregOutput, vregInput1, vregInput2, mask)完成a - b其中向量寬度按數(shù)據(jù)類型VECTOR_REG_WIDTH / sizeof(T)計(jì)算并通過CeilDivision(count, vl)分循環(huán)處理剩余元素CopyOut將結(jié)果寫回 GM 輸出地址最終由DAGSchOutputs組成可調(diào)度的計(jì)算 DAG。Tiling 數(shù)據(jù)本身在 math/assign_sub/op_kernel/arch35/assign_sub_tiling_struct.h 中定義為AssignSubTilingData { EleBaseTilingData baseTiling; }即直接復(fù)用 elementwise 基礎(chǔ) Tiling 結(jié)構(gòu)這也是 AssignSub 與倉庫中其他 elementwise 算子共享同一套調(diào)度框架atvoss/elewise的體現(xiàn)。六、框架適配TensorFlow 算子映射除了圖模式直接構(gòu)圖AssignSub 還通過框架插件完成了與第三方框架的對(duì)接。math/assign_sub/framework/assign_sub_tf_plugin.cpp 中使用REGISTER_CUSTOM_OP將自定義算子與 TensorFlow 側(cè)算子建立映射REGISTER_CUSTOM_OP(AssignSub) .FrameworkType(TENSORFLOW) .OriginOpType(std::vectorge::AscendString{ge::AscendString(AssignSub), ge::AscendString(AssignSubVariableOp)}) .ParseParamsByOperatorFn(AutoMappingByOpFn) .ImplyType(ImplyType::TVM);這里同時(shí)映射了 TensorFlow 的AssignSub與AssignSubVariableOp兩個(gè)原始算子類型后者對(duì)應(yīng)tf.Variable的assign_sub方法參數(shù)通過AutoMappingByOpFn自動(dòng)映射無需手寫逐字段轉(zhuǎn)換邏輯。這意味著在 PyTorch/TensorFlow 生態(tài)中以var.assign_sub(value)這類語義發(fā)起的減法更新可以無縫落到 NPU 上的 AssignSub 內(nèi)核。七、測(cè)試與驗(yàn)證倉庫為 AssignSub 提供了 Host 側(cè)單測(cè)與 ST 用例可作為功能正確性的驗(yàn)證參考Infershape 單測(cè)math/assign_sub/tests/ut/op_host/test_assign_sub_infershape.cpp 構(gòu)造{ -1 }動(dòng)態(tài) shape 的 FP16 輸入斷言推斷輸出 shape 與預(yù)期一致且返回GRAPH_SUCCESSTiling 單測(cè)math/assign_sub/tests/ut/op_host/arch35/test_assign_sub_tiling.cpp 覆蓋 arch35 平臺(tái)上的 Tiling 計(jì)算ST 用例math/assign_sub/tests/st/arch35/ttk_kernel_assign_sub_st.csv 定義了 Kernel 級(jí) ST 用例矩陣配合 math/assign_sub/tests/assets/golden.py 的 golden 數(shù)據(jù)生成腳本可批量校驗(yàn)內(nèi)核輸出。八、小結(jié)AssignSub 雖然只是一個(gè)單輸入對(duì)的原址減法算子但其倉庫實(shí)現(xiàn)完整覆蓋了從算子原型op_graph、Host 定義與推斷op_host、Tiling 計(jì)算arch35、NPU Kernelop_kernel到框架適配framework的全棧鏈路是理解 ops-math 中 elementwise 類算子實(shí)現(xiàn)范式的極佳樣本。實(shí)際使用時(shí)只需記住三個(gè)要點(diǎn)輸入輸出 dtype 與 shape 必須一致、數(shù)據(jù)格式為 ND、通過圖模式GE IR或 TensorFlow 插件路徑發(fā)起調(diào)用即可在支持的 Ascend 產(chǎn)品上正確完成原址減法更新。【免費(fèi)下載鏈接】ops-math本項(xiàng)目是CANN提供的數(shù)學(xué)類基礎(chǔ)計(jì)算算子庫實(shí)現(xiàn)網(wǎng)絡(luò)在NPU上加速計(jì)算。項(xiàng)目地址: https://gitcode.com/cann/ops-math創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考