Выпуск NVIDIA OptiX 9.0 вводит новую функцию под названием «кооперативные векторы», которая включает рабочие процессы ИИ в ядра трассировки лучей. Эта функция использует тензорные ядра NVIDIA RTX для аппаратного ускорения матричных операций и вычислений нейронных сетей во время оттенения. Это открывает новые возможности для технологий рендеринга с помощью ИИ, таких как NVIDIA RTX Neural Shaders и NVIDIA RTX Neural Texture Compression (NTC), и движется в сторону фотореалистичных материалов кино-качества в рендеринге в реальном времени.
Кооперативные векторные API внедряются в OptiX, DirectX, NVAPI, Slang и Vulkan. Этот пост исследует концепции кооперативных векторов, которые применяются во всех API, и демонстрирует пример использования API OptiX.
Почему матричные операции?
Многослойный персептрон (MLP) является основным строительным блоком многих алгоритмов нейронных сетей. Исследования показали, что MLP способны точно воспроизводить эффекты, на которых они обучены. Даже если они достаточно малы для работы в реальном времени, MLP могут обрабатывать интересные эффекты, такие как физически обоснованное оттенение, иногда быстрее и с меньшими затратами памяти, чем традиционные сети оттенения.
Обычно MLP состоит из вектора входных данных, нескольких полностью связанных слоев и вектора выходных данных. Размеры различных векторных слоев не обязательно должны быть одинаковыми.
Каждый слой оценки MLP (инференса) имеет две фазы: взвешенную и смещенную линейную комбинацию значений предыдущего слоя и необязательную нелинейную активационную функцию. Взвешенная и смещенная линейная комбинация сводится к умножению матрицы на вектор, за которым следует сложение вектора смещения, также известное как афинное преобразование.
Состав любого двух афинных преобразований также является афинным преобразованием. Это означает, что при наличии только афинной линейной плюс смещенной фазы на каждом слое, весь MLP всегда можно свести к одному афинному преобразованию. Однако MLP гораздо более выразительны благодаря нелинейным активационным функциям, применяемым после афинного результата каждого слоя.
В полностью связанном MLP каждый нейрон является функцией всех нейронов в предыдущем слое, как показано на Рисунке 1. Полное концептуальное вычисление одного слоя показано на Рисунке 2.
Почему кооперативные векторы?
Одна из целей кооперативных векторов заключается в том, чтобы обеспечить использование NVIDIA Tensor Cores для ускорения матричных операций. Как правило, программная модель CUDA SIMT требует полного ворпа активных потоков для выполнения этой задачи, но модель программирования трассировки лучей обрабатывает потоки независимо и не гарантирует полного ворпа. Кроме того, тензорные ядра предлагают умножение матрицы на матрицу, но каждый поток трассировки лучей нуждается только в умножении вектора на матрицу, что может привести к недостаточному использованию тензорных ядер.
Кроме того, API CUDA требует нацеливания на конкретные версии оборудования и не гарантирует обратную совместимость от архитектуры к архитектуре. Для введения в многопоточный подход CUDA к умножению матриц ознакомьтесь с Руководством пользователя по фоновым данным о умножении матриц.
Кооперативные векторы решают эти ограничения, предоставляя API, который:
- Позволяет выполнять матричные операции с ворпами, содержащими некоторые неактивные потоки
- Обеспечивает совместимость между архитектурами вперед и назад
- Позволяет пользователям задавать векторные данные в одном потоке, одновременно переназначая операцию для более эффективного использования тензорных ядер
Кооперативные векторы могут обрабатывать расхождения данных и выполнения в ворпе, хотя и с некоторым ухудшением производительности. Лучшая производительность будет получена, когда веса MLP во всем ворпе одинаковы, и когда у вас есть полный комплект потоков в ворпе. Использование перестановки выполнения шейдеров (SER) может помочь добиться обоих этих целей.
Поскольку оценка MLP представляет собой серию умножений векторов на матрицы, когда все потоки в ворпе оценивают один и тот же MLP бок о бок, API кооперативного вектора может рассматривать комбинированную аффинную операцию ворпа как умножение матрицы на матрицу плюс смещение. Это то, что означает кооперативность: потоки объединяются, чтобы превратить несколько операций умножения векторов на матрицы в операции умножения матриц на матрицы.
outputMatrix = inputMatrix × weightsMatrix + biasMatrix
Здесь все матрицы, кроме матрицы весов, имеют по 32 ряда, и каждый ряд входных, выходных и смещенных матриц представляет данные для отдельного потока.
Использование кооперативных векторов в OptiX
Кооперативный вектор — это непрозрачный тип вектора, по сути класс массива, который может иметь произвольную длину. OptiX предоставляет реализацию кооперативного вектора, называемого OptixCoopVec. Кооперативные векторы в OptiX поддерживают определенный и ограниченный набор операций, предназначенных для ускорения оценки MLP и небольших нейронных сетей.
В API кооперативного вектора умножение вектора на матрицу со смещением выполняется с помощью функции optixCoopVecMatMul, которая осуществляет аффинную часть оценки слоя. Поскольку использование различных активационных функций на различных этапах часто является желательным, активация применяется отдельно после умножения вектора на матрицу и может быть построена из набора векторных функций, предоставленных API кооперативного вектора.
outputVector = inputVector × matrix + bias
Кооперативные векторы поддерживаются в OptiX на всех устройствах RTX и некоторых серверных GPU. Вы можете запросить поддержку устройства, используя optixDeviceContextGetProperty с OPTIX_DEVICE_PROPERTY_COOP_VEC. Использование кооперативных векторов на неподдерживаемых устройствах приведет к генерации ошибок, поскольку резервная поддержка недоступна.
Пример реализации
Этот раздел исследует API кооперативных векторов для выполнения инференса или оценки слоя MLP в шейдере OptiX. Пример, который мы будем рассматривать, адаптирован из примера optixNeuralTexture в SDK OptiX. Этот пример использует 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. После умножения матриц увеличиваем смещение к вектору смещения, чтобы подготовить его к следующему слою (обратите внимание, что смещениеBias передается по ссылке в этой функции). Затем вызываем активационную функцию для слоя.
Вы могли заметить, что мы передаем указатель на одну и ту же базу весов для нескольких вызовов evalLayer и также используем этот базовый указатель как для весов, так и для смещений. Найти правильные данные на каждом этапе можно, добавляя постоянные значения смещений к базовому указателю, в данном случае это weightsOffsetInBytes или biasOffsetInBytes.
Существует две причины, по которым код пишется таким образом. Первая причина заключается в том, что при чтении файлов в формате NTC API возвращает блок памяти, в котором матрицы весов и векторы смещения плотно упакованы. Таким образом, вы можете использовать простую арифметику для перебора данных для каждого слоя. Вторая причина заключается в том, что API кооперативных векторов использует последовательные вызовы к optixCoopVecMatMul, когда используются одни и те же базовые указатели на веса. Компилятор заметит повторно используемые базовые указатели (даже при использовании различных постоянных смещений) и оптимизирует вашу программу, чтобы предотвратить ненужные операции перемещения и перемещения между слоями.
Наконец, взгляните на активационную функцию:
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.5f) в первом вызове optixCoopVecFFMA. Эта конкретная активационная функция взята из NTC SDK. В общем, при проектировании собственной сети активация может быть такой же простой, как вызов optixCoopVecMax для имитации известной активации ReLU.
Нейронная графика
Кооперативные векторы используются для реализации RTX Neural Shaders и RTX Neural Texture Compression. Эти технологии доступны как часть NVIDIA RTX Kit, набора открытых репозиториев, чтобы упростить использование и интеграцию этих технологий. Для быстрого старта с RTX Kit, включая ссылки и ресурсы для каждого репозитория, смотрите Начните работу с нейронным рендерингом, используя NVIDIA RTX Kit.
Дракон, изображенный на Рисунке 5, требовал сжатия текстуры, чтобы поместиться в 16 ГБ видеопамяти GPU GeForce RTX 5080. Один дракон имеет более 100 текстур 8K UDIM с пятью слоями каждая. Если текстуры будут декомпрессированы из файлов в память, они займут более 32 ГБ видеопамяти, что более чем в два раза превышает доступную память на 5080.
С NTC размер памяти текстур дракона становится гораздо более разумным — менее 3 ГБ, и это примерно в два раза меньше, чем сжатые текстуры BC. Это оставляет много места для остальных текстур в сцене, а также для геометрии, анимации, BVH и шейдеров, позволяя этой огромной сцене производственного масштаба рендериться в реальном времени на одном GPU 5080.
Нейронные шейдеры демонстрируются в RTX Neural Shading SDK, который предоставляет примеры, которые помогут вам научиться обучать свои собственные сети нейронного shading и затем использовать их для выполнения инференса в рамках нормального графического рендеринга. Кооперативные векторы могут стать способом реализации Моделей Нейронных Появлений в Реальном Времени.
Соображения по производительности
Для достижения наилучшей производительности учитывайте следующее:
- Перестановки: Чтобы использовать тензорные ядра, данные перераспределяются по ворпу перед вызовом
optixCoopVecMatMulи снова распределяются после. Между двумя такими вызовами, если вы используете только поддерживаемые векторные операции, это может убрать необходимость в перестановке и повторной перестановке, улучшая производительность. - Полные ворпы: Производительность лучше всего при использовании полных ворпов. Используйте SER для объединения потоков и избегайте вызовов
optixCoopVecMatMulвнутри динамических условных операторов. - Макет памяти: Макет матриц весов значительно влияет на производительность. OptiX поддерживает оптимальные макеты для инференса и обучения, которые следует использовать для достижения наилучшей производительности. Используйте
optixCoopVecMatrixConvertдля преобразования матриц в оптимальные макеты.
Обучение с кооперативными векторами OptiX
OptiX поддерживает обучение с кооперативными векторами, включая прямое и обратное распространение. Для получения дополнительной информации смотрите Руководство по программированию OptiX. В частности, два встроенных устройства optixCoopVecReduceSumAccumulate и optixCoopVecOuterProductAccumulate помогают накопить значения потерь по векторам смещения и матрицам весов соответственно.
Начало работы
Кооперативные векторы являются типом данных и API NVIDIA OptiX для выполнения высокопроизводительных векторных и матричных операций в рамках программ шейдеров OptiX. Эти векторные и матричные операции лежат в основе общих алгоритмов машинного обучения, таких как многослойные персептроны (MLP). Кооперативные векторы упрощают использование тензорных ядер на GPU NVIDIA RTX, которые ранее требовали явной координации между потоками в ворпе. С кооперативными векторами разработчики больше не нуждаются в использовании синхронных многопоточных техник. Они могут выполнять эффективные операции умножения матриц на векторы, которые являются основными для этих алгоритмов нейронных сетей, используя более простой однопоточный стиль программирования.
Кооперативные векторы доступны начиная с NVIDIA OptiX SDK 9.0. API кооперативных векторов также внедряются в DirectX через предварительную версию Agility SDK в конце апреля, Vulkan и Slang, так что вы можете использовать их в любом месте, где поддерживается аппаратная трассировка лучей. Документация для API кооперативных векторов OptiX доступна в Руководстве по программированию OptiX, доступной в онлайн-формате и в формате PDF, распространяемом с SDK и примерами.
В SDK OptiX представлен пример инференса для RTX Neural Texture Compression, называемый optixNeuralTexture, который использует кооперативные векторы для декомпрессии нейронно сжатых текстур на лету во время оттенения, обеспечивая экономию памяти в 20 раз по сравнению с популярным BC5 или BC6 сжатым форматом текстуры, или экономию в 80 раз по сравнению с несжатыми текстурами, используемыми в примере optixMeshViewer.
Будет интересно увидеть новые и интересные случаи использования кооперативных векторов, возникающие со временем. Присоединяйтесь к обсуждению на Форуме разработчиков NVIDIA OptiX, чтобы узнать больше и поделиться своим опытом.
