如何使用Warp在Python環(huán)境中編寫(xiě)CUDA內(nèi)核

通常，實(shí)時(shí)物理模擬代碼是用低級(jí) CUDA C ++編寫(xiě)的，以獲得最佳性能。在這篇文章中，我們將介紹 NVIDIA Warp ，這是一個(gè)新的 Python 框架，可以輕松地用 Python 編寫(xiě)可微圖形和模擬 GPU 代碼。 Warp 提供了編寫(xiě)高性能仿真代碼所需的構(gòu)建塊，但它的工作效率與 Python 等解釋語(yǔ)言相當(dāng)。

在這篇文章的最后，您將學(xué)習(xí)如何使用 Warp 在 Python 環(huán)境中編寫(xiě) CUDA 內(nèi)核，并利用一些內(nèi)置的高級(jí)功能，從而輕松編寫(xiě)復(fù)雜的物理模擬，例如海洋模擬。

安裝

Warp 以 來(lái)自 GitHub 的開(kāi)源庫(kù) 的形式提供?？寺〈鎯?chǔ)庫(kù)后，可以使用本地軟件包管理器進(jìn)行安裝。對(duì)于 pip ，請(qǐng)使用以下命令：

pip install warp

初始化

導(dǎo)入后，必須顯式初始化扭曲：

import warp as wp
wp.init()

推出內(nèi)核

Warp 使用 Python 裝飾器的概念來(lái)標(biāo)記可以在 GPU 上執(zhí)行的函數(shù)。例如，可以編寫(xiě)一個(gè)簡(jiǎn)單的半隱式粒子積分方案，如下所示：

@wp.kernel
def integrate(x: wp.array(dtype=wp.vec3), v: wp.array(dtype=wp.vec3), f: wp.array(dtype=wp.vec3), w: wp.array(dtype=float), gravity: wp.vec3, dt: float): # thread id tid = wp.tid() x0 = x[tid] v0 = v[tid] # Semi-implicit Euler step f_ext = f[tid] inv_mass = w[tid] v1 = v0 + (f_ext * inv_mass + gravity) * dt x1 = x0 + v1 * dt # store results x[tid] = x1 v[tid] = v1 

因?yàn)?Warp 是強(qiáng)類型的，所以應(yīng)該為內(nèi)核參數(shù)提供類型提示。要啟動(dòng)內(nèi)核，請(qǐng)使用以下語(yǔ)法：

 wp.launch(kernel=simple_kernel, # kernel to launch dim=1024, # number of threads inputs=[a, b, c], # parameters device="cuda") # execution device

與 NumPy 等基于張量的框架不同， Warp 使用 kernel-based 編程模型。基于內(nèi)核的編程與底層 GPU 執(zhí)行模型更為匹配。對(duì)于需要細(xì)粒度條件邏輯和內(nèi)存操作的模擬代碼，這通常是一種更自然的表達(dá)方式。然而， Warp 以一種易于使用的方式公開(kāi)了這種以線程為中心的編程模型，它不需要 GPU 體系結(jié)構(gòu)的低級(jí)知識(shí)。

編譯模型

啟動(dòng)內(nèi)核會(huì)觸發(fā)實(shí)時(shí)（ JIT ）編譯管道，該管道會(huì)自動(dòng)從 Python 函數(shù)定義生成 C ++/ CUDA 內(nèi)核代碼。

屬于 Python 模塊的所有內(nèi)核都在運(yùn)行時(shí)編譯到動(dòng)態(tài)庫(kù)和 PTX 中。圖 2 。顯示了編譯管道，其中包括遍歷函數(shù) AST 并將其轉(zhuǎn)換為直線 CUDA 代碼，然后編譯并加載回 Python 進(jìn)程。

圖 2 。 Warp 內(nèi)核的編譯管道

這個(gè) JIT 編譯的結(jié)果被緩存。如果輸入內(nèi)核源代碼不變，那么預(yù)編譯的二進(jìn)制文件將以低開(kāi)銷的方式加載。

記憶模型

Warp 中的內(nèi)存分配通過(guò)warp.array類型公開(kāi)。陣列封裝了可能位于主機(jī)（ CPU ）或設(shè)備（ GPU ）內(nèi)存中的底層內(nèi)存分配。與張量框架不同， Warp 中的數(shù)組是強(qiáng)類型的，并存儲(chǔ)內(nèi)置結(jié)構(gòu)的線性序列（vec3， matrix33， quat，等等）。

您可以從 Python 列表或 NumPy 數(shù)組中構(gòu)造數(shù)組，或使用與 NumPy 和 PyTorch 類似的語(yǔ)法進(jìn)行初始化：

# allocate an uninitizalized array of vec3s v = wp.empty(length=n, dtype=wp.vec3, device="cuda") # allocate a zero-initialized array of quaternions q = wp.zeros(length=n, dtype=wp.quat, device="cuda") # allocate and initialize an array from a numpy array # will be automatically transferred to the specified device v = wp.from_numpy(array, dtype=wp.vec3, device="cuda")

Warp 支持__array_interface__和__cuda_array_interface__協(xié)議，允許在基于張量的框架之間進(jìn)行零拷貝數(shù)據(jù)視圖。例如，要將數(shù)據(jù)轉(zhuǎn)換為 NumPy ，請(qǐng)使用以下命令：

# automatically bring data from device back to host view = device_array.numpy()

特征

Warp 包含幾個(gè)更高級(jí)別的數(shù)據(jù)結(jié)構(gòu)，使實(shí)現(xiàn)模擬和幾何處理算法更容易。

網(wǎng)格

三角形網(wǎng)格在仿真和計(jì)算機(jī)圖形學(xué)中無(wú)處不在。 Warp 提供了一種內(nèi)置類型，用于管理網(wǎng)格數(shù)據(jù)，該數(shù)據(jù)支持幾何查詢，例如最近點(diǎn)、光線投射和重疊檢查。

下面的示例演示如何使用“扭曲”計(jì)算網(wǎng)格上距離輸入位置數(shù)組最近的點(diǎn)。這種類型的計(jì)算是碰撞檢測(cè)中許多算法的基礎(chǔ)（圖 3 ）。 Warp 的網(wǎng)格查詢使實(shí)現(xiàn)此類方法變得簡(jiǎn)單。

@wp.kernel
def project(positions: wp.array(dtype=wp.vec3), mesh: wp.uint64, output_pos: wp.array(dtype=wp.vec3), output_face: wp.array(dtype=int)): tid = wp.tid() x = wp.load(positions, tid) face_index = int(0) face_u = float(0.0) face_v = float(0.0) sign = float(0.0) max_dist = 2.0 if (wp.mesh_query_point(mesh, x, max_dist, sign, face_index, face_u, face_v)): p = wp.mesh_eval_position(mesh, face_index, face_u, face_v) output_pos[tid] = p output_face[tid] = face_index

稀疏卷

稀疏體對(duì)于表示大型域上的網(wǎng)格數(shù)據(jù)非常有用，例如復(fù)雜對(duì)象的符號(hào)距離場(chǎng)（ SDF ）或大規(guī)模流體流動(dòng)的速度。 Warp 支持使用 NanoVDB 標(biāo)準(zhǔn)定義的稀疏卷。使用標(biāo)準(zhǔn) OpenVDB 工具（如 Blender 、 Houdini 或 Maya ）構(gòu)造卷，然后在 Warp 內(nèi)核內(nèi)部采樣。

您可以直接從磁盤或內(nèi)存中的二進(jìn)制網(wǎng)格文件創(chuàng)建卷，然后使用volumes API 對(duì)其進(jìn)行采樣：

wp.volume_sample_world(vol, xyz, mode) # world space sample using interpolation mode
wp.volume_sample_local(vol, uvw, mode) # volume space sample using interpolation mode
wp.volume_lookup(vol, ijk) # direct voxel lookup
wp.volume_transform(vol, xyz) # map point from voxel space to world space
wp.volume_transform_inv(vol, xyz) # map point from world space to volume space

使用卷查詢，您可以以最小的內(nèi)存開(kāi)銷高效地碰撞復(fù)雜對(duì)象。

散列網(wǎng)格

許多基于粒子的模擬方法，如離散元法（ DEM ）或平滑粒子流體動(dòng)力學(xué)（ SPH ），都涉及到在空間鄰域上迭代以計(jì)算力的相互作用。哈希網(wǎng)格是一種成熟的數(shù)據(jù)結(jié)構(gòu)，用于加速這些最近鄰查詢，特別適合 GPU 。

哈希網(wǎng)格由點(diǎn)集構(gòu)成，如下所示：

哈希網(wǎng)格由點(diǎn)集構(gòu)成，如下所示：

grid = wp.HashGrid(dim_x=128, dim_y=128, dim_z=128, device="cuda")
grid.build(points=p, radius=r)

創(chuàng)建散列網(wǎng)格后，可以直接從用戶內(nèi)核代碼中查詢它們，如以下示例所示，該示例計(jì)算所有相鄰粒子位置的總和：

@wp.kernel
def sum(grid : wp.uint64, points: wp.array(dtype=wp.vec3), output: wp.array(dtype=wp.vec3), radius: float): tid = wp.tid() # query point p = points[tid] # create grid query around point query = wp.hash_grid_query(grid, p, radius) index = int(0) sum = wp.vec3() while(wp.hash_grid_query_next(query, index)): neighbor = points[index] # compute distance to neighbor point dist = wp.length(p-neighbor) if (dist <= radius): sum += neighbor output[tid] = sum

圖 5 顯示了粘性材料的 DEM 顆粒材料模擬示例。使用內(nèi)置的哈希網(wǎng)格數(shù)據(jù)結(jié)構(gòu)，您可以在不到 200 行 Python 中編寫(xiě)這樣的模擬，并以交互速率運(yùn)行超過(guò) 100K 個(gè)粒子。

使用扭曲散列網(wǎng)格數(shù)據(jù)可以輕松評(píng)估相鄰粒子之間的成對(duì)力相互作用。

可微性

基于張量的框架，如 PyTorch 和 JAX ，提供了張量計(jì)算的梯度，非常適合于 ML 訓(xùn)練等應(yīng)用。

Warp 的一個(gè)獨(dú)特功能是能夠生成 kernel code 的正向和反向版本。這使得編寫(xiě)可微模擬變得很容易，可以將梯度作為更大訓(xùn)練管道的一部分進(jìn)行傳播。一個(gè)常見(jiàn)的場(chǎng)景是，對(duì)網(wǎng)絡(luò)層使用傳統(tǒng)的 ML 框架，并使用 Warp 實(shí)現(xiàn)允許端到端差異性的模擬層。

當(dāng)需要漸變時(shí)，應(yīng)使用requires_grad=True創(chuàng)建陣列。例如，warp.Tape類可以記錄內(nèi)核啟動(dòng)并回放它們，以計(jì)算標(biāo)量損失函數(shù)相對(duì)于內(nèi)核輸入的梯度：

tape = wp.Tape() # forward pass
with tape: wp.launch(kernel=compute1, inputs=[a, b], device="cuda") wp.launch(kernel=compute2, inputs=[c, d], device="cuda") wp.launch(kernel=loss, inputs=[d, l], device="cuda") # reverse pass
tape.backward(loss=l)

完成后向傳遞后，可通過(guò)Tape對(duì)象中的映射獲得與輸入相關(guān)的梯度：

# gradient of loss with respect to input a
print(tape.gradients[a])

圖 6 。一個(gè)軌跡優(yōu)化的例子，其中球的初始速度被優(yōu)化以擊中黑色目標(biāo)。每行顯示 LBFGS 優(yōu)化步驟的一次迭代的結(jié)果。

總結(jié)

在這篇文章中，我們介紹了 NVIDIA Warp ，這是一個(gè) Python 框架，可以很容易地為 GPU 編寫(xiě)可微模擬代碼。

關(guān)于作者

邁爾斯·麥克林（ Miles Macklin ）是NVIDIA 的首席工程師，致力于模擬技術(shù)。他從哥本哈根大學(xué)獲得計(jì)算機(jī)科學(xué)博士學(xué)位，從事計(jì)算機(jī)圖形學(xué)、基于物理學(xué)的動(dòng)畫(huà)和機(jī)器人學(xué)的研究。他在 ACM SIGGRAPH 期刊上發(fā)表了幾篇論文，他的研究已經(jīng)被整合到許多商業(yè)產(chǎn)品中，包括NVIDIA 的 PhysX 和 ISAAC 健身房模擬器。他最近的工作旨在為 GPU 上的可微編程開(kāi)發(fā)健壯高效的框架。

Fred Oh 是 CUDA 、 CUDA on WSL 和 CUDA Python 的高級(jí)產(chǎn)品營(yíng)銷經(jīng)理。弗雷德?lián)碛屑又荽髮W(xué)戴維斯分校計(jì)算機(jī)科學(xué)和數(shù)學(xué)學(xué)士學(xué)位。他的職業(yè)生涯開(kāi)始于一名 UNIX 軟件工程師，負(fù)責(zé)將內(nèi)核服務(wù)和設(shè)備驅(qū)動(dòng)程序移植到 x86 體系結(jié)構(gòu)。他喜歡《星球大戰(zhàn)》、《星際迷航》和 NBA 勇士隊(duì)。

審核編輯：郭婷