現(xiàn)與 aclnn 調(diào)用實(shí)戰(zhàn))
CANN ops-math IsinPartV1 算子深度解析基于 CatSort 的 isin 核心部件實(shí)現(xiàn)與 aclnn 調(diào)用實(shí)戰(zhà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-mathIsinPartV1 是 CANN ops-math 倉庫 experimental/math/isin_part_v1 目錄下提供的一個(gè)實(shí)驗(yàn)性數(shù)學(xué)算子承擔(dān) isin 功能中判斷成員關(guān)系的核心部件與 Cat、Sort 算子配合即可拼裝出完整的 isin 語義。本文以該算子 README 為主體結(jié)合倉庫內(nèi) op_host、op_kernel、config 與調(diào)用示例源碼完整講解 IsinPartV1 的算子語義、參數(shù)規(guī)格、算法流程、底層實(shí)現(xiàn)與 aclnn 單算子調(diào)用方法幫助讀者在 Atlas A2 訓(xùn)練系列產(chǎn)品 / Atlas 800I A2 推理產(chǎn)品上獨(dú)立完成該算子的理解、構(gòu)建與驗(yàn)證。一、算子定位isin 的核心部件設(shè)計(jì)isin算子的功能是判斷給定張量elements的每個(gè)元素是否出現(xiàn)在另一個(gè)張量test_elements中并返回一個(gè)與elements形狀一致的布爾張量。在 CANN ops-math 中這一功能并未做成一個(gè)整體算子而是被拆解為Cat拼接→ Sort排序→ IsinPartV1成員判定三步IsinPartV1 正是其中負(fù)責(zé)核心判定的部分這也是其名字中 part 的由來。1.1 完整 isin 的數(shù)學(xué)表達(dá)README 給出了 isin 算子完整實(shí)現(xiàn)的公式化描述$$ elementsNum elements.length() \ combine_elements cat(elements, test_elements) \ value, index sort(combine_elements) \ z isin_part_v1(value, index, elementsNum) $$cat將elements與test_elements拼接為一個(gè)張量combine_elements其總長度記為totalLengthsort對拼接結(jié)果按值排序同時(shí)返回排序后的值value與每個(gè)值在拼接序列中的原始位置index其中下標(biāo)0..elementsNum-1對應(yīng)待判定元素elements下標(biāo)elementsNum..totalLength-1對應(yīng)test_elementsisin_part_v1基于排序后相等值必然相鄰這一性質(zhì)掃描value與index只要某個(gè)值既出現(xiàn)在elements段又出現(xiàn)在test_elements段就把輸出z中對應(yīng)位置置為true。由于sort后相等的值相鄰判定只需要在相鄰元素之間進(jìn)行IsinPartV1 因此可以做到接近線性的復(fù)雜度這是該Cat Sort Part組合設(shè)計(jì)優(yōu)于樸素兩兩比較的關(guān)鍵。1.2 從源碼確認(rèn)的輸入輸出關(guān)系在 experimental/math/isin_part_v1/examples/test_aclnn_isin_part_v1.cpp 的調(diào)用示例中注釋明確給出了語義等價(jià)關(guān)系語義等價(jià)于isin(elements[1,1,4,1,1], test_elements[0,1,2,3])示例直接構(gòu)造了三個(gè)輸入value {0, 1, 1, 1, 1, 1, 2, 3, 4}排序后的值elements[1,1,4,1,1]與test_elements[0,1,2,3]拼接后排序index {5, 0, 1, 3, 4, 6, 7, 8, 2}對應(yīng)原始位置0..4為待判定元素5..8為測試集元素elementsNum 5。期望輸出為[1, 1, 0, 1, 1]即 1 出現(xiàn)在測試集中所以為 14 不在測試集中所以為 0——這正好對應(yīng)了 isin 的成員判定語義。示例中直接構(gòu)造三個(gè)輸入、不依賴 Cat/Sort/Cast 等其它算子的做法也再次印證 IsinPartV1 是 isin 鏈路中可獨(dú)立驗(yàn)證的核心部件。二、產(chǎn)品支持情況與約束說明根據(jù) README 的產(chǎn)品支持矩陣產(chǎn)品是否支持Atlas A2 訓(xùn)練系列產(chǎn)品 / Atlas 800I A2 推理產(chǎn)品√約束說明無README 明確標(biāo)注該算子當(dāng)前無額外約束。從倉庫構(gòu)建配置看CMakeLists.txt 中通過add_all_modules_sources以O(shè)PTYPE isin_part_v1、ACLNNTYPE aclnn、COMPUTE_UNIT ascend950、TILING_DIR arch35方式登記該算子同時(shí) isin_part_v1_def.cpp 中為 AICore 增加了ascend910b配置?？梢酝茢嘣撍阕用嫦驎N騰 A2 系列910B 體系NPU 編譯具體以實(shí)際安裝的 CANN 版本與硬件形態(tài)為準(zhǔn)。需要說明的是該算子位于experimental/目錄屬于實(shí)驗(yàn)性算子使用前請關(guān)注對應(yīng)版本的能力說明與配套工具鏈。三、參數(shù)說明README 給出了完整的參數(shù)表此處結(jié)合算子定義源碼進(jìn)一步補(bǔ)充每個(gè)參數(shù)在實(shí)現(xiàn)層面的細(xì)節(jié)參數(shù)名輸入/輸出/屬性描述數(shù)據(jù)類型數(shù)據(jù)格式value輸入elements 與 test_elements 拼接排序后的張量值FLOAT、INT32NDindex輸入elements 與 test_elements 拼接排序后的張量索引INT32NDelementsNum輸入elements 元素?cái)?shù)量標(biāo)量INT64、INT32NDz輸出isin 操作后的輸出張量形狀由 elementsNum 參數(shù)決定BOOLND對參數(shù)表的補(bǔ)充說明依據(jù)源碼確認(rèn)value 與 index 的長度必須一致二者同為拼接后序列長度totalLength的排序結(jié)果。在 isin_part_v1_tiling.cpp 的GetShapeAttrsInfo中會校驗(yàn)輸入 x、輸入 y、輸出 z 三者維度數(shù)一致且逐維檢查inputX.dim outputZ.dim不滿足則直接返回失敗。elementsNum 是決定輸出形狀的關(guān)鍵標(biāo)量它雖然以張量形式傳入但語義上是一個(gè)標(biāo)量。InferShape階段isin_part_v1_infershape.cpp會直接從該輸入張量中讀出數(shù)值若為DT_INT32則取其首個(gè)元素將輸出形狀設(shè)置為[elementsNum]outShape-SetDim(0, elementsNum)即輸出z是長度為elementsNum的一維 BOOL 張量。關(guān)于 elementsNum 的數(shù)據(jù)類型README 參數(shù)表標(biāo)注支持INT64、INT32但當(dāng)前倉庫實(shí)現(xiàn)def.cpp中elementsNum僅注冊ge::DT_INT32infershape的GetNumValue也僅處理DT_INT32示例同樣以ACL_INT32傳入實(shí)際只走 INT32 路徑。使用 INT64 輸入時(shí)請以實(shí)際版本的實(shí)現(xiàn)與校驗(yàn)結(jié)果為準(zhǔn)。value 支持 FLOAT32 與 INT32 兩種 dtype這在 isin_part_v1_binary.json 中體現(xiàn)為兩個(gè)二進(jìn)制條目IsinPartV1_a1532827238e1555db7b997c7bce2928int32 輸入與IsinPartV1_1532827238e1555db7b997c7bce2928float32 輸入index與elementsNum均為 int32、輸出z為 bool全部使用 ND 格式format_match_mode為FormatAgnostic。所有輸入均為required必選參數(shù)且開啟AutoContiguousisin_part_v1_def.cpp框架側(cè)會保證張量內(nèi)存連續(xù)。四、算法原理如何用排序相鄰性判定成員關(guān)系IsinPartV1 的 kernel 實(shí)現(xiàn)位于 op_kernel/isin_part_v1.h其核心思想是排序之后相等的值必然連續(xù)相鄰。因此對于任意一個(gè)值 v只要存在某個(gè)下標(biāo)i使得value[i] v且index[i] elementsNum說明 v 來自 elements同時(shí)存在某個(gè)下標(biāo)j使得value[j] v且index[j] elementsNum說明 v 來自 test_elements那么 v 就在測試集中出現(xiàn)elements 中所有值為 v 的位置都應(yīng)輸出true。4.1 按 block 劃分與邊界修正kernel 按多核并行方式處理拼接后的序列按核劃分區(qū)間Initisin_part_v1.h從 GM 中讀取elementsNum根據(jù)totalLength與核數(shù)GetBlockNum()計(jì)算每個(gè)核負(fù)責(zé)的blockStart/blockEnd最后一個(gè)核處理到totalLength。邊界修正保證等值不跨核排序后相等的值可能恰好被切分到兩個(gè)核上。CopyInisin_part_v1.h會對區(qū)間做兩處調(diào)整若blockStart處元素與前一個(gè)元素相等則向后移動blockStart把相同值完整收進(jìn)本核區(qū)間若blockEnd處元素與后一個(gè)元素相等則向后擴(kuò)展blockEnd。通過這種值邊界對齊保證任意一個(gè)值的所有實(shí)例都落在同一個(gè)核內(nèi)避免跨核比較產(chǎn)生遺漏或重復(fù)。4.2 單核內(nèi)的掃描判定Computeisin_part_v1.h在區(qū)間[blockStart, blockEnd)上做單趟掃描維護(hù)指針pre指向當(dāng)前連續(xù)相等值段的起點(diǎn)當(dāng)發(fā)現(xiàn)index[j] elementsNum當(dāng)前值來自測試集且index[pre] elementsNum段首值來自 elements且value[j] value[pre]時(shí)說明當(dāng)前這個(gè)值同時(shí)出現(xiàn)在 elements 與 test_elements 中觸發(fā)GMout(pre, j-1)GMoutisin_part_v1.h遍歷段內(nèi)所有index elementsNum的下標(biāo)將輸出z的對應(yīng)位置寫入 1當(dāng)value[j] ! value[pre]時(shí)說明進(jìn)入了新的值段重置pre j、清除觸發(fā)標(biāo)記掃描完成后若發(fā)生過寫入調(diào)用zGmCache.Flush()將 UB 緩存回寫到 GM。由于zGmCache使用了SetAtomicAddDataCopyPadisin_part_v1.h即使不同核有重疊回寫也能保證最終結(jié)果正確輸出緩存未命中時(shí)通過MTE3_V事件同步保證先刷舊塊、再寫新塊的順序isin_part_v1.h。4.3 緩存友好的 GM 訪問排序后序列長度可能遠(yuǎn)超 UB 容量kernel 專門實(shí)現(xiàn)了兩個(gè)緩存類GmInCacheisin_part_v1.h為value與index提供按需加載的讀緩存以IN_ELEMENTS_SIZE64KB為粒度分塊DataCopyPad到 UB索引不在緩存范圍內(nèi)時(shí)自動換塊GmOutCacheisin_part_v1.h為輸出z提供按需刷新的寫緩存以O(shè)UT_ELEMENTS_SIZE32KB為粒度先Duplicate清零再逐位置 1換塊前先 Flush。五、Tiling 與配置host 側(cè)如何為 kernel 準(zhǔn)備參數(shù)Tiling 是算子在 host 側(cè)為 device kernel 計(jì)算切分參數(shù)的過程IsinPartV1 的 tiling 實(shí)現(xiàn)在 op_host/isin_part_v1_tiling.cpp主要流程如下獲取平臺信息GetPlatformInfo通過PlatformAscendC獲取 AIV 核數(shù)coreNum與 UB 大小ubSize任一為 0 則報(bào)錯退出校驗(yàn) shape 與 dtypeGetShapeAttrsInfo校驗(yàn)elementsNum必須是 1 維、輸入輸出維度一致、value僅支持DT_FLOAT/DT_INT32、index與elementsNum僅支持DT_INT32totalIdx取第一個(gè)輸入拼接后序列的總元素?cái)?shù)即totalLength空輸入處理totalIdx 0時(shí) tiling 數(shù)據(jù)清零、SetBlockDim(1)直接返回避免 kernel 空跑申請 workspaceGetWorkspaceSize固定申請WS_SYS_SIZE 16MB的系統(tǒng)工作空間寫入currentWorkspace[0]計(jì)算 tile 參數(shù)IsinPartV1TilingFunc以BLOCK_SIZE 32字節(jié)、BUFFER_NUM 2雙緩沖為基準(zhǔn)計(jì)算單 tile 元素?cái)?shù)tileDataNum、總 block 數(shù)blocksTotal再按核數(shù)均分得到 small-core / big-core 兩組數(shù)據(jù)量余數(shù)部分由前tailBlockNum個(gè)核多承擔(dān)一個(gè) block并把smallCoreDataNum、bigCoreDataNum、tileDataNum、smallTailDataNum、bigTailDataNum、finalSmallTileNum、finalBigTileNum、tailBlockNum、totalLength寫入 IsinPartV1TilingData 結(jié)構(gòu)體最后SetBlockDim(finalCoreNum)。Tiling 數(shù)據(jù)結(jié)構(gòu)的 9 個(gè)字段定義在 isin_part_v1_tiling_data.hkernel 側(cè)通過REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT讀取isin_part_v1.cpp。調(diào)度模式tiling key在 isin_part_v1_tiling_key.h 中定義了ELEMENTWISE_TPL_SCH_MODE_0與ELEMENTWISE_TPL_SCH_MODE_1兩種模板參數(shù)組合當(dāng)前默認(rèn)使用 MODE_0對應(yīng) isin_part_v1_simplified_key.ini 中的default0。算子注冊側(cè)的支撐信息算子原型isin_part_v1_def.cpp通過OpDef注冊IsinPartV1聲明 3 個(gè)輸入 1 個(gè)輸出AICore 配置開啟動態(tài)編譯靜態(tài)標(biāo)志、動態(tài) shape/rank 支持并設(shè)置PrecisionReduceFlag(true)shape 推導(dǎo)isin_part_v1_infershape.cpp輸出形狀直接取自elementsNum的數(shù)值注冊為IMPL_OP_INFERSHAPE(IsinPartV1)算子入口isin_part_v1.cpp__global__ __aicore__模板 kernelschMode為模板參數(shù)實(shí)例化NsIsinPartV1::IsinPartV1DTYPE_VALUE后依次執(zhí)行Init與Process。六、aclnn 調(diào)用實(shí)戰(zhàn)README 的調(diào)用說明指明通過 aclnn 接口方式調(diào)用 IsinPartV1 算子調(diào)用樣例為 examples/test_aclnn_isin_part_v1.cpp。下面結(jié)合該示例講解完整的 aclnn 單算子調(diào)用流程ACL 接口風(fēng)格與 CANN 通用單算子調(diào)用一致。6.1 調(diào)用流程總覽標(biāo)準(zhǔn) aclnn 調(diào)用分為六步示例代碼逐一步驟均有注釋環(huán)境初始化aclInit→aclrtSetDevice→aclrtCreateContext→aclrtCreateStream示例固定使用deviceId 0構(gòu)造輸入輸出 Tensor通過aclCreateTensor創(chuàng)建value、index、elementsNum三個(gè)輸入 Tensor 與輸出zTensor示例提供了CreateAclTensor與CreateEmptyTensor兩個(gè)輔助函數(shù)查詢 workspace 大小調(diào)用aclnnIsinPartV1GetWorkspaceSize(value, index, elementsNum, z, wsSize, executor)返回所需 workspace 字節(jié)數(shù)與算子執(zhí)行器句柄申請 workspace 并執(zhí)行aclrtMalloc申請 workspace 后調(diào)用aclnnIsinPartV1(wsAddr, wsSize, executor, stream)同步并取回結(jié)果aclrtSynchronizeStream后通過aclrtMemcpyACL_MEMCPY_DEVICE_TO_HOST拷貝輸出示例的PrintOutResult按%u打印 BOOL 結(jié)果資源清理依次aclDestroyTensor、aclrtFree含 workspace、aclrtDestroyStream、aclrtDestroyContext、aclFinalize。6.2 示例中的關(guān)鍵數(shù)據(jù)示例構(gòu)造的輸入對應(yīng)語義isin(elements[1,1,4,1,1], test_elements[0,1,2,3])參數(shù)值dtypevalue{0, 1, 1, 1, 1, 1, 2, 3, 4}ACL_FLOATindex{5, 0, 1, 3, 4, 6, 7, 8, 2}ACL_INT32elementsNum{5}ACL_INT32輸出 z形狀{5}ACL_BOOL期望輸出[1, 1, 0, 1, 1]——elements 中的 1 出現(xiàn)在 test_elements 中index 為 0、1、3、4 的四個(gè) 1 都命中而 4 未出現(xiàn)在測試集中因此對應(yīng)位置為 0。6.3 動手驗(yàn)證建議編譯示例時(shí)需鏈接 CANN 的 acl 與 aclnn 運(yùn)行庫并包含acl/acl.h與aclnnop/aclnn_isin_part_v1.h頭文件示例已在代碼中引入示例中對每個(gè) ACL 調(diào)用均用CHECK_RET宏檢查返回值失敗即打印錯誤碼退出便于定位問題若想驗(yàn)證完整的 isin 語義可在調(diào)用 IsinPartV1 之前先用 aclnn 的 Cat 與 Sort 接口或框架層圖編譯完成拼接與排序再將value、index、elementsNum喂給aclnnIsinPartV1運(yùn)行環(huán)境為 Atlas A2 訓(xùn)練系列 / Atlas 800I A2 推理產(chǎn)品構(gòu)建登記為COMPUTE_UNIT ascend950CMakeLists.txt實(shí)驗(yàn)性算子請使用對應(yīng) CANN 版本配套的算子包與工具鏈。七、小結(jié)IsinPartV1 展示了 CANN ops-math 中一種典型的功能拆分 基礎(chǔ)算子組合算子設(shè)計(jì)范式用cat與sort把 isin 問題規(guī)約到排序后相鄰等值判定再用一個(gè)輕量核心部件完成 O(N) 的成員關(guān)系輸出。其 host 側(cè) tiling 負(fù)責(zé) shape/dtype 校驗(yàn)、workspace 申請與多核數(shù)據(jù)切分device 側(cè) kernel 則通過值邊界對齊、UB 分塊緩存與原子回寫保證正確性與訪存效率。對于需要在昇騰 NPU 上實(shí)現(xiàn)集合成員判定類算子或研究 CANN 單算子開發(fā)流程的開發(fā)者從 README 出發(fā)配合 tiling 實(shí)現(xiàn)、kernel 實(shí)現(xiàn) 與 aclnn 調(diào)用示例 三份材料對照閱讀是一條完整且高效的學(xué)習(xí)路徑?！久赓M(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 IsinPartV1 算子深度解析：基于 Cat+Sort 的 isin 核心部件實(shí)現(xiàn)與 aclnn 調(diào)用實(shí)戰(zhàn)
CANN ops-math IsinPartV1 算子深度解析基于 CatSort 的 isin 核心部件實(shí)現(xiàn)與 aclnn 調(diào)用實(shí)戰(zhà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-mathIsinPartV1 是 CANN ops-math 倉庫 experimental/math/isin_part_v1 目錄下提供的一個(gè)實(shí)驗(yàn)性數(shù)學(xué)算子承擔(dān) isin 功能中判斷成員關(guān)系的核心部件與 Cat、Sort 算子配合即可拼裝出完整的 isin 語義。本文以該算子 README 為主體結(jié)合倉庫內(nèi) op_host、op_kernel、config 與調(diào)用示例源碼完整講解 IsinPartV1 的算子語義、參數(shù)規(guī)格、算法流程、底層實(shí)現(xiàn)與 aclnn 單算子調(diào)用方法幫助讀者在 Atlas A2 訓(xùn)練系列產(chǎn)品 / Atlas 800I A2 推理產(chǎn)品上獨(dú)立完成該算子的理解、構(gòu)建與驗(yàn)證。一、算子定位isin 的核心部件設(shè)計(jì)isin算子的功能是判斷給定張量elements的每個(gè)元素是否出現(xiàn)在另一個(gè)張量test_elements中并返回一個(gè)與elements形狀一致的布爾張量。在 CANN ops-math 中這一功能并未做成一個(gè)整體算子而是被拆解為Cat拼接→ Sort排序→ IsinPartV1成員判定三步IsinPartV1 正是其中負(fù)責(zé)核心判定的部分這也是其名字中 part 的由來。1.1 完整 isin 的數(shù)學(xué)表達(dá)README 給出了 isin 算子完整實(shí)現(xiàn)的公式化描述$$ elementsNum elements.length() \ combine_elements cat(elements, test_elements) \ value, index sort(combine_elements) \ z isin_part_v1(value, index, elementsNum) $$cat將elements與test_elements拼接為一個(gè)張量combine_elements其總長度記為totalLengthsort對拼接結(jié)果按值排序同時(shí)返回排序后的值value與每個(gè)值在拼接序列中的原始位置index其中下標(biāo)0..elementsNum-1對應(yīng)待判定元素elements下標(biāo)elementsNum..totalLength-1對應(yīng)test_elementsisin_part_v1基于排序后相等值必然相鄰這一性質(zhì)掃描value與index只要某個(gè)值既出現(xiàn)在elements段又出現(xiàn)在test_elements段就把輸出z中對應(yīng)位置置為true。由于sort后相等的值相鄰判定只需要在相鄰元素之間進(jìn)行IsinPartV1 因此可以做到接近線性的復(fù)雜度這是該Cat Sort Part組合設(shè)計(jì)優(yōu)于樸素兩兩比較的關(guān)鍵。1.2 從源碼確認(rèn)的輸入輸出關(guān)系在 experimental/math/isin_part_v1/examples/test_aclnn_isin_part_v1.cpp 的調(diào)用示例中注釋明確給出了語義等價(jià)關(guān)系語義等價(jià)于isin(elements[1,1,4,1,1], test_elements[0,1,2,3])示例直接構(gòu)造了三個(gè)輸入value {0, 1, 1, 1, 1, 1, 2, 3, 4}排序后的值elements[1,1,4,1,1]與test_elements[0,1,2,3]拼接后排序index {5, 0, 1, 3, 4, 6, 7, 8, 2}對應(yīng)原始位置0..4為待判定元素5..8為測試集元素elementsNum 5。期望輸出為[1, 1, 0, 1, 1]即 1 出現(xiàn)在測試集中所以為 14 不在測試集中所以為 0——這正好對應(yīng)了 isin 的成員判定語義。示例中直接構(gòu)造三個(gè)輸入、不依賴 Cat/Sort/Cast 等其它算子的做法也再次印證 IsinPartV1 是 isin 鏈路中可獨(dú)立驗(yàn)證的核心部件。二、產(chǎn)品支持情況與約束說明根據(jù) README 的產(chǎn)品支持矩陣產(chǎn)品是否支持Atlas A2 訓(xùn)練系列產(chǎn)品 / Atlas 800I A2 推理產(chǎn)品√約束說明無README 明確標(biāo)注該算子當(dāng)前無額外約束。從倉庫構(gòu)建配置看CMakeLists.txt 中通過add_all_modules_sources以O(shè)PTYPE isin_part_v1、ACLNNTYPE aclnn、COMPUTE_UNIT ascend950、TILING_DIR arch35方式登記該算子同時(shí) isin_part_v1_def.cpp 中為 AICore 增加了ascend910b配置?？梢酝茢嘣撍阕用嫦驎N騰 A2 系列910B 體系NPU 編譯具體以實(shí)際安裝的 CANN 版本與硬件形態(tài)為準(zhǔn)。需要說明的是該算子位于experimental/目錄屬于實(shí)驗(yàn)性算子使用前請關(guān)注對應(yīng)版本的能力說明與配套工具鏈。三、參數(shù)說明README 給出了完整的參數(shù)表此處結(jié)合算子定義源碼進(jìn)一步補(bǔ)充每個(gè)參數(shù)在實(shí)現(xiàn)層面的細(xì)節(jié)參數(shù)名輸入/輸出/屬性描述數(shù)據(jù)類型數(shù)據(jù)格式value輸入elements 與 test_elements 拼接排序后的張量值FLOAT、INT32NDindex輸入elements 與 test_elements 拼接排序后的張量索引INT32NDelementsNum輸入elements 元素?cái)?shù)量標(biāo)量INT64、INT32NDz輸出isin 操作后的輸出張量形狀由 elementsNum 參數(shù)決定BOOLND對參數(shù)表的補(bǔ)充說明依據(jù)源碼確認(rèn)value 與 index 的長度必須一致二者同為拼接后序列長度totalLength的排序結(jié)果。在 isin_part_v1_tiling.cpp 的GetShapeAttrsInfo中會校驗(yàn)輸入 x、輸入 y、輸出 z 三者維度數(shù)一致且逐維檢查inputX.dim outputZ.dim不滿足則直接返回失敗。elementsNum 是決定輸出形狀的關(guān)鍵標(biāo)量它雖然以張量形式傳入但語義上是一個(gè)標(biāo)量。InferShape階段isin_part_v1_infershape.cpp會直接從該輸入張量中讀出數(shù)值若為DT_INT32則取其首個(gè)元素將輸出形狀設(shè)置為[elementsNum]outShape-SetDim(0, elementsNum)即輸出z是長度為elementsNum的一維 BOOL 張量。關(guān)于 elementsNum 的數(shù)據(jù)類型README 參數(shù)表標(biāo)注支持INT64、INT32但當(dāng)前倉庫實(shí)現(xiàn)def.cpp中elementsNum僅注冊ge::DT_INT32infershape的GetNumValue也僅處理DT_INT32示例同樣以ACL_INT32傳入實(shí)際只走 INT32 路徑。使用 INT64 輸入時(shí)請以實(shí)際版本的實(shí)現(xiàn)與校驗(yàn)結(jié)果為準(zhǔn)。value 支持 FLOAT32 與 INT32 兩種 dtype這在 isin_part_v1_binary.json 中體現(xiàn)為兩個(gè)二進(jìn)制條目IsinPartV1_a1532827238e1555db7b997c7bce2928int32 輸入與IsinPartV1_1532827238e1555db7b997c7bce2928float32 輸入index與elementsNum均為 int32、輸出z為 bool全部使用 ND 格式format_match_mode為FormatAgnostic。所有輸入均為required必選參數(shù)且開啟AutoContiguousisin_part_v1_def.cpp框架側(cè)會保證張量內(nèi)存連續(xù)。四、算法原理如何用排序相鄰性判定成員關(guān)系IsinPartV1 的 kernel 實(shí)現(xiàn)位于 op_kernel/isin_part_v1.h其核心思想是排序之后相等的值必然連續(xù)相鄰。因此對于任意一個(gè)值 v只要存在某個(gè)下標(biāo)i使得value[i] v且index[i] elementsNum說明 v 來自 elements同時(shí)存在某個(gè)下標(biāo)j使得value[j] v且index[j] elementsNum說明 v 來自 test_elements那么 v 就在測試集中出現(xiàn)elements 中所有值為 v 的位置都應(yīng)輸出true。4.1 按 block 劃分與邊界修正kernel 按多核并行方式處理拼接后的序列按核劃分區(qū)間Initisin_part_v1.h從 GM 中讀取elementsNum根據(jù)totalLength與核數(shù)GetBlockNum()計(jì)算每個(gè)核負(fù)責(zé)的blockStart/blockEnd最后一個(gè)核處理到totalLength。邊界修正保證等值不跨核排序后相等的值可能恰好被切分到兩個(gè)核上。CopyInisin_part_v1.h會對區(qū)間做兩處調(diào)整若blockStart處元素與前一個(gè)元素相等則向后移動blockStart把相同值完整收進(jìn)本核區(qū)間若blockEnd處元素與后一個(gè)元素相等則向后擴(kuò)展blockEnd。通過這種值邊界對齊保證任意一個(gè)值的所有實(shí)例都落在同一個(gè)核內(nèi)避免跨核比較產(chǎn)生遺漏或重復(fù)。4.2 單核內(nèi)的掃描判定Computeisin_part_v1.h在區(qū)間[blockStart, blockEnd)上做單趟掃描維護(hù)指針pre指向當(dāng)前連續(xù)相等值段的起點(diǎn)當(dāng)發(fā)現(xiàn)index[j] elementsNum當(dāng)前值來自測試集且index[pre] elementsNum段首值來自 elements且value[j] value[pre]時(shí)說明當(dāng)前這個(gè)值同時(shí)出現(xiàn)在 elements 與 test_elements 中觸發(fā)GMout(pre, j-1)GMoutisin_part_v1.h遍歷段內(nèi)所有index elementsNum的下標(biāo)將輸出z的對應(yīng)位置寫入 1當(dāng)value[j] ! value[pre]時(shí)說明進(jìn)入了新的值段重置pre j、清除觸發(fā)標(biāo)記掃描完成后若發(fā)生過寫入調(diào)用zGmCache.Flush()將 UB 緩存回寫到 GM。由于zGmCache使用了SetAtomicAddDataCopyPadisin_part_v1.h即使不同核有重疊回寫也能保證最終結(jié)果正確輸出緩存未命中時(shí)通過MTE3_V事件同步保證先刷舊塊、再寫新塊的順序isin_part_v1.h。4.3 緩存友好的 GM 訪問排序后序列長度可能遠(yuǎn)超 UB 容量kernel 專門實(shí)現(xiàn)了兩個(gè)緩存類GmInCacheisin_part_v1.h為value與index提供按需加載的讀緩存以IN_ELEMENTS_SIZE64KB為粒度分塊DataCopyPad到 UB索引不在緩存范圍內(nèi)時(shí)自動換塊GmOutCacheisin_part_v1.h為輸出z提供按需刷新的寫緩存以O(shè)UT_ELEMENTS_SIZE32KB為粒度先Duplicate清零再逐位置 1換塊前先 Flush。五、Tiling 與配置host 側(cè)如何為 kernel 準(zhǔn)備參數(shù)Tiling 是算子在 host 側(cè)為 device kernel 計(jì)算切分參數(shù)的過程IsinPartV1 的 tiling 實(shí)現(xiàn)在 op_host/isin_part_v1_tiling.cpp主要流程如下獲取平臺信息GetPlatformInfo通過PlatformAscendC獲取 AIV 核數(shù)coreNum與 UB 大小ubSize任一為 0 則報(bào)錯退出校驗(yàn) shape 與 dtypeGetShapeAttrsInfo校驗(yàn)elementsNum必須是 1 維、輸入輸出維度一致、value僅支持DT_FLOAT/DT_INT32、index與elementsNum僅支持DT_INT32totalIdx取第一個(gè)輸入拼接后序列的總元素?cái)?shù)即totalLength空輸入處理totalIdx 0時(shí) tiling 數(shù)據(jù)清零、SetBlockDim(1)直接返回避免 kernel 空跑申請 workspaceGetWorkspaceSize固定申請WS_SYS_SIZE 16MB的系統(tǒng)工作空間寫入currentWorkspace[0]計(jì)算 tile 參數(shù)IsinPartV1TilingFunc以BLOCK_SIZE 32字節(jié)、BUFFER_NUM 2雙緩沖為基準(zhǔn)計(jì)算單 tile 元素?cái)?shù)tileDataNum、總 block 數(shù)blocksTotal再按核數(shù)均分得到 small-core / big-core 兩組數(shù)據(jù)量余數(shù)部分由前tailBlockNum個(gè)核多承擔(dān)一個(gè) block并把smallCoreDataNum、bigCoreDataNum、tileDataNum、smallTailDataNum、bigTailDataNum、finalSmallTileNum、finalBigTileNum、tailBlockNum、totalLength寫入 IsinPartV1TilingData 結(jié)構(gòu)體最后SetBlockDim(finalCoreNum)。Tiling 數(shù)據(jù)結(jié)構(gòu)的 9 個(gè)字段定義在 isin_part_v1_tiling_data.hkernel 側(cè)通過REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT讀取isin_part_v1.cpp。調(diào)度模式tiling key在 isin_part_v1_tiling_key.h 中定義了ELEMENTWISE_TPL_SCH_MODE_0與ELEMENTWISE_TPL_SCH_MODE_1兩種模板參數(shù)組合當(dāng)前默認(rèn)使用 MODE_0對應(yīng) isin_part_v1_simplified_key.ini 中的default0。算子注冊側(cè)的支撐信息算子原型isin_part_v1_def.cpp通過OpDef注冊IsinPartV1聲明 3 個(gè)輸入 1 個(gè)輸出AICore 配置開啟動態(tài)編譯靜態(tài)標(biāo)志、動態(tài) shape/rank 支持并設(shè)置PrecisionReduceFlag(true)shape 推導(dǎo)isin_part_v1_infershape.cpp輸出形狀直接取自elementsNum的數(shù)值注冊為IMPL_OP_INFERSHAPE(IsinPartV1)算子入口isin_part_v1.cpp__global__ __aicore__模板 kernelschMode為模板參數(shù)實(shí)例化NsIsinPartV1::IsinPartV1DTYPE_VALUE后依次執(zhí)行Init與Process。六、aclnn 調(diào)用實(shí)戰(zhàn)README 的調(diào)用說明指明通過 aclnn 接口方式調(diào)用 IsinPartV1 算子調(diào)用樣例為 examples/test_aclnn_isin_part_v1.cpp。下面結(jié)合該示例講解完整的 aclnn 單算子調(diào)用流程ACL 接口風(fēng)格與 CANN 通用單算子調(diào)用一致。6.1 調(diào)用流程總覽標(biāo)準(zhǔn) aclnn 調(diào)用分為六步示例代碼逐一步驟均有注釋環(huán)境初始化aclInit→aclrtSetDevice→aclrtCreateContext→aclrtCreateStream示例固定使用deviceId 0構(gòu)造輸入輸出 Tensor通過aclCreateTensor創(chuàng)建value、index、elementsNum三個(gè)輸入 Tensor 與輸出zTensor示例提供了CreateAclTensor與CreateEmptyTensor兩個(gè)輔助函數(shù)查詢 workspace 大小調(diào)用aclnnIsinPartV1GetWorkspaceSize(value, index, elementsNum, z, wsSize, executor)返回所需 workspace 字節(jié)數(shù)與算子執(zhí)行器句柄申請 workspace 并執(zhí)行aclrtMalloc申請 workspace 后調(diào)用aclnnIsinPartV1(wsAddr, wsSize, executor, stream)同步并取回結(jié)果aclrtSynchronizeStream后通過aclrtMemcpyACL_MEMCPY_DEVICE_TO_HOST拷貝輸出示例的PrintOutResult按%u打印 BOOL 結(jié)果資源清理依次aclDestroyTensor、aclrtFree含 workspace、aclrtDestroyStream、aclrtDestroyContext、aclFinalize。6.2 示例中的關(guān)鍵數(shù)據(jù)示例構(gòu)造的輸入對應(yīng)語義isin(elements[1,1,4,1,1], test_elements[0,1,2,3])參數(shù)值dtypevalue{0, 1, 1, 1, 1, 1, 2, 3, 4}ACL_FLOATindex{5, 0, 1, 3, 4, 6, 7, 8, 2}ACL_INT32elementsNum{5}ACL_INT32輸出 z形狀{5}ACL_BOOL期望輸出[1, 1, 0, 1, 1]——elements 中的 1 出現(xiàn)在 test_elements 中index 為 0、1、3、4 的四個(gè) 1 都命中而 4 未出現(xiàn)在測試集中因此對應(yīng)位置為 0。6.3 動手驗(yàn)證建議編譯示例時(shí)需鏈接 CANN 的 acl 與 aclnn 運(yùn)行庫并包含acl/acl.h與aclnnop/aclnn_isin_part_v1.h頭文件示例已在代碼中引入示例中對每個(gè) ACL 調(diào)用均用CHECK_RET宏檢查返回值失敗即打印錯誤碼退出便于定位問題若想驗(yàn)證完整的 isin 語義可在調(diào)用 IsinPartV1 之前先用 aclnn 的 Cat 與 Sort 接口或框架層圖編譯完成拼接與排序再將value、index、elementsNum喂給aclnnIsinPartV1運(yùn)行環(huán)境為 Atlas A2 訓(xùn)練系列 / Atlas 800I A2 推理產(chǎn)品構(gòu)建登記為COMPUTE_UNIT ascend950CMakeLists.txt實(shí)驗(yàn)性算子請使用對應(yīng) CANN 版本配套的算子包與工具鏈。七、小結(jié)IsinPartV1 展示了 CANN ops-math 中一種典型的功能拆分 基礎(chǔ)算子組合算子設(shè)計(jì)范式用cat與sort把 isin 問題規(guī)約到排序后相鄰等值判定再用一個(gè)輕量核心部件完成 O(N) 的成員關(guān)系輸出。其 host 側(cè) tiling 負(fù)責(zé) shape/dtype 校驗(yàn)、workspace 申請與多核數(shù)據(jù)切分device 側(cè) kernel 則通過值邊界對齊、UB 分塊緩存與原子回寫保證正確性與訪存效率。對于需要在昇騰 NPU 上實(shí)現(xiàn)集合成員判定類算子或研究 CANN 單算子開發(fā)流程的開發(fā)者從 README 出發(fā)配合 tiling 實(shí)現(xiàn)、kernel 實(shí)現(xiàn) 與 aclnn 調(diào)用示例 三份材料對照閱讀是一條完整且高效的學(xué)習(xí)路徑?！久赓M(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），僅供參考