Cooperative Vectors Introduction

A brief history

In 2021 we started investigating use cases for Neural Networks (NN) in our rendering engine. Our first successful attempt was with Neural Materials (NM).

Initially we only needed support for inference, since training of the NM was done "offline" in PyTorch. At the time, hardware accelerated inference was only supported through early vendor specific extensions on vulkan (Cooperative Matrix). Therefore, we built our own infrastructure for NN inference. This was built on top of our render graph, and fully in compute shaders (hlsl) without the use of any extension, to be able to deploy on all our target platforms and backends. One year down the line we saw impressive results from Neural Radiance Caching (NRC), which required runtime training of (mostly small, 16, 32 or 64 features wide) NNs. This led to the expansion of our framework to support inference and training pipelines.

While GPU’s had hardware dedicated to accelerate NN inference, there was not a way of accessing these features in shaders, especially cross-platform. For example, on desktop hardware we have different options per vendor: NVIDIA Tensor Cores, at the time, were only accessible through CUDA, and Intel's Xe Matrix Extensions (XMX) were vendor-specific, while **AMD **had Wave Matrix Multiply Accumulate (WMMA) instructions that use the "normal" shader cores.

As mentioned above, the Cooperative Matrix extension SPV_NV_cooperative_matrix,** **now elevated to KHR, was introduced as an abstraction for (tiled) matrix-to-matrix operations with the ability to access hardware accelleration, if present. This extension added new shader instructions to interpret buffers as matrices and to multiply them together. Similarly, DirectX added a very similar feature called WaveMatrix, initally intended for Shader Model (SM) 6.8, but then itnever released out of preview.

As a result of research on Neural Materials and Neural Texture Compression (NTC) by NVIDIA, they introduced a feature in Optix and a Vulkan extension: Cooperative Vector VK_NV_cooperative_vector. The solution proposed in the NTC paper compresses a set of textures for a material (like albedo, normal, roughness and metallic) into a NN and a learned texture-like representation. Both presented an interesting challenge: each material would have its own network, thus resulting in a situation where adjacent pixels on the screen might sample different textures, requiring an evaluation of a different network and therefore a different set of weights. This is not currently possible with Cooperative Matrix, which are meant for non-divergent work.

Similarly NM, have the same issue, where different pixels might require different sets of weights. The way we solved it in our inital implementation was to bucket queries to the same materials and run multiple dispatches, one per material. This solution is not ideal, but works in practice, whilst being cumbersome and quite involved, ideally this should just be a branch in your shaders. Cooperative Vector solves this challenge by shifting interface from a matrix-matrix (in Cooperative Matrix) to a vector-matrix operation.

What is a Cooperative Vector?

Let's clarify some terminology first: a cooperative vector, or long vector in Dx/HLSL, is the data structure itself. To avoid confusion, from now onwards, we will use long vector,** **when referring to the data structure. On the other hand, we will use Cooperative Vector to refer to the operations and hardware acceleration mechanisms that work with long vectors.

The "Cooperative" in Cooperative Vector refers to an implementation detail of the hardware acceleration, where individual vector-matrix multiply requests submitted by threads in a wave are combined into a matrix-matrix operation accelerated collectively for the wave. This name doesn't appear in HLSL code itself, just vector types and operations like vector-matrix multiplication as shown in the examples below.

So what are they exactly? Long vectors in HLSL are simply vector data structures that can be much longer than the usual 4 elements we've been used to in shaders. Nothing more, nothing less. They're a transparent data structure where the user can read and write individual elements via index, do operations on vectors like addition, multiplication, max (ReLU), etc.

Thread / Execution Model

Now that we have an idea of what are Cooperative Vectors, where can they be used? At the moment, on both Vulkan and DirectX, they should be accessible in all shader stages, although we mostly use them in compute shaders . That being said, since the idea is to provide a way to support divergent data; in pixel and compute shaders we will have each pixel/thread hold its own long vector.

The Matrices

Matrices come in two flavors:

Plain Layouts:

• Row Major
• Column Major

Optimal Layouts:

• MulOptimal (Inference)
• OuterProductOptimal (Training/Accumulation)

Accessibility in Shaders

Both Vulkan and DirectX provide interfaces that you can use to point to a memory address, like a buffer, and use it as a matrix.

For plain layouts (row/column major), we can address elements if we know the data type, matrix size, and layout. For optimal layouts, however, the memory addresses are completely opaque: we cannot do any read or write to a specific element, since the layout will be defined by the driver.

Inference and Training of a Network

One of the simplest NN, the Multi Layer Perceptron (MLP), is built as a sequence of linear layers and activations. Each layer computation can be performed by a single matrix-matrix or vector-matrix operation and addition of a bias and finally an activation, like ReLU.

Inference

For inference, we can use either row/column layout or MulOptimal (inference optimal) for our weight matrices. Obviously, in the majority of cases, you'll want MulOptimal.