量向上取整運(yùn)算的用法與底層實(shí)現(xiàn)解析)
CANN PyPTOpl.simt.ceilSIMT 標(biāo)量向上取整運(yùn)算的用法與底層實(shí)現(xiàn)解析【免費(fèi)下載鏈接】pyptoPyPTO發(fā)音: pai p-t-oParallel Tensor/Tile Operation編程范式。項(xiàng)目地址: https://gitcode.com/cann/pypto導(dǎo)讀pypto_pro.language.simt.ceil即pl.simt.ceil是 CANN PyPTO 為 SIMT單指令多線程編程范式提供的標(biāo)量數(shù)學(xué)函數(shù)之一用于將浮點(diǎn)源操作數(shù)向正無窮方向舍入為不小于其本身的最小整數(shù)值。本文以 SIMT-API 標(biāo)量計(jì)算文檔 為主體完整梳理該 API 的函數(shù)原型、參數(shù)約束、產(chǎn)品支持情況與調(diào)用示例并結(jié)合倉庫中pl.simt命名空間的 API 聲明、IR 構(gòu)建器與解析器源碼講清ceil從前端調(diào)用到 IR 生成的完整鏈路。讀完本文你可以在 PyPTO SIMT 內(nèi)核中正確編寫基于pl.simt.ceil的線程級標(biāo)量運(yùn)算例如按 block_size 向上取整計(jì)算線程塊數(shù)量。功能說明與數(shù)學(xué)語義pl.simt.ceil對單個(gè)浮點(diǎn)標(biāo)量執(zhí)行向上取整ceiling其數(shù)學(xué)定義如下$$ result \lceil value \rceil $$即返回不小于源操作數(shù)的最小整數(shù)值向正無窮方向舍入輸入3.2→ 輸出4.0輸入-3.2→ 輸出-3.0注意不是向零舍入輸入4.0已是整數(shù)→ 輸出4.0。該語義與 CUDA 中的ceilf、C 標(biāo)準(zhǔn)庫ceil一致是向上取整而非四舍五入。倉庫中pl.simt命名空間還提供了語義各異的兄弟函數(shù)便于對照理解函數(shù)舍入規(guī)則源碼位置pl.simt.ceil向正無窮方向取整python/pypto_pro/language/_simt_api.pypl.simt.floor向負(fù)無窮方向取整同上pl.simt.trunc向零方向取整同上pl.simt.rint四舍五入到最近整數(shù)同上pl.simt.round遠(yuǎn)離零方向的四舍五入同上從 API 聲明文件 的 docstring 可見ceil被描述為 Round a floating-point Scalar up to an integer value與本文檔的定義完全對應(yīng)。產(chǎn)品支持情況根據(jù)文檔聲明pl.simt.ceil在不同昇騰硬件上的支持情況如下產(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)品不支持需要特別說明的是當(dāng)前倉庫中該 API 的支持范圍以文檔標(biāo)注為準(zhǔn)并非所有昇騰平臺都支持 SIMT 標(biāo)量ceil。在編寫可移植的內(nèi)核代碼時(shí)應(yīng)結(jié)合目標(biāo)硬件確認(rèn)支持情況避免在 Atlas A2/A3 系列上直接使用該接口。函數(shù)原型pypto_pro.language.simt.ceil( value: Scalar, ) - Scalar在 PyPTO 內(nèi)核源碼中通常通過import pypto_pro.language as pl簡寫為pl.simt.ceil(value)調(diào)用。參數(shù)說明參數(shù)輸入/輸出說明value輸入源操作數(shù)Scalar類型支持DT_FP16、DT_BF16和DT_FP32。若源數(shù)據(jù)來自Tensor或Tile需先通過下標(biāo)訪問取出元素后再傳入。兩個(gè)關(guān)鍵約束需要牢記類型約束value必須是浮點(diǎn)標(biāo)量類型FP16/BF16/FP32。整數(shù)類型標(biāo)量本身已是整數(shù)值無需也不支持調(diào)用本函數(shù)。元素訪問約束該函數(shù)接收的是Scalar不能直接傳入Tensor或Tile整體。文檔明確指出Tensor 或 Tile 元素需通過下標(biāo)訪問后傳入即先執(zhí)行類似tensor[0, tid]的下標(biāo)運(yùn)算得到標(biāo)量再交給ceil。返回值說明返回不小于源操作數(shù)的最小整數(shù)值數(shù)據(jù)類型與輸入value一致即 FP16 輸入返回 FP16BF16 輸入返回 BF16FP32 輸入返回 FP32。函數(shù)不改變數(shù)據(jù)的位寬與類型只改變數(shù)值。約束說明僅在 SIMT 上下文中調(diào)用pl.simt.ceil只能在由pypto_pro.language.vector_function(modesimt)定義的SIMT 入口函數(shù)或SIMT 輔助函數(shù)中調(diào)用。這一點(diǎn)與 PyPTO 的編程模型設(shè)計(jì)直接相關(guān)。從 decorator.py 的解析實(shí)現(xiàn) 可以看到vector_function裝飾器支持modesimt與modesimd兩種模式其中modesimt表示以多線程方式執(zhí)行、每個(gè)線程操作一個(gè)標(biāo)量本文ceil所在的范式max_threads參數(shù)僅在modesimt時(shí)被允許設(shè)置用于聲明該 SIMT 函數(shù)的線程規(guī)模若mode取值非法既非simt也非simd解析階段會(huì)直接拋出InvalidVal。因此pl.simt.ceil這類標(biāo)量運(yùn)算天然服務(wù)于線程級thread-per-scalar的 SIMT 編程模型與處理張量級運(yùn)算的 SIMD 向量編程模型對應(yīng)pl.simd.*/ 向量函數(shù)分工不同。調(diào)用示例按 block_size 向上取整計(jì)算線程塊數(shù)量原文檔給出的示例非常經(jīng)典用ceil把總長度 / 塊大小向上取整從而得到覆蓋全部數(shù)據(jù)的線程塊數(shù)量。完整代碼如下import pypto_pro.language as pl pl.vector_function(modesimt, max_threads256) def block_count( lengths: pl.Tensor[[1, 256], pl.DT_FP32], output: pl.Tensor[[1, 256], pl.DT_FP32], block_size: pl.DT_FP32, ): tid pl.simt.linear_thread_idx() output[0, tid] pl.simt.ceil(lengths[0, tid] / block_size) pl.jit() def simt_ceil_kernel( lengths: pl.Tensor[[1, 256], pl.DT_FP32], output: pl.Tensor[[1, 256], pl.DT_FP32], block_size: pl.DT_FP32, ): with pl.section_vector(): block_count256逐行拆解這個(gè)示例pl.vector_function(modesimt, max_threads256)聲明block_count是一個(gè) SIMT 輔助函數(shù)最多 256 個(gè)線程。它既不是宿主側(cè)的可執(zhí)行函數(shù)也不是 SIMD 向量函數(shù)而是供內(nèi)核內(nèi)部調(diào)用的線程級函數(shù)模板。tid pl.simt.linear_thread_idx()獲取當(dāng)前線程在 block 內(nèi)的一維扁平化線程號。該調(diào)用在 simt_ops.py 的 IR 構(gòu)建器 中被構(gòu)建為simt.linear_thread_idx操作語義是返回當(dāng)前 block 內(nèi) x-major 展開的線程索引。lengths[0, tid] / block_size通過下標(biāo)訪問lengths張量的第tid個(gè)元素得到Scalar除以block_size得到待取整的浮點(diǎn)標(biāo)量——這正是文檔Tensor 或 Tile 元素需通過下標(biāo)訪問后傳入約束的體現(xiàn)。pl.simt.ceil(...)對商向上取整將結(jié)果寫入output[0, tid]。pl.jit()with pl.section_vector():外層用pl.jit()編譯整個(gè)內(nèi)核pl.section_vector()是 PyPTO 的向量區(qū)段上下文用于把 SIMT 輔助函數(shù)的調(diào)用組織進(jìn)向量計(jì)算流水區(qū)。若在解析期出現(xiàn)pl.section_vector()/pl.section_cube()之外的區(qū)段用法解析器會(huì)報(bào)錯(cuò)見 parser/_control_flow_parser.py 中對應(yīng)檢查邏輯。block_count256以 256 線程為規(guī)模調(diào)用 SIMT 輔助函數(shù)。底層實(shí)現(xiàn)從pl.simt.ceil到 IR 的完整鏈路pl.simt.ceil并非普通的 Python 運(yùn)行期函數(shù)而是由 PyPTO 前端 AST 解析器在編譯期攔截并轉(zhuǎn)換為 IR 的聲明式 API。理解這一機(jī)制有助于避免誤用第一步API 聲明類型提示與文檔。在 python/pypto_pro/language/_simt_api.py 中ceil被定義為Simt類的靜態(tài)方法并用_api_decl裝飾。該文件頭部注釋明確說明這些聲明僅用于 IDE 跳轉(zhuǎn)、導(dǎo)入期拼寫檢查、類型校驗(yàn)與 docstring 展示在 PyPTO 內(nèi)核內(nèi)部由 AST 解析器在內(nèi)核執(zhí)行前攔截pl.simt.xxx調(diào)用在內(nèi)核之外直接調(diào)用會(huì)拋出RuntimeError。第二步IR 構(gòu)建器。在 python/pypto_pro/ir/op/simt_ops.py 中ceil通過_create_math_call(ceil, value, spanspan)構(gòu)建 IR。_create_math_call同文件 L107-L110的作用是把任意顯式 SIMT 標(biāo)量數(shù)學(xué)操作統(tǒng)一構(gòu)造為simt.op_name形式的 IR 調(diào)用節(jié)點(diǎn)def _create_math_call(op_name: str, *operands: Expr, span: Span | None) - Call: Create an IR call for one explicit SIMT scalar math operation. actual_span _get_span_or_capture(span) return _ir_core.create_op_call(fsimt.{op_name}, list(operands), {}, actual_span)即pl.simt.ceil(value)最終生成simt.ceil這一 IR 操作源碼位置信息span也會(huì)一并保留便于后續(xù)報(bào)錯(cuò)定位。第三步調(diào)用點(diǎn)解析。SIMT 輔助函數(shù)的調(diào)用由 parser/_call_parser.py 處理其中明確約束pl.vector_function(modesimt)輔助函數(shù)的調(diào)用只接受位置參數(shù)、不支持遞歸調(diào)用且返回值只能是None或單個(gè)標(biāo)量值。這與ceil的標(biāo)量入?yún)?標(biāo)量出參簽名天然契合。第四步后端代碼生成。simt.ceil這一 IR 操作最終由 C 側(cè)代碼生成器翻譯為昇騰設(shè)備上的標(biāo)量指令。從倉庫結(jié)構(gòu)看SIMT/標(biāo)量運(yùn)算的代碼生成實(shí)現(xiàn)在 framework/src/codegen 目錄如codegen_scalar.cpp、codegen_vector.cpp等ceil作為浮點(diǎn)標(biāo)量舍入運(yùn)算屬于其中的基礎(chǔ)數(shù)學(xué)指令族。與張量級pypto.ceil的區(qū)別需要區(qū)分的是PyPTO 還提供張量級的pypto.ceil(input: Tensor) - Tensor見 python/pypto/op/math.py 與 python/pypto/tensor.py 的Tensor.ceil方法。兩者的區(qū)別在于維度pl.simt.ceilpypto.ceil操作對象單個(gè)Scalar標(biāo)量整個(gè)Tensor編程范式SIMT線程級張量級算子Tiling/向量化自動(dòng)展開調(diào)用上下文SIMT 入口/輔助函數(shù)內(nèi)常規(guī) PyPTO 內(nèi)核含 loop/view/assemble 流程倉庫中的系統(tǒng)測試 python/tests/st/operation/vector/test_ceil.py 展示了張量級pypto.ceil的用法通過pypto.loop遍歷、pypto.view切出子視圖、pypto.ceil(input0_view)計(jì)算、再pypto.assemble寫回輸出。而本文檔聚焦的pl.simt.ceil則是 SIMT 線程編程中逐線程處理單個(gè)標(biāo)量的選擇。二者相輔相成向量化 Kernel 用pypto.ceil線程化 SIMT Kernel 用pl.simt.ceil。典型應(yīng)用場景pl.simt.ceil在 SIMT 內(nèi)核中的典型用途包括線程塊/分組數(shù)量計(jì)算如示例所示用ceil(length / block_size)保證分組向上取整、不遺漏任何數(shù)據(jù)邊界對齊將浮點(diǎn)尺寸向上取整到整數(shù)值配合對齊邏輯使用SIMT 內(nèi)部的逐線程標(biāo)量處理任何需要在每個(gè)線程內(nèi)對 FP16/BF16/FP32 標(biāo)量執(zhí)行向上取整的數(shù)學(xué)邏輯??偨Y(jié)pl.simt.ceil是 CANN PyPTO SIMT 編程范式中一個(gè)語義簡單但用途明確的標(biāo)量數(shù)學(xué) API向正無窮方向取整、支持 FP16/BF16/FP32、僅限 SIMT 上下文使用且目前僅在上昇騰 950PR/950DT 產(chǎn)品上支持。通過本文可以掌握其函數(shù)簽名、參數(shù)與返回值約束、完整的調(diào)用示例以及從前端 AST 解析、simt.ceilIR 構(gòu)建到后端標(biāo)量指令生成的實(shí)現(xiàn)鏈路如需進(jìn)一步探索可閱讀 SIMT-API 目錄下其他標(biāo)量計(jì)算文檔 或 pl.simt 命名空間源碼。【免費(fèi)下載鏈接】pyptoPyPTO發(fā)音: pai p-t-oParallel Tensor/Tile Operation編程范式。項(xiàng)目地址: https://gitcode.com/cann/pypto創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考