使用Nuba扩展在Python中编写光线跟踪应用程序

光线跟踪是一种渲染算法，它可以通过模拟光如何传输以及与不同材质的交互来生成照片级真实感图像。如今，它被广泛应用于游戏开发、电影制作和物理模拟中，将图像带入生活。

然而，光线跟踪算法计算量大，需要在 GPU 上进行硬件加速才能实现实时性能。

为了利用光线跟踪的硬件功能，人们发明了各种工具链和语言来满足需要，例如 openGL 和着色语言。

通常，这些软件工具链的构建过程会给 Python 开发人员带来重大挑战。为了减轻困难并为编写光线跟踪内核提供熟悉的环境， NVIDIA 为 PyOptiX 开发了 Numba 扩展。这种扩展使图形研究人员和应用程序开发人员能够减少从构思到实现的时间，并缩短每次迭代的开发周期。

在本文中，我将概述 NVIDIA 光线跟踪引擎 PyOptiX ，并解释 Python JIT 编译器 Numba 如何加速 Python 代码。最后，通过一个完整的光线跟踪示例，我将引导您完成使用 PyOptiX 的 Nuba 扩展的步骤，并用 Python 编写一个加速的光线跟踪内核。

什么是 NVIDIA OptiX 和 PyOptiX ？

NVIDIA RTX 技术使光线跟踪成为许多现代渲染管道中的默认渲染算法。由于对独特外观的需求是无限的，因此需要灵活定制渲染管道。

NVIDIA RTX 光线跟踪管道是可定制的。通过配置光在各种材质上的传输、反射和折射方式，可以在对象上实现独特的外观，例如有光泽、有光泽或半透明。通过配置光线的生成方式，可以相应地更改视图的视野和透视效果。

为了满足这一需求， NVIDIA 开发了 NVIDIA OptiX ，这是一种光线跟踪引擎，可用于配置硬件加速的光线跟踪管道。 PyOptiX 是 NVIDIA OptiX Python 接口。此接口为 Python 开发人员提供了与使用 C ++编写的 NVIDIA OptiX 开发人员相同的功能。

内核函数

要自定义图像方面，可以使用内核函数，也称为内核方法或内核。您可以将内核视为一组将数据输入转换为所需形式的算法。本地 NVIDIA OptiX 开发人员可以使用 CUDA 编写内核。使用 Nuba 扩展，您可以在 Python 中编写光线跟踪内核。

Numba 和 Numba 的性能更高。库达大学

光线跟踪是一种计算密集型算法。虽然理论上可以使用标准 C Python 解释器运行光线跟踪内核，但渲染常规光线跟踪图像需要几天的时间。此外， NVIDIA OptiX 要求内核可以在 GPU 设备上运行，以便与其余渲染管道集成。

使用 Numba ，一个实时的 GPU 函数编译器，您可以使用 Python 硬件执行并加速您的 Python 光线跟踪内核。 Numba 解析 Python 功能代码并将其转换为有效的机器代码。在较高层次上，该过程分为七个步骤：

该函数的字节码由字节码编译器生成。

分析了字节码。生成控制流图（ CFG ）和数据流图（ DFG ）。

通过字节码、 CFG 和 DFG ，可以生成 Numba 中间表示（ IR ）。

根据函数输入的类型，推断每个 IR 变量的类型。

Nuba IR 被重写，并得到 Python 特定的优化。

Numba IR 降低到 LLVM IR ，并执行更一般的优化。

LLVM IR 由 LLVM 后端使用，并生成优化的 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 函数转换为本机代码，所以在 Numba CUDA 内核中， Python 用户拥有同等的权限，就像他们在用本机 CUDA 编写内核一样。此代码显示可在设备上执行的点产品。有关更多信息，请参阅 Numba Examples 。

介绍 PyOptiX 的 Nuba 扩展

要自定义光线跟踪管道的特定阶段，必须将 Nuba 内核转换为 NVIDIA OptiX 引擎可以理解的内容。 NVIDIA 为 PyOptiX 开发了 Numba 扩展以实现这一目标。

扩展包括自定义类型定义和内部函数降维。 NVIDIA OptiX 附带一组内部类型：

OptixTraversableHandle

OptixVisibilityMask

SbtDataPointer

功能，如optix.Trace

为了让 Nuba 对这些新类型和方法执行类型推断，您必须注册这些类型并在编译用户内核之前提供这些方法的实现。目前， NVIDIA 正在扩展支持的类型和内部函数，以添加更多示例。

通过向 Numba 公开这些类型和内部函数，您现在可以编写内核，它不仅针对 GPU ，而且可以专门针对 GPU 进行光线跟踪内核。与 Numba CUDA 结合使用，您可以编写功率相等的光线跟踪内核，就像为 NVIDIA OptiX 编写本机 CUDA 光线跟踪内核一样。

在下一节中，我将介绍一个带有 PyOptiX-Numba 扩展的 Hello-World 示例。在此之前，让我快速回顾一些光线跟踪算法的基础知识。

射线追踪基础

假设您使用相机拍摄图像。场景中的光源发射光线，光线沿直线传播。当光线击中物体时，它会从表面反射，最终到达相机传感器。

从较高的层次来看，光线跟踪算法将遍历到达图像平面的所有光线，以在场景中确定光线的相交位置和相交内容。找到交点后，可以采用各种着色技术来确定交点的颜色。然而，也有一些射线不会击中场景中的任何东西。在这种情况下，这些光线被视为“丢失”目标。

使用 PyOptiX 的 Numba 扩展对三角形进行光线跟踪的步骤

在下面的示例中，我将展示 PyOptiX 的 Numba 扩展如何帮助您编写自定义内核，以定义光线生成、光线命中和光线未命中时的光线行为。

场景设置

我将您看到的视图建模为一个图像平面，它通常略位于相机前面。相机被建模为三维空间中的一个点和一组相互正交的向量。

图 2 三角形渲染示例的场景设置

照相机

相机建模为三维中的一个点。摄像机的三个矢量， U 、 V 和 W 、 用于显示侧面、向上和正面方向 。这唯一地确定了相机的位置和方向。

为了简化后续光线生成的计算， U 和 V 矢量不是单位矢量。相反，它们的长度与图像的纵横比成比例匹配。最后， W 向量的长度是相机和图像平面之间的距离。

射线生成内核

射线生成内核是该算法的核心。射线原点和方向在此处生成，然后传递给跟踪调用。它的强度从其他内核中检索出来，并作为图像数据写入。在本节中，我将讨论在此内核中生成光线的方法。

使用相机和图像平面，可以生成光线。采用以图像中心为原点的坐标系约定。图像像素中坐标的符号表示其相对于原点的相对位置，其大小表示距离。使用此属性，将相机的 U 和 V 矢量与像素位置的相应元素相乘，然后将它们相加。结果是从图像中心指向像素的向量。

最后，将该向量添加到 W 或前向量，这将生成一条光线，该光线从相机位置开始，穿过图像平面上的像素。图 3 显示了一条光线的分解，该光线起源于相机，并穿过图像平面中的点（ x 、 y ）。

图 3 穿过像素的光线分解 （ x ， y ）

在代码中，可以使用 optix 的两个内在函数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的助手函数，该函数计算光线的原点和方向向量。

接下来，将生成的光线传递给内部函数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 内部数据结构设置的，则可以调用 NVIDIA OptiX 内在optix.GetTriangleBarycentrics来检索命中点的重心坐标。

要使颜色更有趣，请将此坐标插入该像素的颜色中。颜色的蓝色通道设置为 1.0 。光线的强度应传递给光线生成内核进行进一步的后处理，并写入图像。

NVIDIA OptiX 通过有效负载寄存器在内核之间共享数据。使用setPayload功能将有效负载寄存器的值设置为光线强度。默认情况下，有效负载寄存器是整数类型。使用 CUDA 内部函数float_as_int将浮点值解释为整数，而不更改位。

@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是在设置渲染管道期间在着色器绑定表中指定的一些数据。现在，请注意，这是一组硬编码的浮点数，表示场景的背景色。

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

将强度转换为颜色并写入图像

现在，您已经为所有光线定义了颜色。颜色在光线生成内核中作为payload_pack数据结构从optix.trace调用中检索。还记得在 ray hit 和 ray miss 内核中，必须将浮点数的位解释为整数吗？使用int_as_float功能还原此步骤。

现在，您可以直接将这些值写入图像，它仍然看起来很棒。再多做一步，对原始像素值执行后处理步骤，这对于更复杂场景中的出色图像非常重要。

您检索到的值只是光线的原始强度，它与光线携带的能量级别成线性比例。虽然这符合你的物理世界模型，但人眼不会以线性方式对光刺激作出反应。相反，它遵循输入的映射，通过幂函数进行响应。

为此，对强度进行 gamma correction 测试。此外，大多数查看此图像结果的用户都在观看具有 sRGB 颜色空间的监视器。假设光线跟踪世界中的值位于 CIE-XYZ color space 中，并应用颜色空间转换。最后，将颜色值量化为 8 位无符号整数。

下面的代码示例显示了用于后期处理颜色强度并将其写入光线生成内核中的像素阵列的辅助函数。

@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 函数转换为与渲染管道兼容的设备代码。 NVIDIA 将这两个库组合到 PyOptiX 的 Nuba 扩展中，使您能够在完整的 Python 环境中编写加速光线跟踪应用程序。

结合 Python 已经拥有的丰富而活跃的环境，您现在可以解锁构建光线跟踪应用程序的真正能力，硬件加速。 下载演示 亲自体验 PyOptiX 的 Numba 扩展！

下一步是什么？

PyOptiX Numba 扩展正处于开发阶段， NVIDIA 正在努力添加更多示例，并使 NVIDIA OptiX 原语的键入更加灵活和 Pythonic 。

关于作者

Michael Yh Wang 是 NVIDIA Rapids 的软件工程师。目前，他将自己的工程技能贡献给了 cuDF 、 cuSpatial 和 Numba 。在加入 NVIDIA 之前，他获得了耶鲁大学的理学硕士学位。他早期的经验包括在一个独立电影项目中担任视觉效果主管，并在 WAIC 2020 hackathon 竞赛中获得第一名。 Michael 对软件工程、计算机图形算法和编译器技术有浓厚的兴趣。他相信，在未来，通过编译器和语言创新，加速计算将更容易为公众所接受。

审核编辑：郭婷