CANN SHMEM 的 ACLGraph 集成實踐基于圖捕獲與重放的多卡 allGather 精度驗證【免費下載鏈接】shmemCANN SHMEM 是面向昇騰平臺的多機(jī)多卡內(nèi)存通信庫基于OpenSHMEM 標(biāo)準(zhǔn)協(xié)議實現(xiàn)跨設(shè)備的高效內(nèi)存訪問與數(shù)據(jù)同步。項目地址: https://gitcode.com/cann/shmem導(dǎo)讀本文面向在昇騰平臺上使用 CANN SHMEM 的開發(fā)者圍繞倉庫中的 aclgraph_demo 樣例完整講解如何將allGather集合通信算子編排進(jìn) ACLGraphAscendCL 圖模式模型并通過「首次循環(huán)圖捕獲、后續(xù)循環(huán)圖重放」的方式驗證算子在圖內(nèi)逐次循環(huán)中的功能與精度。讀完本文你將掌握 ACLGraph 圖捕獲/重放的 API 使用流程、SHMEM 集合通信算子入?yún)⒌膱D適配改造方法magic參數(shù)從int值改為設(shè)備側(cè)地址__gm__ int *、樣例的編譯運行方式與參數(shù)含義以及設(shè)備側(cè) allGather 內(nèi)核基于信號量同步的實現(xiàn)原理。一、樣例背景為什么要在 ACLGraph 中集成 SHMEM 集合通信ACLGraph對應(yīng)aclmdlRI*系列接口是 AscendCL 提供的圖執(zhí)行模式先將一段由算子構(gòu)成的 Host 側(cè)調(diào)用序列「捕獲」為圖模型model后續(xù)循環(huán)只需重放該模型即可避免反復(fù)下發(fā)算子降低調(diào)度開銷。而 CANN SHMEM 作為面向昇騰多機(jī)多卡的內(nèi)存通信庫其集合通信算子如allGather通常以設(shè)備側(cè) kernel 形式在 Stream 上異步執(zhí)行天然可以與 ACLGraph 的算子編排模型結(jié)合。aclgraph_demo 的定位就是驗證在 ACLGraph 圖中混合編排普通的向量add算子與 SHMEM 的allGather集合通信算子并確認(rèn)這種混合圖在多次循環(huán)重放中功能與精度均正確。倉庫中同樣基于aclmdlRI接口的還有 rdma_aclgraph_demo可見「ACLGraph SHMEM 通信」是倉庫著力驗證的一類重要使用場景。二、ACLGraph 圖結(jié)構(gòu)與算子編排樣例構(gòu)造的 ACLGraph 模型包含四個算子節(jié)點按順序執(zhí)行為add → allGather → allGather → add具體數(shù)據(jù)流如下與 main.cpp 中捕獲區(qū)間的調(diào)用順序一一對應(yīng)第一個add將輸入input_ptr與input_b_ptr逐元素相加結(jié)果寫入input_c_ptr通過aclrtMemcpyAsync將input_c_ptr的結(jié)果設(shè)備到設(shè)備拷回input_ptr作為第一個allGather的輸入第一個allGather對input_ptr做全收集結(jié)果寫入output_ptr第二個allGather對input_b_ptr做全收集結(jié)果寫入output_b_ptr第二個add將兩個allGather的輸出相加結(jié)果寫入output_c_ptr。在 Host 側(cè)這四步被包在aclmdlRICaptureBegin/aclmdlRICaptureEnd之間完成圖捕獲main.cppif (zz 0) { aclmdlRICaptureBegin(stream, ACL_MODEL_RI_CAPTURE_MODE_RELAXED); run_vector_addT(trans_size, input_ptr, input_b_ptr, input_c_ptr, stream); aclrtMemcpyAsync(input_ptr, trans_size * sizeof(T), input_c_ptr, trans_size * sizeof(T), ACL_MEMCPY_DEVICE_TO_DEVICE, stream); allgather_demoT(BLOCK_NUM, stream, ffts_addr, (uint8_t*)input_ptr, (uint8_t*)output_ptr, (uint8_t*)ptr, trans_size, magic_ptr); allgather_demoT(BLOCK_NUM, stream, ffts_addr, (uint8_t*)input_b_ptr, (uint8_t*)output_b_ptr, (uint8_t*)ptr, trans_size, magic_ptr2); run_vector_addT(n_pes * trans_size, output_ptr, output_b_ptr, output_c_ptr, stream); aclmdlRICaptureEnd(stream, model); aclmdlRIExecuteAsync(model, stream); } else { aclmdlRIExecuteAsync(model, stream); }捕獲完成后model保存了整條算子鏈后續(xù)每次循環(huán)只調(diào)用aclmdlRIExecuteAsync(model, stream)完成重放。樣例默認(rèn)loop_times 3main.cpp即 1 次捕獲 2 次重放每次循環(huán)結(jié)束后都會執(zhí)行一次精度校驗。三、核心適配改造magic入?yún)膇nt值改為設(shè)備地址為了適配 ACLGraphallGather算子用于同步控制的magic入?yún)⑿枰獜摹窰ost 側(cè)傳入的int值」改為「設(shè)備側(cè)地址__gm__ int *」。這樣做的原因非常直接被捕獲進(jìn)圖內(nèi)的 kernel 參數(shù)在重放時是固定的只有通過指針引用設(shè)備內(nèi)存才能在每次重放前由 Host 修改該內(nèi)存中的值從而讓圖內(nèi)通信邏輯感知到循環(huán)次數(shù)的變化。從源碼可以清楚看到這一改造的落點內(nèi)核側(cè)簽名中magic已是設(shè)備地址__gm__ int* magic并在內(nèi)核入口通過int magic_data *magic;讀取aclgraph_demo_kernel.cppHost 側(cè)在每次循環(huán)開始前先計算新的 magic 值并拷貝到設(shè)備內(nèi)存再觸發(fā)圖執(zhí)行main.cppfor (int zz 0; zz loop_times; zz) { magic; magic_value magic * MAGIC_MULTIPLIER; // MAGIC_MULTIPLIER 1024 magic_value_2 (magic loop_times 10) * MAGIC_MULTIPLIER; *magic_host_ptr magic_value; aclrtMemcpy(magic_ptr, sizeof(int), magic_host_ptr, sizeof(int), ACL_MEMCPY_HOST_TO_DEVICE); *magic_host_ptr magic_value_2; aclrtMemcpy(magic_ptr2, sizeof(int), magic_host_ptr, sizeof(int), ACL_MEMCPY_HOST_TO_DEVICE); ... }magic在圖中承擔(dān)同步角色設(shè)備側(cè)內(nèi)核用magic_data構(gòu)造信號signal值通過aclshmemx_signal_op/aclshmem_signal_wait_until完成跨 PE 的同步握手詳見下文第五節(jié)。由于每次循環(huán)magic單調(diào)遞增可以有效區(qū)分不同輪次的同步信號避免重放時舊信號被誤當(dāng)作新一輪完成信號。這種「參數(shù)指針化 循環(huán)前更新」的模式正是 SHMEM 算子被安全嵌入 ACLGraph 圖模型的關(guān)鍵技巧。四、編譯與運行4.1 編譯構(gòu)建在倉庫根目錄shmem/下執(zhí)行構(gòu)建腳本并攜帶-examples選項即可編出所有示例含 aclgraph_demoA2 / A3 平臺bash scripts/build.sh -examplesAscend950 平臺bash scripts/build.sh -soc_type Ascend950 -examples構(gòu)建完成后可執(zhí)行文件輸出到build/bin/aclgraph_demo動態(tài)庫輸出到build/librun.sh 中的LD_LIBRARY_PATH即指向這兩個目錄。4.2 運行 demo在examples/aclgraph_demo目錄下執(zhí)行# 完成 PEs 卡下的 ACLGraphadd allGather allGather add并在每次循環(huán)中驗證精度。 # 支持 PEs : [2, 4, 8] bash run.sh -pes ${PEs}運行腳本會依次完成三件事run.sh生成 golden 數(shù)據(jù)調(diào)用python3 ./scripts/data_gen.py PE_SIZE TEST_TYPE生成每個 PE 的輸入與期望輸出設(shè)置運行環(huán)境導(dǎo)出SHMEM_UID_SESSION_ID127.0.0.1:8899并設(shè)置LD_LIBRARY_PATH指向構(gòu)建產(chǎn)物拉起多進(jìn)程按GNPU_NUM個進(jìn)程數(shù)用msprof以--application方式啟動build/bin/aclgraph_demo每個進(jìn)程對應(yīng)一個 PE。4.3 run.sh 可選參數(shù)run.sh除-pes外還支持以下參數(shù)默認(rèn)值見 run.sh參數(shù)含義默認(rèn)值說明-pesPE進(jìn)程/rank數(shù)量2必須為正整數(shù)樣例驗證 PEs 2 / 4 / 8-ipport通信服務(wù)監(jiān)聽地址tcp://127.0.0.1:8766非空字符串-gnpus本機(jī) NPU 數(shù)量進(jìn)程數(shù)8必須為正整數(shù)若大于PE_SIZE會自動收斂為PE_SIZE-fnpu起始 NPU 編號0非負(fù)整數(shù)-type測試數(shù)據(jù)類型int當(dāng)前僅支持int-tool性能采集工具msprof僅支持msprof腳本內(nèi)置了嚴(yán)格的參數(shù)校驗非法參數(shù)、空值、非整數(shù)、超出支持范圍的值都會報錯退出當(dāng)GNPU_NUM PE_SIZE時自動把GNPU_NUM賦值為PE_SIZE并打印提示。4.4 直接運行二進(jìn)制也可跳過腳本直接運行二進(jìn)制參數(shù)順序固定見 main.cpp./build/bin/aclgraph_demo n_pes pe_id ipport g_npus f_pe f_npu data_type loop_times位置參數(shù)含義樣例值n_pesPE 總數(shù)2 / 4 / 8pe_id當(dāng)前進(jìn)程的 PE 編號0 ~ n_pes-1ipport通信地址tcp://127.0.0.1:8998main 默認(rèn)值g_npus本機(jī) NPU 數(shù)8f_pe/f_npu起始 PE / NPU 偏移0 / 0data_type數(shù)據(jù)類型僅支持intloop_times循環(huán)次數(shù)含捕獲3進(jìn)程內(nèi)會按device_id pe_id % g_npus f_npu綁定 NPUmain.cpp隨后依次執(zhí)行aclInit、aclrtSetDevice、aclshmemx_init_attr初始化模式為ACLSHMEMX_INIT_WITH_DEFAULT并在結(jié)束時按aclshmem_finalize→aclrtResetDevice→aclFinalize的順序清理。五、主機(jī)側(cè)實現(xiàn)細(xì)節(jié)初始化、捕獲與逐輪校驗5.1 初始化與公共參數(shù)main.cpp開頭定義了若干編譯期常量main.cpp理解它們有助于把握樣例的資源邊界constexpr int64_t SYNC_FLAG_INTERVAL 16; // 同步 flag 間隔 constexpr int64_t UB_DMA_MAX_SIZE 190 * 1024; // UB 單次 DMA 搬運上限約 190 KB constexpr int64_t GVA_BUFF_MAX_SIZE 100 * 1024 * 1024; // 對稱內(nèi)存緩沖上限100 MB constexpr uint32_t MAGIC_MULTIPLIER 1024; // magic 信號值放大倍數(shù) constexpr uint32_t DATA_SIZE_THRESHOLD 2097152; // 大數(shù)據(jù)/小數(shù)據(jù)內(nèi)核分界2 MB constexpr uint32_t BLOCK_NUM_SMALL_DATA 8; // 小數(shù)據(jù)內(nèi)核 AIV 核數(shù) constexpr uint32_t BLOCK_NUM_LARGE_DATA 16; // 大數(shù)據(jù)內(nèi)核 AIV 核數(shù)測試共分 2 個 case數(shù)據(jù)長度為262144 * (1 i)i 0, 1即 256K 與 512K 個int元素main.cpp。根據(jù)trans_size * sizeof(T)是否小于 2 MB 閾值選擇 8 核或 16 核 AIV 的內(nèi)核啟動配置main.cpp。5.2 設(shè)備內(nèi)存布局每個 case 內(nèi)樣例使用aclrtMalloc分配輸入/輸出緩沖并使用 SHMEM 對稱堆接口分配同步與數(shù)據(jù)緩沖main.cppvoid *ptr aclshmem_malloc(aiv_num * SYNC_FLAG_INTERVAL * sizeof(T) GVA_BUFF_MAX_SIZE / sizeof(T));這段aclshmem_malloc分配的內(nèi)存同時承載兩塊用途起始aiv_num * SYNC_FLAG_INTERVAL個int32用作同步 flag 區(qū)每核一個 flag 槽位后續(xù)最多GVA_BUFF_MAX_SIZE字節(jié)用作GVA 數(shù)據(jù)區(qū)即 allGather 各 PE 寫入、讀取的對稱內(nèi)存中轉(zhuǎn)區(qū)。5.3 逐輪精度校驗每次循環(huán)在aclrtSynchronizeStream(stream)之后將output_c_ptr拷回 Host與 golden 對比main.cppfor (int zzz 0; zzz n_pes * trans_size; zzz) { if (static_castint(output_host[zzz]) ! static_castint(golden_host[zzz]) * (zz 3)) { std::cout ERROR each loop: ... std::endl; status -1; break; } }期望值是golden * (zz 3)其中zz為循環(huán)序號0、1、2golden為所有 PE 輸入拼接的全收集正確結(jié)果。該期望值與數(shù)據(jù)流嚴(yán)格自洽第一次add把input更新為2 × 原輸入隨后兩個 allGather 分別貢獻(xiàn)2 × golden與1 × golden第二次add得3 × golden此后輸入隨循環(huán)不斷累加期望值隨之遞增為4 × golden、5 × golden?？梢娦ｒ炌瑫r覆蓋了兩點allGather 圖內(nèi)功能正確以及每次重放時經(jīng)設(shè)備地址更新的magic參數(shù)確實生效。每個 case 結(jié)束時還會調(diào)用aclshmemx_get_prof(nullptr, true)觸發(fā)性能數(shù)據(jù)落盤并執(zhí)行內(nèi)存與模型銷毀清理。六、設(shè)備側(cè) allGather 內(nèi)核實現(xiàn)信號量同步與大小數(shù)據(jù)分路內(nèi)核入口ShmemAllGather_int依據(jù)數(shù)據(jù)量分派兩條實現(xiàn)路徑aclgraph_demo_kernel.cpp小于 2 MB 走all_gather_small_data否則走all_gather_big_data。兩條路徑共享同一套「兩步走」流水Step 1本地寫遠(yuǎn)端每個 AIV 核把本 PE 輸入input的一段數(shù)據(jù)通過aclshmemx_mte_put_nbi寫入自己的對稱內(nèi)存區(qū)GVA 數(shù)據(jù)區(qū)隨后通過aclshmemx_signal_op(..., ACLSHMEM_SIGNAL_SET, my_rank)發(fā)出遞增信號Step 2遠(yuǎn)端讀本地各核通過aclshmem_signal_wait_until(..., ACLSHMEM_CMP_EQ, magic_data)等待對端信號到達(dá)再通過aclshmemx_mte_get_nbi從對端對稱內(nèi)存區(qū)把數(shù)據(jù)拉回output的對應(yīng)分片位置。小數(shù)據(jù)路徑直接以elements / aiv_num切分?jǐn)?shù)據(jù)配合aclshmem_quiet()與AscendC::SyncAll()保證本端寫任務(wù)完成后再發(fā)信號aclgraph_demo_kernel.cpp大數(shù)據(jù)路徑則將數(shù)據(jù)按GVA_BUFF_MAX_SIZE分塊循環(huán)處理并把 AIV 核分成「發(fā)送半?yún)^(qū)」與「接收半?yún)^(qū)」兩組接收組使用 ping-pong 雙緩沖ping_buff/pong_buff輪轉(zhuǎn)拉取數(shù)據(jù)同時用flags_ub1/flags_ub2兩級 UB 標(biāo)志位跟蹤各數(shù)據(jù)塊的就緒進(jìn)度避免等待空轉(zhuǎn)aclgraph_demo_kernel.cpp。可以推斷這套內(nèi)核的同步信號語義正是上文「magic 入?yún)⒅羔樆垢脑斓闹苯邮芤嬲適agic_data參與信號值的構(gòu)造與比較Host 側(cè)每輪循環(huán)更新設(shè)備內(nèi)存中的magic即可讓重放圖內(nèi)每一輪的信號匹配新值從而實現(xiàn)圖內(nèi)通信算子的多輪安全復(fù)用。內(nèi)核啟動封裝allgather_demoT位于 aclgraph_demo_kernel.h以block_dimAIV 核數(shù)、stream、FFTS 配置地址、輸入輸出及magic設(shè)備地址為參數(shù)。圖內(nèi)的兩個向量add算子由倉庫內(nèi)的 AscendC 算子實現(xiàn)add_custom.hadd_custom內(nèi)核按塊切分?jǐn)?shù)據(jù)使用深度為 2 的TPipe流水隊列執(zhí)行 CopyIn →AscendC::Add→ CopyOut并單獨處理尾塊Host 封裝run_vector_add依據(jù)platform_ascendc查詢到的 UB 大小與 AIV 核數(shù)自動計算numBlocks、blockLength與tileSize從而讓一個簡單的 add 算子也能在全卡規(guī)模上并行。七、golden 數(shù)據(jù)生成與校驗基準(zhǔn)run.sh啟動前會自動調(diào)用 scripts/data_gen.py 生成基準(zhǔn)數(shù)據(jù)避免手工準(zhǔn)備 bin 文件固定隨機(jī)種子np.random.seed(42)保證多節(jié)點、多次運行的數(shù)據(jù)可復(fù)現(xiàn)每個 PE 的輸入input_gm_i.bin用np.random.uniform(0.0, 10.0)生成全收集的正確結(jié)果golden.bin為所有 PE 輸入的按序拼接目錄按golden/allgather_data_len_pe_size/組織主程序運行時按同樣的路徑拼接規(guī)則加載input_gm_pe_id.bin與golden.binmain.cpptype_map預(yù)留了int32/float16/bfloat16等類型的映射但當(dāng)前樣例的run.sh與主程序均限定int類型。八、構(gòu)建與運行的適用前提構(gòu)建命令中的-soc_type Ascend950僅用于 Ascend950 平臺A2 / A3 平臺直接使用默認(rèn)目標(biāo)即可運行需要ASCEND_HOME_PATH環(huán)境變量指向已安裝的 CANN 工具鏈run.sh用它拼接lib64路徑多進(jìn)程由msprof --application拉起每個進(jìn)程綁定一塊 NPU若PE_SIZE超過單機(jī) NPU 數(shù)需調(diào)整-gnpus與-fnpu以適配實際機(jī)器拓?fù)錁永敵瞿夸洖閑xamples/aclgraph_demo/output/msprof 落盤位置可在其中檢查采集到的性能與日志數(shù)據(jù)。總結(jié)aclgraph_demo 為「ACLGraph 圖模式 SHMEM 集合通信算子」的融合使用提供了完整的參考實現(xiàn)通過將allGather的magic同步參數(shù)改造為設(shè)備側(cè)地址__gm__ int *成功讓帶通信算子的圖模型在捕獲后支持多輪重放且每輪的信號同步與精度均正確。其主機(jī)側(cè)覆蓋了 ACL 初始化、圖捕獲/重放、golden 校驗與資源回收的完整生命周期設(shè)備側(cè)則給出了基于signal_op/signal_wait_until與 UB ping-pong 緩沖的大小數(shù)據(jù)雙路 allGather 內(nèi)核實現(xiàn)。開發(fā)者可直接復(fù)用本文的編譯、運行與參數(shù)說明將該模式遷移到其他集合通信算子如 reduce、allreduce 等與 ACLGraph 的集成場景中?！久赓M下載鏈接】shmemCANN SHMEM 是面向昇騰平臺的多機(jī)多卡內(nèi)存通信庫基于OpenSHMEM 標(biāo)準(zhǔn)協(xié)議實現(xiàn)跨設(shè)備的高效內(nèi)存訪問與數(shù)據(jù)同步。項目地址: https://gitcode.com/cann/shmem創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考