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 開發者論壇 的對話，了解更多信息并發布您的體驗。
?
