使用Nuba擴(kuò)展在Python中編寫光線跟蹤應(yīng)用程序

光線跟蹤是一種渲染算法，它可以通過模擬光如何傳輸以及與不同材質(zhì)的交互來生成照片級真實感圖像。如今，它被廣泛應(yīng)用于游戲開發(fā)、電影制作和物理模擬中，將圖像帶入生活。

然而，光線跟蹤算法計算量大，需要在 GPU 上進(jìn)行硬件加速才能實現(xiàn)實時性能。

為了利用光線跟蹤的硬件功能，人們發(fā)明了各種工具鏈和語言來滿足需要，例如 openGL 和著色語言。

通常，這些軟件工具鏈的構(gòu)建過程會給 Python 開發(fā)人員帶來重大挑戰(zhàn)。為了減輕困難并為編寫光線跟蹤內(nèi)核提供熟悉的環(huán)境， NVIDIA 為 PyOptiX 開發(fā)了 Numba 擴(kuò)展。這種擴(kuò)展使圖形研究人員和應(yīng)用程序開發(fā)人員能夠減少從構(gòu)思到實現(xiàn)的時間，并縮短每次迭代的開發(fā)周期。

在本文中，我將概述 NVIDIA 光線跟蹤引擎 PyOptiX ，并解釋 Python JIT 編譯器 Numba 如何加速 Python 代碼。最后，通過一個完整的光線跟蹤示例，我將引導(dǎo)您完成使用 PyOptiX 的 Nuba 擴(kuò)展的步驟，并用 Python 編寫一個加速的光線跟蹤內(nèi)核。

什么是 NVIDIA OptiX 和 PyOptiX ？

NVIDIA RTX 技術(shù)使光線跟蹤成為許多現(xiàn)代渲染管道中的默認(rèn)渲染算法。由于對獨特外觀的需求是無限的，因此需要靈活定制渲染管道。

NVIDIA RTX 光線跟蹤管道是可定制的。通過配置光在各種材質(zhì)上的傳輸、反射和折射方式，可以在對象上實現(xiàn)獨特的外觀，例如有光澤、有光澤或半透明。通過配置光線的生成方式，可以相應(yīng)地更改視圖的視野和透視效果。

為了滿足這一需求， NVIDIA 開發(fā)了 NVIDIA OptiX ，這是一種光線跟蹤引擎，可用于配置硬件加速的光線跟蹤管道。 PyOptiX 是 NVIDIA OptiX Python 接口。此接口為 Python 開發(fā)人員提供了與使用 C ++編寫的 NVIDIA OptiX 開發(fā)人員相同的功能。

內(nèi)核函數(shù)

要自定義圖像方面，可以使用內(nèi)核函數(shù)，也稱為內(nèi)核方法或內(nèi)核。您可以將內(nèi)核視為一組將數(shù)據(jù)輸入轉(zhuǎn)換為所需形式的算法。本地 NVIDIA OptiX 開發(fā)人員可以使用 CUDA 編寫內(nèi)核。使用 Nuba 擴(kuò)展，您可以在 Python 中編寫光線跟蹤內(nèi)核。

Numba 和 Numba 的性能更高。庫達(dá)大學(xué)

光線跟蹤是一種計算密集型算法。雖然理論上可以使用標(biāo)準(zhǔn) C Python 解釋器運行光線跟蹤內(nèi)核，但渲染常規(guī)光線跟蹤圖像需要幾天的時間。此外， NVIDIA OptiX 要求內(nèi)核可以在 GPU 設(shè)備上運行，以便與其余渲染管道集成。

使用 Numba ，一個實時的 GPU 函數(shù)編譯器，您可以使用 Python 硬件執(zhí)行并加速您的 Python 光線跟蹤內(nèi)核。 Numba 解析 Python 功能代碼并將其轉(zhuǎn)換為有效的機(jī)器代碼。在較高層次上，該過程分為七個步驟：

該函數(shù)的字節(jié)碼由字節(jié)碼編譯器生成。

分析了字節(jié)碼。生成控制流圖（ CFG ）和數(shù)據(jù)流圖（ DFG ）。

通過字節(jié)碼、 CFG 和 DFG ，可以生成 Numba 中間表示（ IR ）。

根據(jù)函數(shù)輸入的類型，推斷每個 IR 變量的類型。

Nuba IR 被重寫，并得到 Python 特定的優(yōu)化。

Numba IR 降低到 LLVM IR ，并執(zhí)行更一般的優(yōu)化。

LLVM IR 由 LLVM 后端使用，并生成優(yōu)化的 GPU 機(jī)器代碼。

圖 1 Numba 編譯管道的高級視圖

圖 1 顯示了前面提到的編譯管道的圖形概述。這篇關(guān)于 Numba 編譯器管道的快速教程只提供了對 Numba 內(nèi)部架構(gòu)的一點了解。

下面的代碼顯示了一個示例 GPU 內(nèi)核，該內(nèi)核計算兩個 3 元素向量的點積。

@cuda.jit(device=True)
def dot(a, b): return a.x * b.x + a.y * b.y + a.z * b.z

因為 Numba 可以將任何 Python 函數(shù)轉(zhuǎn)換為本機(jī)代碼，所以在 Numba CUDA 內(nèi)核中， Python 用戶擁有同等的權(quán)限，就像他們在用本機(jī) CUDA 編寫內(nèi)核一樣。此代碼顯示可在設(shè)備上執(zhí)行的點產(chǎn)品。有關(guān)更多信息，請參閱 Numba Examples 。

介紹 PyOptiX 的 Nuba 擴(kuò)展

要自定義光線跟蹤管道的特定階段，必須將 Nuba 內(nèi)核轉(zhuǎn)換為 NVIDIA OptiX 引擎可以理解的內(nèi)容。 NVIDIA 為 PyOptiX 開發(fā)了 Numba 擴(kuò)展以實現(xiàn)這一目標(biāo)。

擴(kuò)展包括自定義類型定義和內(nèi)部函數(shù)降維。 NVIDIA OptiX 附帶一組內(nèi)部類型：

OptixTraversableHandle

OptixVisibilityMask

SbtDataPointer

功能，如optix.Trace

為了讓 Nuba 對這些新類型和方法執(zhí)行類型推斷，您必須注冊這些類型并在編譯用戶內(nèi)核之前提供這些方法的實現(xiàn)。目前， NVIDIA 正在擴(kuò)展支持的類型和內(nèi)部函數(shù)，以添加更多示例。

通過向 Numba 公開這些類型和內(nèi)部函數(shù)，您現(xiàn)在可以編寫內(nèi)核，它不僅針對 GPU ，而且可以專門針對 GPU 進(jìn)行光線跟蹤內(nèi)核。與 Numba CUDA 結(jié)合使用，您可以編寫功率相等的光線跟蹤內(nèi)核，就像為 NVIDIA OptiX 編寫本機(jī) CUDA 光線跟蹤內(nèi)核一樣。

在下一節(jié)中，我將介紹一個帶有 PyOptiX-Numba 擴(kuò)展的 Hello-World 示例。在此之前，讓我快速回顧一些光線跟蹤算法的基礎(chǔ)知識。

射線追蹤基礎(chǔ)

假設(shè)您使用相機(jī)拍攝圖像。場景中的光源發(fā)射光線，光線沿直線傳播。當(dāng)光線擊中物體時，它會從表面反射，最終到達(dá)相機(jī)傳感器。

從較高的層次來看，光線跟蹤算法將遍歷到達(dá)圖像平面的所有光線，以在場景中確定光線的相交位置和相交內(nèi)容。找到交點后，可以采用各種著色技術(shù)來確定交點的顏色。然而，也有一些射線不會擊中場景中的任何東西。在這種情況下，這些光線被視為“丟失”目標(biāo)。

使用 PyOptiX 的 Numba 擴(kuò)展對三角形進(jìn)行光線跟蹤的步驟

在下面的示例中，我將展示 PyOptiX 的 Numba 擴(kuò)展如何幫助您編寫自定義內(nèi)核，以定義光線生成、光線命中和光線未命中時的光線行為。

場景設(shè)置

我將您看到的視圖建模為一個圖像平面，它通常略位于相機(jī)前面。相機(jī)被建模為三維空間中的一個點和一組相互正交的向量。

圖 2 三角形渲染示例的場景設(shè)置

照相機(jī)

相機(jī)建模為三維中的一個點。攝像機(jī)的三個矢量， U 、 V 和 W 、 用于顯示側(cè)面、向上和正面方向 。這唯一地確定了相機(jī)的位置和方向。

為了簡化后續(xù)光線生成的計算， U 和 V 矢量不是單位矢量。相反，它們的長度與圖像的縱橫比成比例匹配。最后， W 向量的長度是相機(jī)和圖像平面之間的距離。

射線生成內(nèi)核

射線生成內(nèi)核是該算法的核心。射線原點和方向在此處生成，然后傳遞給跟蹤調(diào)用。它的強(qiáng)度從其他內(nèi)核中檢索出來，并作為圖像數(shù)據(jù)寫入。在本節(jié)中，我將討論在此內(nèi)核中生成光線的方法。

使用相機(jī)和圖像平面，可以生成光線。采用以圖像中心為原點的坐標(biāo)系約定。圖像像素中坐標(biāo)的符號表示其相對于原點的相對位置，其大小表示距離。使用此屬性，將相機(jī)的 U 和 V 矢量與像素位置的相應(yīng)元素相乘，然后將它們相加。結(jié)果是從圖像中心指向像素的向量。

最后，將該向量添加到 W 或前向量，這將生成一條光線，該光線從相機(jī)位置開始，穿過圖像平面上的像素。圖 3 顯示了一條光線的分解，該光線起源于相機(jī)，并穿過圖像平面中的點（ x 、 y ）。

圖 3 穿過像素的光線分解 （ x ， y ）

在代碼中，可以使用 optix 的兩個內(nèi)在函數(shù)optix.GetLaunchIndex和optix.GetLaunchDimensions檢索圖像平面的像素索引和圖像尺寸。接下來，像素索引被歸一化為［-1.0 ， 1.0］。下面的代碼示例顯示了 Nuba CUDA 內(nèi)核中的這種邏輯。

@cuda.jit(device=True, fast_math=True)
def computeRay(idx, dim): U = params.cam_u V = params.cam_v W = params.cam_w # Normalizing coordinates to [-1.0, 1.0] d = float32(2.0) * make_float2( float32(idx.x) / float32(dim.x), float32(idx.y) / float32(dim.y) ) - float32(1.0) origin = params.cam_eye direction = normalize(d.x * U + d.y * V + W) return origin, direction def __raygen__rg():
 # Look up your location within the launch grid
 idx = optix.GetLaunchIndex() dim = optix.GetLaunchDimensions()  # Map your launch idx to a screen location and create a ray from the camera # location through the screen ray_origin, ray_direction = computeRay(make_uint3(idx.x, idx.y, 0), dim)

此代碼示例顯示了computeRay的助手函數(shù)，該函數(shù)計算光線的原點和方向向量。

接下來，將生成的光線傳遞給內(nèi)部函數(shù)optix.Trace。這將初始化光線跟蹤算法。底層 optiX 引擎遍歷基本體，計算場景中的交點，最后返回光線的強(qiáng)度。下面的代碼示例顯示了對optix.Trace的調(diào)用。

# In __raygen__rg
 payload_pack = optix.Trace( params.handle, ray_origin, ray_direction, float32(0.0), # Min intersection distance float32(1e16), # Max intersection distance float32(0.0), # rayTime -- used for motion blur OptixVisibilityMask(255),  # Specify always visible uint32(OPTIX_RAY_FLAG_NONE), uint32(0), # SBT offset -- Refer to OptiX Manual for SBT uint32(1),  # SBT stride -- Refer to OptiX Manual for SBT uint32(0),  # missSBTIndex -- Refer to OptiX Manual for SBT )

射線命中內(nèi)核

在光線命中內(nèi)核中，您可以編寫代碼來確定光線的每個通道的強(qiáng)度。如果三角形頂點是使用 NVIDIA OptiX 內(nèi)部數(shù)據(jù)結(jié)構(gòu)設(shè)置的，則可以調(diào)用 NVIDIA OptiX 內(nèi)在optix.GetTriangleBarycentrics來檢索命中點的重心坐標(biāo)。

要使顏色更有趣，請將此坐標(biāo)插入該像素的顏色中。顏色的藍(lán)色通道設(shè)置為 1.0 。光線的強(qiáng)度應(yīng)傳遞給光線生成內(nèi)核進(jìn)行進(jìn)一步的后處理，并寫入圖像。

NVIDIA OptiX 通過有效負(fù)載寄存器在內(nèi)核之間共享數(shù)據(jù)。使用setPayload功能將有效負(fù)載寄存器的值設(shè)置為光線強(qiáng)度。默認(rèn)情況下，有效負(fù)載寄存器是整數(shù)類型。使用 CUDA 內(nèi)部函數(shù)float_as_int將浮點值解釋為整數(shù)，而不更改位。

@cuda.jit(device=True, fast_math=True)
def setPayload(p): optix.SetPayload_0(float_as_int(p.x)) optix.SetPayload_1(float_as_int(p.y)) optix.SetPayload_2(float_as_int(p.z)) def __closesthit__ch():  # When a built-in triangle intersection is used, a number of fundamental # attributes are provided by the NVIDIA OptiX API, including barycentric coordinates. barycentrics = optix.GetTriangleBarycentrics() setPayload(make_float3(barycentrics, float32(1.0)))

射線未命中內(nèi)核

“光線未命中”內(nèi)核設(shè)置未命中場景中任何對象的光線的顏色。在這里，您可以將它們設(shè)置為背景色。

bg_color是在設(shè)置渲染管道期間在著色器綁定表中指定的一些數(shù)據(jù)。現(xiàn)在，請注意，這是一組硬編碼的浮點數(shù)，表示場景的背景色。

def __miss__ms(): miss_data = MissDataStruct(optix.GetSbtDataPointer()) setPayload(miss_data.bg_color)

將強(qiáng)度轉(zhuǎn)換為顏色并寫入圖像

現(xiàn)在，您已經(jīng)為所有光線定義了顏色。顏色在光線生成內(nèi)核中作為payload_pack數(shù)據(jù)結(jié)構(gòu)從optix.trace調(diào)用中檢索。還記得在 ray hit 和 ray miss 內(nèi)核中，必須將浮點數(shù)的位解釋為整數(shù)嗎？使用int_as_float功能還原此步驟。

現(xiàn)在，您可以直接將這些值寫入圖像，它仍然看起來很棒。再多做一步，對原始像素值執(zhí)行后處理步驟，這對于更復(fù)雜場景中的出色圖像非常重要。

您檢索到的值只是光線的原始強(qiáng)度，它與光線攜帶的能量級別成線性比例。雖然這符合你的物理世界模型，但人眼不會以線性方式對光刺激作出反應(yīng)。相反，它遵循輸入的映射，通過冪函數(shù)進(jìn)行響應(yīng)。

為此，對強(qiáng)度進(jìn)行 gamma correction 測試。此外，大多數(shù)查看此圖像結(jié)果的用戶都在觀看具有 sRGB 顏色空間的監(jiān)視器。假設(shè)光線跟蹤世界中的值位于 CIE-XYZ color space 中，并應(yīng)用顏色空間轉(zhuǎn)換。最后，將顏色值量化為 8 位無符號整數(shù)。

下面的代碼示例顯示了用于后期處理顏色強(qiáng)度并將其寫入光線生成內(nèi)核中的像素陣列的輔助函數(shù)。

@cuda.jit(device=True, fast_math=True)
def toSRGB(c):
 # Use float32 for constants
 invGamma = float32(1.0) / float32(2.4) powed = make_float3( fast_powf(c.x, invGamma), fast_powf(c.y, invGamma), fast_powf(c.z, invGamma), ) return make_float3( float32(12.92) * c.x if c.x < float32(0.0031308) else float32(1.055) * powed.x - float32(0.055), float32(12.92) * c.y if c.y < float32(0.0031308) else float32(1.055) * powed.y - float32(0.055), float32(12.92) * c.z if c.z < float32(0.0031308) else float32(1.055) * powed.z - float32(0.055), ) @cuda.jit(device=True, fast_math=True)
def make_color(c): srgb = toSRGB(clamp(c, float32(0.0), float32(1.0))) return make_uchar4( quantizeUnsigned8Bits(srgb.x), quantizeUnsigned8Bits(srgb.y), quantizeUnsigned8Bits(srgb.z), uint8(255), ) # In __raygen__rg result = make_float3( int_as_float(payload_pack.p0), int_as_float(payload_pack.p1), int_as_float(payload_pack.p2), )  # Record results in your output raster params.image[idx.y * params.image_width + idx.x] = make_color(result)

總結(jié)

PyOptiX 允許您使用 Python 設(shè)置光線跟蹤渲染管道。 Nuba 將 Python 函數(shù)轉(zhuǎn)換為與渲染管道兼容的設(shè)備代碼。 NVIDIA 將這兩個庫組合到 PyOptiX 的 Nuba 擴(kuò)展中，使您能夠在完整的 Python 環(huán)境中編寫加速光線跟蹤應(yīng)用程序。

結(jié)合 Python 已經(jīng)擁有的豐富而活躍的環(huán)境，您現(xiàn)在可以解鎖構(gòu)建光線跟蹤應(yīng)用程序的真正能力，硬件加速。 下載演示 親自體驗 PyOptiX 的 Numba 擴(kuò)展！

下一步是什么？

PyOptiX Numba 擴(kuò)展正處于開發(fā)階段， NVIDIA 正在努力添加更多示例，并使 NVIDIA OptiX 原語的鍵入更加靈活和 Pythonic 。

關(guān)于作者

Michael Yh Wang 是 NVIDIA Rapids 的軟件工程師。目前，他將自己的工程技能貢獻(xiàn)給了 cuDF 、 cuSpatial 和 Numba 。在加入 NVIDIA 之前，他獲得了耶魯大學(xué)的理學(xué)碩士學(xué)位。他早期的經(jīng)驗包括在一個獨立電影項目中擔(dān)任視覺效果主管，并在 WAIC 2020 hackathon 競賽中獲得第一名。 Michael 對軟件工程、計算機(jī)圖形算法和編譯器技術(shù)有濃厚的興趣。他相信，在未來，通過編譯器和語言創(chuàng)新，加速計算將更容易為公眾所接受。

審核編輯：郭婷