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

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

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

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

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

什么是 NVIDIA OptiX 和 PyOptiX ？

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

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

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

內核函數(shù)

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

Numba 和 Numba 的性能更高。庫達大學

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

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

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

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

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

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

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

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

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

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

圖 1 顯示了前面提到的編譯管道的圖形概述。這篇關于 Numba 編譯器管道的快速教程只提供了對 Numba 內部架構的一點了解。

下面的代碼顯示了一個示例 GPU 內核，該內核計算兩個 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ù)轉換為本機代碼，所以在 Numba CUDA 內核中， Python 用戶擁有同等的權限，就像他們在用本機 CUDA 編寫內核一樣。此代碼顯示可在設備上執(zhí)行的點產品。有關更多信息，請參閱 Numba Examples 。

介紹 PyOptiX 的 Nuba 擴展

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

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

OptixTraversableHandle

OptixVisibilityMask

SbtDataPointer

功能，如optix.Trace

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

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

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

射線追蹤基礎

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

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

使用 PyOptiX 的 Numba 擴展對三角形進行光線跟蹤的步驟

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

場景設置

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

圖 2 三角形渲染示例的場景設置

照相機

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

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

射線生成內核

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

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

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

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

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

@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ù)計算光線的原點和方向向量。

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

# 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 )

射線命中內核

在光線命中內核中，您可以編寫代碼來確定光線的每個通道的強度。如果三角形頂點是使用 NVIDIA OptiX 內部數(shù)據結構設置的，則可以調用 NVIDIA OptiX 內在optix.GetTriangleBarycentrics來檢索命中點的重心坐標。

要使顏色更有趣，請將此坐標插入該像素的顏色中。顏色的藍色通道設置為 1.0 。光線的強度應傳遞給光線生成內核進行進一步的后處理，并寫入圖像。

NVIDIA OptiX 通過有效負載寄存器在內核之間共享數(shù)據。使用setPayload功能將有效負載寄存器的值設置為光線強度。默認情況下，有效負載寄存器是整數(shù)類型。使用 CUDA 內部函數(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)))

射線未命中內核

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

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

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

將強度轉換為顏色并寫入圖像

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

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

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

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

下面的代碼示例顯示了用于后期處理顏色強度并將其寫入光線生成內核中的像素陣列的輔助函數(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)

總結

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

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

下一步是什么？

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

關于作者

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

審核編輯：郭婷