原位縮放累加(a·x+y)的實(shí)現(xiàn)原理與實(shí)戰(zhàn))
PTO-ISA TAXPY 指令深度解析Tile 級(jí)原位縮放累加a·xy的實(shí)現(xiàn)原理與實(shí)戰(zhàn)【免費(fèi)下載鏈接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.項(xiàng)目地址: https://gitcode.com/cann/pto-isa導(dǎo)讀本文以 CANN pto-isa 倉(cāng)庫(kù)中的 docs/isa/TAXPY.md 為核心完整講解 PTOParallel Tile Operation虛擬指令集中的TAXPY指令——它在 Tile 上執(zhí)行原位縮放累加AXPY$a \cdot x y$。你將掌握 TAXPY 的數(shù)學(xué)語(yǔ)義、C 內(nèi)建接口簽名與參數(shù)約定、支持的數(shù)據(jù)類型組合、在向量流水線PIPE_V上的底層實(shí)現(xiàn)差異A5 與 A2/A3 兩個(gè)后端路徑、CPU_SIM 參考實(shí)現(xiàn)以及如何編寫可在 NPU 上運(yùn)行的完整 ST 測(cè)試內(nèi)核。讀完本文你可以在自己的 kernel 中安全、正確地使用TAXPY完成帶標(biāo)量縮放的累加運(yùn)算。TAXPY 指令概述TAXPY對(duì)單個(gè) Tile 執(zhí)行原位縮放累加將源 Tilesrc0$x$按標(biāo)量scalar$a$縮放后累加到目標(biāo) Tiledst$y$上結(jié)果寫回dst$$ \mathrm{dst}{i,j} \leftarrow \mathrm{scalar} \cdot \mathrm{src0}{i,j} \mathrm{dst}_{i,j} $$三個(gè)操作數(shù)的角色非常明確見 docs/isa/TAXPY.mddst既是累加輸入$y$也是輸出調(diào)用前必須已初始化src0只讀$x$逐元素參與運(yùn)算scalar標(biāo)量縮放系數(shù)$a$類型為TileDataSrc::DType。這與普通的TADDSTile 加標(biāo)量不同TAXPY是兩個(gè) Tile 之間帶標(biāo)量系數(shù)的融合運(yùn)算輸入輸出共用一個(gè) TileRMW讀-修改-寫天然適合在已有累加基上做帶權(quán)累加。數(shù)學(xué)語(yǔ)義基于有效區(qū)域的逐元素 RMW指令語(yǔ)義定義在 Tile 的有效區(qū)域valid region內(nèi)。對(duì)有效區(qū)域中的每個(gè)元素(i, j)$$ \mathrm{dst}{i,j}^{\text{new}} \mathrm{scalar} \cdot \mathrm{src0}{i,j} \mathrm{dst}_{i,j}^{\text{old}} $$各操作數(shù)的數(shù)據(jù)流角色操作數(shù)角色說明dst讀-修改-寫RMW讀入舊值作為累加基 $y$寫回 $\mathrm{scalar} \cdot x y$src0只讀逐元素貢獻(xiàn)被縮放的 $x$scalar標(biāo)量縮放系數(shù) $a$類型為TileDataSrc::DType除非另有說明語(yǔ)義在有效區(qū)域內(nèi)定義目標(biāo)相關(guān)行為標(biāo)記為實(shí)現(xiàn)定義implementation-defined。這一語(yǔ)義決定了兩個(gè)關(guān)鍵約束dst必須先初始化否則累加基是未定義數(shù)據(jù)dst與src0的有效形狀必須一一對(duì)應(yīng)逐元素映射。C 內(nèi)建接口簽名、聲明位置與參數(shù)約定TAXPY的公共聲明位于 include/pto/common/pto_instr.hpp對(duì)外包含頭為pto/pto-inst.hpp內(nèi)部實(shí)現(xiàn)聲明在pto/common/pto_instr.hpptemplate typename TileDataDst, typename TileDataSrc, typename... WaitEvents PTO_INST RecordEvent TAXPY(TileDataDst dst, TileDataSrc src0, typename TileDataSrc::DType scalar, WaitEvents ...events);參數(shù)方向含義dst輸入/輸出累加基與結(jié)果 Tile$y$讀-修改-寫Vecsrc0輸入縮放源 Tile$x$只讀Vec有效形狀與dst相同scalar輸入標(biāo)量縮放系數(shù)$a$類型為TileDataSrc::DTypeevents...輸入等待事件WaitEvents指令前隱式 event synchronization從源碼看公開包裝函數(shù)先調(diào)用detail::PtoWaitEvents(events...)完成隱式事件同步再經(jīng)MAP_INSTR_IMPL(TAXPY, dst, src0, scalar)宏分發(fā)到各后端的TAXPY_IMPL實(shí)現(xiàn)參見 include/pto/common/pto_instr.hpp 的宏定義。返回值類型為RecordEvent可用于指令間依賴編排。Tile 尺寸與數(shù)據(jù)類型對(duì)于有效形狀 $M \times N$ 的 TileTiledtype有效形狀TileType說明dsthalf或float$M \times N$Vec(UB)累加基 結(jié)果RMWsrc0half或float$M \times N$Vec(UB)縮放源逐元素dst與src0的有效行數(shù)、有效列數(shù)必須完全相同。支持的 dtype 組合dstdtypesrc0dtypescalardtype說明halfhalfhalf同類型路徑直接vaxpyfloatfloatfloat同類型路徑直接vaxpyfloathalfhalf差異路徑src0拓寬為 FP32 后累加dst與src0必須 dtype 一致或dst為float且src0為half允許 half→float 的拓寬累加。dst為half而src0為float的組合非法由實(shí)現(xiàn)內(nèi)static_assert在編譯期攔截。這一 dtype 規(guī)則并非只在文檔中聲明而是被真實(shí)編碼進(jìn)源碼。在 include/pto/npu/a5/TAxpy.hpp 與 include/pto/npu/a2a3/TAxpy.hpp 的TAXPY_IMPL中均有兩條static_assertstatic_assert(std::is_same_vT, half || std::is_same_vT, float, TAXPY: Invalid data type); static_assert(std::is_same_vT, U || (std::is_same_vT, float std::is_same_vU, half), TAXPY: The data type of dst must be consistent with src or dst is float while src is half.);同時(shí)static_assert(TileDataDst::Loc TileType::Vec, ...)保證兩個(gè) Tile 都位于 UB統(tǒng)一緩沖區(qū)、向量流水線并附帶運(yùn)行期PTO_ASSERT檢查src0與dst的 valid row/col 一致。底層實(shí)現(xiàn)原理向量流水線上的 vaxpyTAXPY在向量流水線PIPE_V上執(zhí)行核心是vaxpy$a \cdot x y$向量?jī)?nèi)建。不同后端在實(shí)現(xiàn)路徑上有明顯差異這正是 PTO 跨平臺(tái)設(shè)計(jì)的關(guān)鍵體現(xiàn)。同類型路徑A5 后端在 include/pto/npu/a5/TAxpy.hpp 的AxpyInstrSame中按CeilDivision(validCol, elementsPerRepeat)計(jì)算 repeat 次數(shù)elementsPerRepeat CCE_VL / sizeof(T)然后逐行、逐 repeatvlds加載src0與dstNORM 模式用CreatePredicateT(sreg)構(gòu)造尾部謂詞掩碼屏蔽不足一個(gè) repeat 的列執(zhí)行vaxpy(vreg2, vreg0, scalar, preg)vsts寫回dst。差異類型路徑A5 后端UNPK_B16 vcvt在 include/pto/npu/a5/TAxpy.hpp 的AxpyInstrDiff中當(dāng)dst為float、src0為half時(shí)vlds(..., UNPK_B16)以半精度解包模式加載src0vcvt(reg_src_tmp, vreg0, preg, PART_EVEN)將 half 拓寬為 FP32再以(T)scalar執(zhí)行vaxpy并寫回。即 A5 上差異路徑需要顯式“解包 類型轉(zhuǎn)換 縮放累加”三步這與文檔中“src0 拓寬為 FP32 后累加”的描述完全對(duì)應(yīng)。A2/A3 后端count 模式與 norm 模式自適應(yīng)在 include/pto/npu/a2a3/TAxpy.hpp 中A2/A3 的實(shí)現(xiàn)先計(jì)算常量dstStride / dstBlockSizeElem 255 || srcStride / srcBlockSizeElem 255判定repeat-stride 是否溢出uint8_t上限 255validCol / elementsPerRepeat validRow判定列 repeat 數(shù)與行數(shù)的關(guān)系。隨后二選一count 模式AxpyCountModeinclude/pto/npu/a2a3/TAxpy.hppset_mask_count()SetVectorCount(validCol)逐行執(zhí)行一次vaxpy(dstPtr, src0Ptr, scalar, 0, 1, 1, 8, 4)——注意差異類型下 src 一個(gè) repeat 只占 4 個(gè) blockdst 占 8 個(gè) block由vaxpy原生處理 half→float 的 4-block src / 8-block dst 映射norm 模式AxpyNormMode/AxpyNormModeTailinclude/pto/npu/a2a3/TAxpy.hpp先按REPEAT_MAX切分行塊行內(nèi)按dstElementsPerRepeat循環(huán)切列剩余列用SetContMaskByDTypeU(numRemainAfterLoop)連續(xù)掩碼處理尾部。選擇邏輯useCountMode repeatStrideOverflow || validCol / elementsPerRepeat validRow保證任意有效形狀包括大行數(shù)、超長(zhǎng)列、窄行都能被覆蓋這與文檔中“按 repeat-stride 是否溢出、以及列數(shù)與行數(shù)的關(guān)系在 count 模式與 norm 模式間選擇”的描述一一對(duì)應(yīng)。CPU_SIM 參考實(shí)現(xiàn)模擬半精度舍入在 include/pto/cpu/TBinSOps.hpp 的 CPU_SIMTAXPY_IMPL中half/half路徑刻意模擬硬件上的兩步半精度行為volatile auto product static_casttypename TileDataDst::DType(src.data()[srcIdx] * scalar); volatile auto result static_casttypename TileDataDst::DType(dst.data()[dstIdx] product); dst.data()[dstIdx] result;即先把乘積舍入為half再與dst相加并把和再次舍入為half而不是按主機(jī)浮點(diǎn)精度一次性計(jì)算完整表達(dá)式。volatile保證兩次舍入不被編譯器優(yōu)化合并從而與 NPU 上的可觀測(cè)行為保持一致。float路徑則直接按dst src * scalar計(jì)算。約束一覽約束適用范圍原因dst、src0必須為TileType::Vec所有目標(biāo)在 UB向量流水線上執(zhí)行dst與src0有效形狀相同$M \times N$所有目標(biāo)逐元素一一對(duì)應(yīng)dstdtype ∈ {half,float}所有目標(biāo)vaxpy支持的浮點(diǎn)字寬dst/src0dtype 一致或 (float,half)所有目標(biāo)僅允許 half→float 拓寬累加dst調(diào)用前必須已初始化所有目標(biāo)dst作為累加基 $y$ 被讀入完整實(shí)戰(zhàn)示例NPU 上的 ST 內(nèi)核倉(cāng)庫(kù)在 tests/npu/a5/src/st/testcase/taxpy/taxpy_kernel.cpp 提供了完整的 ST 測(cè)試內(nèi)核展示了TAXPY在真實(shí) kernel 中的完整用法#include pto/pto-inst.hpp #include pto/common/constants.hpp #include acl/acl.h using namespace pto; template typename T, int kTRows_, int kTCols_, int vRows, int vCols __global__ AICORE void runTAxpy(__gm__ T __out__* out, __gm__ T __in__* src0, float scalar) { using DynShapeDim5 Shape1, 1, 1, vRows, vCols; using DynStridDim5 pto::Stride1, 1, 1, vCols, 1; using GlobalData GlobalTensorT, DynShapeDim5, DynStridDim5; using TileData TileTileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1; TileData src0Tile(vRows, vCols); TileData dstTile(vRows, vCols); TASSIGN(src0Tile, 0x0); TASSIGN(dstTile, 0x20000); GlobalData src0Global(src0); GlobalData dstGlobal(out); TLOAD(src0Tile, src0Global); TLOAD(dstTile, dstGlobal); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); TAXPY(dstTile, src0Tile, (T)scalar); // dst scalar * src0 dst set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); TSTORE(dstGlobal, dstTile); out dstGlobal.data(); }要點(diǎn)解讀dst必須先用TLOAD從 GM 載入作為累加基 $y$這與“調(diào)用前必須初始化”的約束一致TASSIGN為兩個(gè) Tile 指定 UB 上的起始地址set_flag/wait_flag是顯式的 MTE2→V 流水線事件同步TAXPY自身的WaitEvents參數(shù)在這里未被用到但包裝函數(shù)內(nèi)部也會(huì)做隱式同步結(jié)果通過TSTORE寫回 GM。該測(cè)試文件還覆蓋了多種形狀與 dtype 組合的實(shí)例化見 taxpy_kernel.cpp包括half的 64×64、63×63、1×16384、2048×16 以及float的 8×8、15×15 等從側(cè)面驗(yàn)證了“任意有效形狀均可覆蓋”的實(shí)現(xiàn)目標(biāo)。更多完整 ST 示例可參考tests/npu/a5/src/st/testcase/taxpy/A5tests/npu/a2a3/src/st/testcase/taxpy/A2/A3tests/npu/kirin9030/src/st/testcase/taxpy/Kirin9030tests/cpu/st/testcase/taxpy/CPU 參考實(shí)現(xiàn)相關(guān)資源指令文檔docs/isa/TAXPY.md / docs/isa/TAXPY_zh.md公共聲明與包裝函數(shù)include/pto/common/pto_instr.hppA5 后端實(shí)現(xiàn)include/pto/npu/a5/TAxpy.hppA2/A3 后端實(shí)現(xiàn)include/pto/npu/a2a3/TAxpy.hppCPU_SIM 參考實(shí)現(xiàn)include/pto/cpu/TBinSOps.hpp公開包含頭include/pto/pto-inst.hpp如需深入理解 Tile 編程模型與vaxpy之外的向量指令體系可繼續(xù)閱讀 docs/coding/ProgrammingModel.md 與 docs/isa/README.md?!久赓M(fèi)下載鏈接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.項(xiàng)目地址: https://gitcode.com/cann/pto-isa創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考