借助最新版本的 Warp 1.5.0 ，開發者現在可以使用 Python 中基于圖塊的新編程基元。這些新工具利用 cuBLASDx 和 cuFFTDx ，在 Python 內核中為開發者提供高效的矩陣乘法和 Fourier 變換，從而加速仿真和科學計算。在這篇博文中，我們將介紹這些新功能，并展示如何使用它們來優化應用。Warp 1.5.0 中提供的基于圖塊的編程模型目前處于預覽階段，在即將推出的版本中，性能和 APIs 可能會發生變化。
簡介?
在過去十年中，GPU 硬件已從單純的 SIMT (單指令多線程) 執行模型發展為高度依賴協作操作來提高效率的模型。隨著 Tensor Core 數學單元在整體 GPU 計算中的作用越來越大，高效且高效的編程變得越來越重要。高級 API 如 BLAS 提供的抽象概念可以面向各種高性能低級指令。但是，這些 API 通常難以與用戶程序集成，并且可能會在庫調用之間強制將結果傳回全局內存，從而降低效率。相反，直接在 C++/CUDA 級別對 Tensor Core 進行編程非常復雜，需要仔細管理單元之間的數據流。
為解決這些問題，我們開發了基于圖塊的編程模型，例如 OpenAI Triton 和 C++ AMP 中模型。與純 SIMT 模型不同，基于圖塊的編程允許開發者在圖塊上表示多個線程可以協同執行的操作，從而提高效率和生產力。
在 Warp 1.5.0 版本中，我們擴展了 Warp 基于內核的編程模型，將基于平鋪的操作包括在內，旨在讓 Warp 開發者能夠充分利用現代 GPU 硬件的全部功能。這些擴展程序：
- 提供編程模型，使開發者能夠從 SIMT 順利過渡到基于圖塊的執行。
- 減少對手動索引、共享內存管理和指針運算的需求。
- 支持反向傳播和訓練的自動微分。
此外，Warp 還利用 cuBLASDx 和 cuFFTDx 執行矩陣乘法和 Fast Fourier Transform (FFT) 圖塊運算。這些 NVIDIA 設備端數學庫與 Warp 的圖塊編程模型相結合，可在單個內核中無縫融合 Tensor Core 加速的 GEMM、FFT 和其他圖塊運算，從而減少內存 I/O 和內核啟動用度，同時更大限度地提高算術強度。借助這種方法，對于需要密集線性代數的應用 (例如機器人前向動力學)，我們的性能可以比傳統的線性代數或張量框架高出 4 倍。
扭曲圖塊基元?
Warp 中的新平鋪基元包括構建、加載/存儲、線性代數和映射/歸約運算，從而自然擴展了現有基于內核的編程模型。
建筑?
可以使用 NumPy 式操作在 Warp 內核中構建圖塊，如下所示：
import warp as wp@wp.kerneldef compute(): # construct a 16x16 tile of zeroed 32-bit floats a = wp.tile_zeros(m=16, n=16, dtype=wp.float32) # construct a 16x16 tile of 16-bit floats initialized to 1.0 b = wp.tile_ones(m=16, n=16, dtype=wp.float16) |
在 Warp 中，tiles 是二維數組，可能包含標量、向量、矩陣或結構化數據類型作為元素。與 Warp 數組或 PyTorch 張量不同，其中的維度是動態的并在運行時指定，tile 的維度 (例如，上述示例中的 16×16) 必須是編譯時已知的常量。此外，與線程的本地 SIMT 數據不同，tile 數據存儲在整個 CUDA 塊中，既可以存儲在寄存器中，也可以存儲在共享內存中。有關 tile 構建例程的完整列表，請單擊此處訪問 GitHub。
加載/存儲?
Warp 為全局內存中的平鋪數據提供顯式加載/存儲操作。塊中的所有線程協同執行這些操作，確保全局內存與共享或寄存器內存之間的高效數據傳輸。在以下示例中，系統從全局內存加載兩個數據圖塊，并對其求和，然后將結果存儲回全局內存。用戶無需明確管理共享內存分配或存儲。
import warp as wp@wp.kerneldef compute(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)): # cooperatively load input tiles a = wp.tile_load(A, i=0, j=0, m=16, n=16) b = wp.tile_load(B, i=0, j=0, m=16, n=16) # compute sum c = a + b # cooperatively store sum to global memory wp.tile_store(C, i=0, j=0, t=c)A = wp.ones((16,16), dtype=float)B = wp.ones((16,16), dtype=float)C = wp.empty((16,16), dtype=float)wp.launch_tiled(compute, dim=1, inputs=[A, B, C], device="cuda:0", block_dim=64) |
除了加載/存儲操作外，Warp 還支持 wp.tile_atomic_add() 等原子操作。有關內存操作的完整列表，請參閱以下 文檔 。
矩陣乘法?
基于圖塊的編程的主要優勢之一是能夠執行協作矩陣乘法。Warp 1.5.0 引入了通用乘積累加基元 wp.tile_matmul()，允許開發者執行協同矩陣乘法。這在幕后利用 cuBLASDx ，根據元素類型、矩陣大小和數據布局，cuBLASDx 將自動發送適當的 Tensor Core MMA 指令，以實現最佳性能。
我們來看看在 Warp 中使用基于圖塊的編程來執行矩陣乘法的示例：
import warp as wpTILE_M = wp.constant(32)TILE_N = wp.constant(64)TILE_K = wp.constant(64)@wp.kerneldef gemm_tiled(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)): i, j = wp.tid() # allocate output tile sum = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=float) count = int(K / TILE_K) # iterate over inner dimension for k in range(count): a = wp.tile_load(A, i, k, m=TILE_M, n=TILE_K) b = wp.tile_load(B, k, j, m=TILE_K, n=TILE_N) # perform gemm + accumulate wp.tile_matmul(a, b, sum) # store result wp.tile_store(C, i, j, sum)# test with 1024^2 inputsM, N, K = 1024, 1024, 1024A = wp.ones((M, K), dtype=float)B = wp.ones((K, N), dtype=float)C = wp.empty((M, N), dtype=float)# launch kernel with 128 threads per-blockwp.launch_tiled(gemm_tiled, dim=(int(M//TILE_M), int(N//TILE_N)), inputs=[A, B, C], block_dim=128) |
在本示例中，我們定義了執行平鋪矩陣乘法的核 gemm_tiled()。內核循環處理全局內存中的 2D 數據片，將其加載到共享內存圖塊中，使用 wp.tile_matmul() 執行矩陣乘法，在共享內存中累加結果，然后將結果存儲回全局內存。
下圖顯示了上述 GEMM 內核在 NVIDIA A100 80GB SXM（時鐘鎖定至其最大值）上與 cuBLAS 12.4 相比在各種 FP32 矩陣大小下的性能百分比。對于小問題，我們發現性能與 cuBLAS 具有競爭力，這可能是因為我們使用自動調整來為這種小尺寸找到最佳參數，而且啟動開銷在成本中占更大的比例。對于更大的問題，由于目前 tile 結果始終存儲在共享內存中，因此性能會降低。然而，即使對于這個簡單的示例，我們也會看到較大矩陣的 cuBLAS 性能約為 70–80%。未來版本的 Warp 和 cuBLASDx 將在寄存器中保留 GEMM 的輸出，從而提高性能。
下圖中，我們看看圖塊大小對單個問題大小的整體性能的影響。整體性能是圖塊維度和塊維度的函數，前者決定問題的分解方式，后者決定向每個子問題分配多少線程。在這里，我們可以看到，對于 M = N = K = 1024 問題，使用 TILE_M = 32、TILE_N = 64、TILE_K = 64 和 128 線程的平鋪維度獲得了最佳性能。Warp 的動態編程和運行時內核創建允許用戶輕松執行超參數自動調整，如基準腳本示例所示。
請參閱此 參考 資料，了解平鋪線性代數基元的完整列表。
映射/歸約?
Warp 1.5.0 還包含 map/reduce 基元，支持開發者在圖塊上執行歸約和元素級運算。這些基元對于 LayerNorm 和 SoftMax 等任務至關重要，這些任務需要對不同數量進行高效歸約。
以下示例展示了如何使用每行一個 CUDA 塊和 wp.tile_sum() 計算數組行中所有元素的總和，以執行協作歸約。
import warp as wp@wp.kerneldef row_sum(input: wp.array2d(dtype=float), output: wp.array1d(dtype=float)): # obtain our block index i = wp.tid() # load a row of 256 elements from global memory t = wp.tile_load(input[i], i=0, n=256) # cooperatively sum elements s = wp.tile_sum(t) # store sum to output wp.tile_store(output, i, s) a = wp.ones((1024, 256), dtype=float)b = wp.empty(1024, dtype=float)wp.launch_tiled(row_sum, dim=[a.shape[0]], inputs=[a, b], block_dim=64) |
Warp 還支持自定義歸約運算符，在本例中，我們使用內置的 wp.tile_reduce() 和 wp.mul() 來計算因子，不過也可以使用用戶定義的 @wp.func 歸約運算符。
import warp as wp@wp.kerneldef factorial(): t = wp.tile_arange(1, 10, dtype=int) s = wp.tile_reduce(wp.mul, t) # prints "tile(m=1, n=1, storage=register) = [[362880]]" print(s)wp.launch(factorial, dim=[16], inputs=[], block_dim=16) |
此處 提供了 map/reduce 基元的完整列表。
案例研究?
融合神經網絡?
基于圖塊的編程還支持高效實現融合多層感知器（MLP）。以下是在 Warp 中使用基于圖塊的編程的融合 MLP 示例：
import warp as wpDIM_IN = wp.constant(4)DIM_HID = wp.constant(32)DIM_OUT = wp.constant(3)@wp.kerneldef mlp_fused(weights_0: wp.array2d(dtype=wp.float16), weights_1: wp.array2d(dtype=wp.float16), loss: wp.array(dtype=float)): t = wp.tid() # construct simple positional encoding x = wp.vec4h(wp.sin(x), wp.cos(x), wp.sin(x*2.0), wp.cos(x*2.0)) # tile input across block to create feature vectors f = wp.tile(x) # fully connected layer 0 w0 = wp.tile_load(weights_0, 0, 0, m=DIM_HID, n=DIM_IN) z = wp.tile_map(relu, wp.tile_matmul(w0, f)) # fully connected layer 1 w1 = wp.tile_load(weights_1, 0, 0, m=DIM_OUT, n=DIM_HID) z = wp.tile_map(relu, wp.tile_matmul(w1, z)) # loss function l = wp.tile_sum(z) wp.atomic_add(loss, 0, l)wp.launch(mlp_fused, dim=(1,), inputs=[weights_0, weights_1, loss], block_dim=128) |
在此示例中 ，mlp_fused() 內核通過加載權重、執行矩陣乘法、使用 wp.tile_map() 應用激活函數以及計算損失來評估簡單的兩層神經網絡，所有這些操作均在單個內核中完成。下圖展示了使用此方法對圖像進行編碼的示例。由于 Warp 支持自動微分，因此我們可以直接評估和訓練網絡權重，以學習從圖像坐標 (x,y) 映射到像素顏色 (RGB) 的函數。完整示例請參見此處。
信號處理?
線程束平鋪運算集成 cuFFTDx，用于內核內的正向和反向 FFT，在數據塊上提供高效的 Fourier-transform 運算。以下是在 Warp 中使用基于圖塊的 FFT 來使用濾波器計算卷積的示例：
import warp as wp@wp.kerneldef conv_tiled(x: wp.array2d(dtype=wp.vec2d), y: wp.array2d(dtype=wp.vec2d), z: wp.array2d(dtype=wp.vec2d)):i, j = wp.tid()# load signal and filtera = wp.tile_load(x, i, j, m=TILE_M, n=TILE_N)f = wp.tile_load(y, i, j, m=TILE_M, n=TILE_N)# compute Fourier transform of input signalwp.tile_fft(a)# compute filter in frequency spacec = wp.tile_map(cplx_prod, a, b)# convert back to realwp.tile_ifft(c)wp.tile_store(z, i, j, c) |
在此示例中，conv_tiled() 內核 (沿最后一個維度) 執行一塊數據的前向 FFT，應用過濾器，然后計算反向 FFT。在幕后，cuFFTDx 用于實現。完整示例可在此處找到。下圖顯示了對噪聲輸入信號應用濾波器的輸出。
機器人前向動力學?
基于圖塊的編程也非常有利于需要密集線性代數的模擬應用。在機器人仿真中，復合剛體算法 (CRBA) 方法用于計算關節機制的前向動力學。在 CRBA 方法中，需要以下三矩陣乘積，其中內矩陣 M 是塊稀疏對角線質量矩陣：
構建完成后，系統矩陣使用 Cholesky 分解進行分解，并使用前向和反向替換求解。我們可以使用 Warp 的 tile 基元來利用 M 的稀疏性，將此問題的批處理版本表述如下：
import warp as wp@wp.kerneldef foward_dynamics( J_arr: wp.array3d(dtype=float), M_arr: wp.array3d(dtype=float), R_arr: wp.array3d(dtype=float), H_arr: wp.array3d(dtype=float), L_arr: wp.array3d(dtype=float),): batch = wp.tid() J = wp.tile_load(J_arr[batch], 0, 0, m=wp.static(6 * num_joints), n=num_dofs) P = wp.tile_zeros(m=wp.static(6 * num_joints), n=num_dofs, dtype=float) # compute P = M*J where M is a 6x6 block diagonal mass matrix for i in range(int(num_joints)): # 6x6 block matrices are on the diagonal M_body = wp.tile_load(M_arr[batch], i, i, m=6, n=6) # load a 6xN row from the Jacobian J_body = wp.tile_view(J, i * 6, 0, m=6, n=num_dofs) # compute weighted row P_body = wp.tile_matmul(M_body, J_body) # assign to the P slice wp.tile_assign(P, i * 6, 0, P_body) # compute H = J^T*P H = wp.tile_matmul(wp.tile_transpose(J), P) # cholesky L L^T = (H + diag(R)) R = wp.tile_load(R_arr[batch], 0, 0, m=num_dofs, n=1, storage="shared") H += wp.tile_diag(R) L = wp.tile_cholesky(H) wp.tile_store(L_arr[batch], 0, 0, L)# launch kernel with 64 threads per-robotwp.launch_tiled(forward_dynamics, dim=(num_robots,), inputs=[J_arr, M_arr, R_arr, H_arr, L_arr], block_dim=64) |
在本例中，forward_dynamics() 內核執行 CRBA 方法，加載 Jacobian 矩陣和質量矩陣的圖塊，并計算其積以形成系統矩陣 H 及其 Cholesky 分解。在此特定用例中，Torch 需要啟動十幾個內核，而 Warp 實現則需要一個完全融合的內核。這可減少全局內存往返次數和啟動開銷，從而顯著提高性能。
對于在 NVIDIA A100 80GB GPU 上運行的前向動態內核，1,024 個四足機器人的性能如下，所有計時均以毫秒為單位 (越低越好)：
這里 提供了完整的示例，用于展示 Cholesky 分解和反向替換的擴展程序即將推出。
未來發展?
Warp 和 MathDx 的未來版本將包括：
- 對逐行歸約運算符的額外支持
- 從 Lambda 函數創建圖塊
- 數據類型和布局轉換
- 提高 GEMM 運算的性能
- 其他線性代數基元，包括各種矩陣分解算法。
了解詳情?
Warp 1.5.0 中基于圖塊的編程提供了一種強大而靈活的 GPU 編程方法，使開發者能夠顯著提升其應用的性能。通過利用 cuBLASDx 和 cuFFTDx，Warp 1.5.0 可實現 GEMM 和 FFT 操作的無縫融合，從而減少內存 I/O 和內核啟動用度。
要開始使用 Warp 的 Tile 操作，請在 Python 環境中使用以下命令安裝 Warp：
pip install warp-lang |
要運行融合的 MLP 示例，請使用以下命令：
python -m warp.examples.tile.example_mlp.py |
相關資源?
如需詳細了解 Warp 1.5.0 和 NVIDIA Math device acceleration (Dx) 庫，請訪問以下鏈接：
- Warp GTC Presentation: https://www.nvidia.com/en-us/on-demand/session/gtc24-s63345/
- 扭曲源 ： https://github.com/nvidia/warp
- Warp Tile 文檔：https://nvidia.github.io/warp/modules/tiles.html
- cuBLASDx : http://www.open-lab.net/cublasdx-downloads
- cuFFTDx: http://www.open-lab.net/cufftdx-downloads
致謝?
感謝 Pawe? Grabowski、Doris Pan、Neil Lindquist、Jakub Szuppe、?ukasz Ligowski、Sergey Maydanov 和 ?ukasz Wawrzyniak 對此帖子和項目的貢獻。
?
