Taichi 性能調(diào)優(yōu)完全指南循環(huán)指令、數(shù)據(jù)布局、TLS/BLS 局部存儲(chǔ)優(yōu)化與離線緩存【免費(fèi)下載鏈接】taichiProductive, portable, and performant GPU programming in Python.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/ta/taichiTaichi 是面向 Python 的生產(chǎn)力、可移植性與高性能 GPU 編程框架其編譯器會(huì)自動(dòng)并行化最外層 for 循環(huán)。本文圍繞官方性能調(diào)優(yōu)文檔系統(tǒng)講解ti.loop_config循環(huán)指令的線程級(jí)調(diào)優(yōu)、GPU 線程層級(jí)背景、稀疏場(chǎng)數(shù)據(jù)布局、TLS 線程局部存儲(chǔ)與 BLS 塊局部存儲(chǔ)優(yōu)化原理以及ti.init()中離線緩存Offline Cache的配置策略幫助開發(fā)者在 MPM、流體、渲染等典型負(fù)載上榨取最后幾個(gè)百分點(diǎn)的性能。For-loop 指令細(xì)粒度控制循環(huán)并行度Taichi 內(nèi)核ti.kernel會(huì)自動(dòng)將最外層作用域的 for 循環(huán)并行化編譯器默認(rèn)按目標(biāo)架構(gòu)自動(dòng)設(shè)定調(diào)度參數(shù)。但對(duì)于追求極限性能的開發(fā)者Taichi 提供了ti.loop_configAPI 來微調(diào)循環(huán)指令。例如為循環(huán)指定合適的block_dim在 mpm3d.py 這類 MPM 模擬中可能帶來接近 3 倍的性能提升。ti.loop_config用于設(shè)置下一個(gè)for 循環(huán)的指令可用指令包括parallelize設(shè)置 CPU 上使用的線程數(shù)block_dim設(shè)置 GPU 上一個(gè) block 內(nèi)的線程數(shù)serialize若設(shè)為Truefor 循環(huán)將串行執(zhí)行且可在循環(huán)體內(nèi)寫 break 語句僅對(duì) range/ndrange for 生效等價(jià)于將parallelize設(shè)為 1從 python/taichi/lang/misc.py 的loop_config實(shí)現(xiàn)可以看到這些指令最終分別映射到底層 AST builder 的parallelize(v)、block_dim(dim)、strictly_serialize()調(diào)用當(dāng)parallelize 1時(shí)還會(huì)自動(dòng)調(diào)用strictly_serialize()強(qiáng)制串行化這與文檔中serialize 等價(jià)于 parallelize1的描述完全一致。ti.kernel def break_in_serial_for() - ti.i32: a 0 ti.loop_config(serializeTrue) for i in range(100): # This loop runs serially a i if i 10: break return a break_in_serial_for() # returns 55n 128 val ti.field(ti.i32, shapen) ti.kernel def fill(): ti.loop_config(parallelize8, block_dim16) # If the kernel is run on the CPU backend, 8 threads will be used to run it # If the kernel is run on the CUDA backend, each block will have 16 threads. for i in range(n): val[i] i此外loop_config還支持另外兩個(gè)可選指令見源碼 docstringblock_dim_adaptive是否允許后端自適應(yīng)設(shè)置 block_dim默認(rèn)開啟注意自適應(yīng) block_dim 僅在 CPU 后端生效其他后端會(huì)給出警告bit_vectorize是否對(duì) quant_array 上的 struct for 啟用位向量化例如將 32 個(gè) 1-bit 元素一次性拷貝背景GPU 的線程層級(jí)為了理解 for 循環(huán)如何被并行化需要先了解當(dāng)代 GPU 架構(gòu)上的線程層級(jí)Thread Hierarchy。從細(xì)粒度到粗粒度計(jì)算單元依次為iteration迭代、thread線程、block線程塊、grid網(wǎng)格。Iterationfor 循環(huán)的循環(huán)體即一次迭代每次迭代對(duì)應(yīng) for 循環(huán)中不同的i值Thread迭代被歸類為線程。線程是最小的并行單元一個(gè)線程內(nèi)的所有迭代串行執(zhí)行。為最大化并行效率通常采用一次迭代對(duì)應(yīng)一個(gè)線程的策略Block線程被組織成塊block塊內(nèi)所有線程并行執(zhí)行且可共享塊內(nèi)局部存儲(chǔ)block local storage即 CUDA 的 shared memoryGrid塊被組成網(wǎng)格grid網(wǎng)格是從主機(jī)端啟動(dòng)的最小單元網(wǎng)格內(nèi)所有塊并行執(zhí)行。在 Taichi 中每一個(gè)被并行化的 for 循環(huán)都表示為一個(gè) grid這里采用 CUDA 的術(shù)語其他后端如 OpenGL、Metal遵循類似的線程層級(jí)結(jié)構(gòu)。示例調(diào)整 for 循環(huán)的塊級(jí)并行度開發(fā)者可以在 for 循環(huán)前通過ti.loop_config調(diào)整其屬性指令只作用于緊接著的下一個(gè)循環(huán)ti.kernel def func(): for i in range(8192): # no decorator, use default settings ... ti.loop_config(block_dim128) # change the property of next for-loop: for i in range(8192): # will be parallelized with block_dim128 ... for i in range(8192): # no decorator, use default settings ...數(shù)據(jù)布局Data Layouts由于 Taichi 將數(shù)據(jù)結(jié)構(gòu)與計(jì)算分離開發(fā)者可以自由實(shí)驗(yàn)不同的數(shù)據(jù)布局。與其他編程語言一樣選擇高效的數(shù)據(jù)布局可能顯著提升性能。關(guān)于 Taichi 中高級(jí)數(shù)據(jù)布局如結(jié)構(gòu)體數(shù)組 SoA、數(shù)組結(jié)構(gòu) AoS、稠密/指針等 SNode 層級(jí)、布局推導(dǎo)與手動(dòng)布局的完整討論請(qǐng)參閱 Fields (advanced) 章節(jié)。局部存儲(chǔ)優(yōu)化Local Storage OptimizationsTaichi 內(nèi)置若干利用快速內(nèi)存如 CUDA shared memory、L1 cache的加速優(yōu)化。簡單來說只要可行Taichi 就會(huì)把對(duì)全局內(nèi)存慢的訪問替換為對(duì)局部內(nèi)存快的訪問并在結(jié)束時(shí)把局部內(nèi)存如 CUDA shared memory中的數(shù)據(jù)寫回全局內(nèi)存。這類變換保持原始程序的語義不變。線程局部存儲(chǔ)TLSTLS 主要用于優(yōu)化并行歸約parallel reduction。當(dāng) Taichi 在ti.kernel中識(shí)別出全局歸約模式時(shí)會(huì)在代碼生成階段自動(dòng)應(yīng)用 TLS 優(yōu)化其做法與常見 GPU 歸約實(shí)現(xiàn)類似。以下用 CUDA 術(shù)語演示一個(gè)例子x ti.field(ti.f32, shape1000000) s ti.field(ti.f32, shape()) ti.kernel def sum(): for i in x: s[None] x[i] sum()Taichi 的并行循環(huán)在內(nèi)部基于Grid-Stride Loops實(shí)現(xiàn)每個(gè)物理 CUDA 線程可能處理x中的多個(gè)元素也就是說為sum啟動(dòng)的線程數(shù)可以少于x的 shape。TLS 優(yōu)化正是利用這一點(diǎn)不再直接、原子地將x[i]加到全局內(nèi)存目標(biāo)s[None]上而是進(jìn)入線程時(shí)預(yù)分配一個(gè)線程局部緩沖區(qū)將x的值非原子地累加進(jìn)該緩沖區(qū)在退出線程前再把緩沖區(qū)結(jié)果原子地加回s[None]。如果每個(gè)線程處理N個(gè)元素原子加的次數(shù)就降為原來的1/N。此外最后對(duì)全局內(nèi)存s[None]的原子加還利用了 CUDA 的 warp 級(jí)內(nèi)建指令進(jìn)一步優(yōu)化進(jìn)一步減少所需原子加次數(shù)。當(dāng)前 TLS 的適用范圍支持add、sub、min、max這四類歸約算子作用于0D標(biāo)量/向量/矩陣ti.field尚不支持ti.ndarray。在 Nvidia GeForce RTX 3090 上對(duì)一維 8M float 的 Taichi field 做全局 max 歸約的基準(zhǔn)對(duì)比TLS 關(guān)閉5.2 × 1e3 usTLS 開啟5.7 × 1e1 usTLS 帶來了約100 倍加速且 TLS 歸約求和性能與 CUDA 實(shí)現(xiàn)相當(dāng)。塊局部存儲(chǔ)BLS適用場(chǎng)景對(duì)于最后一層為denseSNode 的稀疏場(chǎng)即層級(jí)結(jié)構(gòu)匹配ti.root.(sparse SNode).denseTaichi 會(huì)將每個(gè)dense容器或denseblock分配給一個(gè) CUDA 線程塊。BLS 優(yōu)化專門針對(duì)這類場(chǎng)生效。BLS 旨在利用 CUDA shared memory 加速模板計(jì)算stencil computation。優(yōu)化流程如下用戶通過ti.block_local標(biāo)注想要緩存的場(chǎng)集合在編譯期Taichi 嘗試識(shí)別這些標(biāo)注場(chǎng)相對(duì)于denseblock 的訪問范圍若識(shí)別成功Taichi 生成代碼先把范圍內(nèi)可訪問的數(shù)據(jù)全部加載到block local緩沖區(qū)CUDA shared memory再將循環(huán)體中對(duì)相應(yīng)槽位的訪問全部替換為該緩沖區(qū)下面演示 BLS 的用法a是一個(gè) block size 為4x4的稀疏場(chǎng)a ti.field(ti.f32) b ti.field(ti.f32) # a has a block size of 4x4 ti.root.pointer(ti.ij, 32).dense(ti.ij, 4).place(a) ti.kernel def foo(): # Taichi will cache a into the CUDA shared memory ti.block_local(a) for i, j in a: print(a[i - 1, j], a[i, j 2])每個(gè)循環(huán)迭代訪問的元素相對(duì)于自身坐標(biāo)的偏移分別是[-1, 0]和[0, 2]。因此對(duì)于一個(gè)從[M, N]含到[M 4, N 4]不含的完整 block其相對(duì)于該 block 的訪問范圍是[M - 1, M 4) x [N, N 6)由[M (-1), M 4) x [N, N 4 2)推導(dǎo)得出。全局坐標(biāo)i, j與緩沖區(qū)局部索引的映射關(guān)系如下圖作為用戶你無需關(guān)心這些底層細(xì)節(jié)——Taichi 自動(dòng)完成所有推斷以及全局/塊局部映射預(yù)分配一塊5x6的 CUDA shared memory 緩沖區(qū)將a的內(nèi)容預(yù)加載進(jìn)緩沖區(qū)然后把循環(huán)體內(nèi)的所有a全局內(nèi)存訪問替換為緩沖區(qū)訪問。上面的例子沒有修改a本身而如果被塊緩存的場(chǎng)發(fā)生了寫入Taichi 會(huì)生成把緩沖區(qū)寫回全局內(nèi)存的代碼。注意BLS 并非沒有代價(jià)。BLS 面向的是具有大量重疊全局內(nèi)存訪問的模板計(jì)算。如果不符合這一條件預(yù)加載與回寫反而可能降低性能。此外新一代 Nvidia GPU 已縮小了全局內(nèi)存與 shared memory 在只讀訪問上的差距目前我們觀察到 BLS 對(duì)于緩存原子操作的目標(biāo)地址更有效。作為經(jīng)驗(yàn)法則建議通過基準(zhǔn)測(cè)試來判斷是否啟用 BLS。從 python/taichi/lang/misc.py 的block_local實(shí)現(xiàn)可見它會(huì)把場(chǎng)成員標(biāo)記為SNodeAccessFlag.block_local且當(dāng)優(yōu)化級(jí)別為 0 時(shí)會(huì)強(qiáng)制提升到opt_level 1以啟用 BLS 分析這也印證了 BLS 需要依賴編譯器優(yōu)化管線。離線緩存Offline CacheTaichi 內(nèi)核在第一次被調(diào)用時(shí)會(huì)被隱式編譯。為降低后續(xù)調(diào)用的開銷編譯結(jié)果會(huì)保存在在線的內(nèi)存緩存中只要內(nèi)核未改變即可立即加載并啟動(dòng)。但應(yīng)用退出后該緩存即失效重啟程序后 Taichi 必須重新編譯所有內(nèi)核例程并重建在線內(nèi)存緩存——由于編譯開銷Taichi 函數(shù)的首次啟動(dòng)通常較慢。離線緩存功能解決了這一問題它將編譯緩存轉(zhuǎn)儲(chǔ)到磁盤供后續(xù)運(yùn)行使用從而大幅降低重復(fù)運(yùn)行時(shí)的首次啟動(dòng)開銷。Taichi 默認(rèn)構(gòu)建并維護(hù)離線緩存同時(shí)提供若干ti.init()選項(xiàng)用于配置其行為offline_cache: bool啟用或禁用離線緩存。默認(rèn)值Trueoffline_cache_file_path: str存放離線緩存文件的目錄。Windows 上默認(rèn)C:\taichi_cache\ticache\Unix 系系統(tǒng)默認(rèn)~/.cache/taichi/ticache/。目錄會(huì)自動(dòng)創(chuàng)建offline_cache_max_size_of_files: int32緩存文件的最大字節(jié)數(shù)。默認(rèn)值100MB。當(dāng)緩存文件大小超過該限制時(shí)觸發(fā)清理進(jìn)程offline_cache_cleaning_policy: str緩存中過時(shí)文件的替換策略?？蛇x值never、version、lru、fifo。默認(rèn)值lrunever從不清理無論offline_cache_max_size_of_files配置如何都保留所有緩存文件version僅丟棄與內(nèi)核函數(shù)相關(guān)的舊版本緩存文件lru丟棄最近最少使用的緩存文件fifo丟棄最早加入的緩存文件例如在ti.init()中開啟離線緩存并指定緩存目錄ti.init(archti.gpu, offline_cacheTrue, offline_cache_file_path./ticache, offline_cache_max_size_of_files2 * 1024 * 1024 * 1024, offline_cache_cleaning_policylru)在 python/taichi/lang/misc.py 中可以看到前端默認(rèn)將cfg.offline_cache True與文檔所述默認(rèn)啟用一致。驗(yàn)證效果將示例程序運(yùn)行兩次并觀察啟動(dòng)開銷結(jié)果如下圖所示藍(lán)色為禁用離線緩存橙色為首次啟用離線緩存綠色為第二次啟用可以看到MPM99、MPM3D、Cornell Box 等示例在第二次運(yùn)行綠色柱時(shí)啟動(dòng)時(shí)間均趨近于 0 秒而禁用緩存時(shí)則需要數(shù)秒至數(shù)十秒的編譯時(shí)間。注意如果你的代碼行為異?？赏ㄟ^設(shè)置環(huán)境變量TI_OFFLINE_CACHE0或在ti.init()中設(shè)置offline_cacheFalse來禁用離線緩存并向 Taichi 倉庫提交 issue 反饋。小結(jié)與實(shí)踐建議循環(huán)級(jí)調(diào)優(yōu)使用ti.loop_config(parallelize..., block_dim...)為 CPU 指定線程數(shù)、為 GPU 指定每塊線程數(shù)需要 break 時(shí)用serializeTrue。注意指令只作用于緊接著的下一個(gè) for 循環(huán)。數(shù)據(jù)布局善用 Taichi 數(shù)據(jù)與計(jì)算分離的特性參考 Fields (advanced) 選擇高效布局。歸約負(fù)載優(yōu)先使用 0Dti.field上的add/sub/min/max歸約讓 TLS 自動(dòng)生效ti.ndarray目前不支持 TLS。模板計(jì)算對(duì)符合ti.root.(sparse).dense層級(jí)且訪存重疊度高的 stencil 核函數(shù)用ti.block_local配合基準(zhǔn)測(cè)試決定是否啟用 BLS。重復(fù)啟動(dòng)場(chǎng)景保持默認(rèn)的離線緩存開啟必要時(shí)通過ti.init()或TI_OFFLINE_CACHE環(huán)境變量調(diào)整緩存目錄、大小上限與清理策略?！久赓M(fèi)下載鏈接】taichiProductive, portable, and performant GPU programming in Python.項(xiàng)目地址: https://gitcode.com/GitHub_Trending/ta/taichi創(chuàng)作聲明：本文部分內(nèi)容由AI輔助生成（AIGC），僅供參考