NVIDIA OptiX 9.0 的发布引入了一项名为 Cooperative Vectors 的新功能，可将 AI 工作流作为光线追踪内核的一部分。该功能利用 NVIDIA RTX Tensor Cores 在着色过程中进行硬件加速的矩阵运算和神经网络计算。这解锁了 NVIDIA RTX Neural Shaders 和 NVIDIA RTX Neural Texture Compression (NTC) 等 AI 渲染技术，并在实时渲染中进一步向电影级逼真材质迈进。
协作向量 API 已在 OptiX 、 DirectX 、 NVAPI 、 Slang 和 Vulkan 中推出。本文将探讨适用于所有 API 的协作向量背后的概念，并通过使用 OptiX API 的示例进行工作。
为何选择矩阵运算？
多层感知器 (MLP) 是许多神经网络算法的基本构建模块。研究表明，MLP 能够忠实地再现其训练所用的效果。即使 MLP 足够小，可以实时运行，它们也能够处理有趣的效果，例如基于物理性质的着色，有时比传统着色网络更快，内存占用更小。
MLP 通常由一个输入向量、几个全连接层和一个输出向量组成。不同层向量的大小不必相同。
MLP 评估 (推理) 的每一层都有两个阶段：前一层值的加权和偏差线性组合，以及可选的 非线性激活函数 。加权和偏差线性组合归结为矩阵向量乘法，然后添加偏差向量，也称为 Affine Transform 。
任何两个 affine 变换的合成都是 affine 变换。这意味着，如果每层只有 affine 线性相加偏置相位，则整个 MLP 始终可以简化为单个 affine 变换。由于在每层的 affine 结果后应用了非线性激活函数，MLP 的表现力远比这更强。
在全连接 MLP 中，每个神经元都是前一层中所有神经元的函数，如 Figure 1 所示。单层的完整概念计算如 Figure 2 所示。
为何选择协作向量？
“协作向量的目标之一是支持使用 NVIDIA Tensor Cores 来加速矩阵运算。通常情况下，CUDA SIMT 编程模型需要完整的活动线程束来执行此操作，但光线追踪编程模型独立处理线程，无法保证完整的线程束。此外，Tensor Cores 提供矩阵 – 矩阵乘法，但每个光线追踪线程只需要向量 – 矩阵乘法，这将未充分利用 Tensor Cores。”
此外，CUDA API 需要针对特定硬件版本，并且不能保证在各个架构之间向前兼容。如需了解矩阵乘法的 CUDA 多线程方法，请查看《 Matrix Multiplication Background User’s Guide 》。
协作向量通过提供以下 API 来解决这些限制：
- 允许使用包含一些不活动线程的 warp 进行矩阵运算
- 提供跨架构的向前和向后兼容性
- 支持用户在单个线程中指定向量数据，同时重新映射运算，以更高效地利用 Tensor Cores
协作向量可以处理线程束中的数据和执行离散，但性能会有所下降。当线程束中的 MLP 权重相同，并且线程束中有完整的线程互补时，即可获得最佳性能。使用 Shader Execution Reordering (SER) 可以帮助实现这两个目标。
由于评估 MLP 是一系列向量-matrix 乘法，因此当线程束中的所有线程并排评估同一 MLP 时，协作 vector API 可以将组合线程束的仿射运算视为 matrix-matrix 乘法和 bias。这就是协作意味着：线程捆绑在一起，将多个向量-matrix 运算转化为 matrix-matrix 运算。
outputMatrix = inputMatrix × weightsMatrix + biasMatrix
此处，除权重矩阵外，所有矩阵均为 32 行高，输入、输出和 bias 矩阵的每一行均表示单独线程的数据。
在 OptiX 中使用 Cooperative Vectors
协作向量是不透明向量类型，本质上是数组类，可以具有任意长度。OptiX 提供了一种名为 OptixCoopVec 的协作向量的实现。OptiX 中的协作向量支持一组特定且有限的运算，旨在帮助加速 MLP 和小型神经网络的评估。
在协作向量 API 中，使用函数 optixCoopVecMatMul 完成带偏差的向量矩阵乘法，该函数执行层评估的仿射部分。由于通常需要在不同阶段使用不同的激活函数，因此在向量 – 矩阵乘法后单独应用激活函数，并且可以通过协作向量 API 提供的一组向量函数来构建激活函数。
outputVector = inputVector × matrix + bias
所有 RTX 设备和特定服务器级 GPU 上的 OptiX 均支持协作向量。您可以使用 optixDeviceContextGetProperty 和 OPTIX_DEVICE_PROPERTY_COOP_VEC 查询设备支持。在不支持的设备上使用协作向量将生成错误，因为没有后备支持可用。
实现示例
本节将探讨用于在 OptiX 着色器中执行推理或 MLP 层评估的协作向量 API。我们将学习的示例改编自 OptiX SDK 中的 optixNeuralTexture 示例。此示例使用 NTC SDK ，它可以处理训练 (压缩) 纹理，以指定的文件格式存储权重和偏差，并演示使用不同着色语言在各种情况下的推理 (解压缩) 。您可以使用 NTC SDK 压缩自己的纹理，然后使用 optixNeuralTexture 进行查看。
使用协作向量的代码大量使用 C++ 模板。模板通过提供静态数组大小和编译时已知的静态数据类型，帮助编译器生成高性能代码。提前定义常用类型，使代码更易于阅读。
using T_OUT = OptixCoopVec<float, 16 /*output channels*/>;
...
T_OUT texel = inferTexel<T_OUT> (
latents, weights, x, y, ... );
如此一来，您便可使用 T_OUT 类型的快捷键来代替模板化的 OptixCoopVec<> 类型。函数 evalMLP 针对屏幕上的给定像素评估完整的 MLP。在伪代码术语中，它将设置 MLP 的输入，然后评估 MLP 的每一层，最后返回最后一层的输出：
template <class T_OUT>
evalMLP( T_OUT& outLayer, latents, mlpWeights, x, y )
{
using T_IN = OptixCoopVec<half, 48 /* input vec size */ >;
using T_HID = OptixCoopVec<half, 64 /* hidden layer size */ >;
T_IN networkInputs = prepareNetworkInputs_FP16<T_IN>(x, y, latents);
T_HID hiddenOut1 = evalLayer<T_IN, T_HID>(
networkInputs, mlpWeights, 0, scaleBiasOffset, hiddenOut1);
T_HID hiddenOut2 = evalLayer<T_HID, T_HID>(
hiddenOut1, mlpWeights, weightOffset1, scaleBiasOffset, hiddenOut2);
T_HID hiddenOut3 = evalLayer<T_HID, T_HID>(
hiddenOut2, mlpWeights, weightOffset2, scaleBiasOffset, hiddenOut3 );
outLayer = evalLayer<T_HID, T_OUT>(
hiddenOut3, mlpWeights, weightOffset3, scaleBiasOffset, outLayer);
return true;
}
请注意每层评估的输出如何成为下一层评估的输入。详细了解层评估：
template <class T_IN, class T_OUT> evalLayer(
T_IN& inputArray,
uint8_t* weights,
uint32_t weightsOffsetInBytes,
uint32_t& biasOffsetInBytes,
T_OUT& outputArray )
{
outputArray = optixCoopVecMatMul <
T_OUT,
T_IN,
OPTIX_COOP_VEC_ELEM_TYPE_FLOAT8_E4M3, // inputInterpretation
MAT_LAYOUT, // matrixLayout
false, // transpose
T_OUT::size, // N
T_IN::size, // K
OPTIX_COOP_VEC_ELEM_TYPE_FLOAT8_E4M3, // matrixElementType
OPTIX_COOP_VEC_ELEM_TYPE_FLOAT16 // biasElementType
>(
inputArray, // inputVector
weights, // matrix base ptr
weightsOffsetInBytes, // matrix offset
weights, // bias base ptr, same as weights
biasOffsetInBytes // bias offset
);
// increment offset to the next layer
biasOffsetInBytes += T_OUT::size * sizeof( T_OUT::value_type );
outputArray = activate<T_OUT>( outputArray );
}
单层计算只不过是对 optixCoopVecMatMul 的封装。矩阵乘法后，将偏移量增加到偏置向量，以便为下一层做好准备 (请注意，偏置偏移量是在此函数中通过参考传递的) 。然后调用层上的激活函数。
您可能在这些代码示例中注意到，我们将相同的 weights base pointer 传递给多次调用 evalLayer 的函数，并将此基指针用于 weights 和 biases。通过向基本指针（本例中为 weightsOffsetInBytes 或 biasOffsetInBytes）添加常量偏移值，在每个步骤中查找正确的数据。
使用这种方式编写代码有两个原因。第一个原因是，在读取 NTC 格式的文件时，API 会返回一个内存块，其中 weights matrices 和 bias vectors 均被紧密打包，您可以使用简单的运算来迭代每层的数据。第二个原因是，当重复使用相同的 weights base pointers 时，cooperative vectors API 会利用对 optixCoopVecMatMul 的连续调用。编译器将注意到重复使用的 base pointers (即使使用不同的常量偏移量) ，并将优化您的程序，以防止层之间不必要地发生 shuffling 和 unshuffling 操作。
最后，我们来看看 activation function：
template<class T_IN>
VecT activate(const T_IN& x, bool scaleActivation=true)
{
T_IN tmp = optixCoopVecFFMA( x, T_IN( 1.0f/3.0f ), T_IN( 0.5f ) );
tmp = optixCoopVecMin( optixCoopVecMax( tmp, 0.0f ), 1.f ); // clamp
T_IN result = optixCoopVecMin( x, 3.0f );
result = optixCoopVecMul( result, tmp );
if( scaleActivation )
result = optixCoopVecFFMA( result, T_IN(invStep), T_IN(bias) );
return result;
}
由于 MLP 激活函数对层输出向量的每个元素应用非线性映射，因此上述代码中调用的协作向量函数均为向量运算，而非矩阵运算。与应用层权重矩阵相比，激活操作通常要小得多，且成本更低。有一组有限的内置向量函数可与通常在 MLP 激活函数中找到的协作向量一起使用，例如 tanh、log2、exp2、min、max、ffma 等。
一些协作向量函数具有采用标量参数的变体，但并非全部。在某些情况下，您需要从标量中创建具有常量值的向量。在这里，这是使用 OptixCoopVec 构造函数完成的，例如，使用 T_IN (0.5 f) 参数进行第一次 optixCoopVecFFMA 调用。这种特定的激活函数来自 NTC SDK 。通常，在设计您自己的网络时，激活可以像调用 optixCoopVecMax 来模拟众所周知的 ReLU 激活一样简单。
神经图形
协作向量用于实现 RTX Neural Shaders 和 RTX Neural Texture Compression 。这些技术可作为 NVIDIA RTX Kit 的一部分提供，NVIDIA RTX Kit 是一套开源资源库，旨在简化这些技术的使用和集成。如需快速了解 RTX Kit (包括每个资源库的链接和资源) ，请参阅 使用 NVIDIA RTX Kit 开始使用神经网络渲染 。
图 5 中描绘的巨龙需要进行纹理压缩，才能容纳 GeForce RTX 5080 GPU 的 16 GB 显存。单是巨龙就有超过 100 个 8K UDIM 纹理，每个纹理有五个层。如果将纹理从文件解压缩到内存中，它们将消耗超过 32 GB 的 VRAM，是 GeForce RTX 5080 上可用内存的两倍多。
借助 NTC，巨龙纹理的内存占用在不到 3 GB 的情况下变得更加合理，大约是 BC 压缩纹理的一半。这为场景中的其余纹理以及几何图形、动画、BVH 和着色器留下了充足的空间，使这个大规模生产规模的场景能够在单个 5080 GPU 上实时渲染。
RTX Neural Shading SDK 中展示了神经网络着色器，其中提供的示例可帮助您学习如何训练自己的神经网络着色，然后在正常图形渲染过程中使用这些着色器执行推理。协作向量可以作为实现 Real-Time Neural Appearance Models 的一种方法。
性能注意事项
请考虑以下事项以获得最佳性能：
- 混洗和取消混洗： 为利用 Tensor Core，在调用
optixCoopVecMatMul之前对数据进行混洗，然后取消混洗。在两次此类调用之间，如果您仅使用支持的向量运算，则无需进行解洗和重洗，从而提高性能。 - 全线程束 ：使用全线程束时性能最佳。使用 SER 合并线程，并避免在动态条件内调用
optixCoopVecMatMul。 - 显存布局：权重矩阵布局会显著影响性能。OptiX 支持推理和训练的最佳布局，应使用这些布局以获得最佳性能。使用
optixCoopVecMatrixConvert将矩阵转换为最佳布局。
使用 OptiX 协作向量进行训练
OptiX 支持使用协作向量进行训练，包括向前和向后传播。有关更多信息，请参阅 OptiX 编程指南 。特别是，两个设备端内部函数 optixCoopVecReduceSumAccumulate 和 tg_ 23 分别有助于累积偏置向量和权重矩阵的损失值。
开始使用
协作向量是 NVIDIA OptiX 数据类型和 API，用于在 OptiX 着色器程序中执行高性能向量和矩阵运算。这些向量和矩阵运算是多层感知器 (MLPs) 等常见机器学习算法的核心。协作向量有助于在 NVIDIA RTX GPU 上使用 Tensor Cores，而这之前需要在线程束中的线程之间进行显式协调。借助协作向量，开发者不再需要使用同步多线程技术。它们可以使用更简单的单线程编程风格，执行对这些神经网络算法至关重要的高效矩阵向量乘法运算。
协作向量可从 NVIDIA OptiX SDK 9.0 开始使用。我们还通过 4 月底的 Agility SDK 预览版、 Vulkan 和 Slang 将协作向量 API 引入 DirectX ，以便您可以在支持硬件加速光线追踪的任何地方使用它们。OptiX 协作向量 API 文档可在 OptiX 编程指南 中获取，可在线获取，也可通过 SDK 和示例分发 PDF 格式。
在 OptiX SDK 中，有一个 RTX Neural Texture Compression 推理示例 (称为 optixNeuralTexture) ，该示例使用 cooperative vectors 在着色过程中实时解压缩神经压缩纹理，与热门的 BC5 或 BC6 压缩纹理格式 相比，可节省 20 倍的内存，与 optixMeshViewer 示例中使用的未压缩纹理相比，可节省 80 倍的纹理占用空间。
随着时间的推移，我们很高兴看到协作向量出现新的有趣用例。加入 OptiX NVIDIA 开发者论坛 的对话，了解更多信息并发布您的体验。
