RTX Neural Shading (RTXNS)

RTXNS uses a modular architecture for neural networks, similar to PyTorch's torch.nn.Module. The core abstraction is the IModule interface, which defines a differentiable forward method.

Modules are combined into architectures using a pattern analogous to torch.nn.Sequential. This allows for generic, reusable training functions that can operate on any model conforming to a specific IModule signature (e.g., a specific number of inputs and outputs) without knowing the internal architecture.

Training typically follows these steps:

1. Generate random input (e.g., UV coordinates).
2. Sample a target value (e.g., from a reference texture).
3. Evaluate the model using a differentiable forward method.
4. Compute loss gradient.
5. Backpropagate the gradient through the network parameters using bwd_diff.

interface IModule<T : __BuiltinFloatingPointType, let NumInputs : int, let NumOutputs : int>
{
    [BackwardDifferentiable]
    CoopVec<T, NumOutputs> forward(CoopVec<T, NumInputs> inputParams);
}

The LossAccumulation module provides a GPU-friendly system for accumulating floating-point loss components into a RWByteAddressBuffer.

LossConfig

Describes how loss components are stored:

• componentCount: Number of components to accumulate (clamped to LOSS_ACCUM_MAX_COMPONENTS, which is 128).
• baseByteOffset: Starting byte offset in the target RWByteAddressBuffer.

Key Functions

• AtomicAddF32Portable: Automatically selects between native InterlockedAddF32 or a CAS-based fallback (AtomicAddFloat) for compatibility.
• Zero<N>: Initializes an array of $N$ elements to 0.0f.
• AddInPlace<N>: Adds src[N] into dst[N] for the first componentCount elements.
• AccumulateComponents<N>: Performs atomic floating-point addition of an array of $N$ components into the global buffer at config.baseByteOffset.

Typical Workflow

1. Maintain per-thread or per-sample component arrays.
2. Perform local reduction.
3. Call AccumulateComponents to merge local results into the global GPU buffer.

public struct LossConfig
{
    uint componentCount;
    uint baseByteOffset;
    public uint GetComponentCount() { return min(componentCount, LOSS_ACCUM_MAX_COMPONENTS); }
}

// Accumulation usage
AccumulateComponents<MAX_COMPONENTS>(accumulationBuffer, gLossShared[0], lossConfig);

The library uses rtxns::MatrixLayout to manage how weights and biases are organized. Layouts are categorized into two types:

1. Host Layouts (rtxns::MatrixLayout::RowMajor, rtxns::MatrixLayout::ColumnMajor): Hardware and API agnostic. Use these for storing networks to disk or sharing between different GPUs/APIs.
2. Device Optimal Layouts (rtxns::MatrixLayout::InferencingOptimal, rtxns::MatrixLayout::TrainingOptimal): Opaque, hardware-specific formats. These often include specific alignment and padding for performance and should not be transferred between different GPUs.

Typical Lifecycle:

1. Initialize in a Host Layout (new or from file).
2. Upload to GPU and convert to rtxns::MatrixLayout::TrainingOptimal using rtxns::NetworkUtilities::ConvertWeights.
3. After training, convert to rtxns::MatrixLayout::InferencingOptimal for deployment.
4. To save, convert back to a Host Layout before writing to file.

The MLP module provides two primary structures for executing neural networks depending on your use case.

InferenceMLP

Used for inference-only workloads. It provides a forward pass function but does not support backpropagation.

TrainingMLP

Used for training workloads. It provides both a forward pass and a backward pass. Unlike the manual LinearOp_Backward, TrainingMLP uses Slang's auto-differentiation functionality to automatically generate the backward propagation function.

Both structures require specifying the floating-point type T, number of HIDDEN_LAYERS, INPUTS, HIDDEN layer size, OUTPUTS size, and the matrixLayout / componentType for the weights.

// Inference usage pattern
CoopVec<T, OUTPUTS> forward<Act, FinalAct>(CoopVec<T, INPUTS> inputParams, Act act, FinalAct finalAct);

// Training usage pattern
void backward<Act, FAct>(CoopVec<T, INPUTS> ip, Act act, FAct fact, CoopVec<T, OUTPUTS> loss);

To create deep architectures where the output of one module serves as the input to another, use the ModuleChain type. ModuleChain implements the IModule interface by composing two existing modules.

ModuleChain requires:

• T: The floating point type.
• NumInputs: The input dimension of the first module.
• NumHidden: The output dimension of the first module (which must match the input dimension of the second).
• NumOutputs: The output dimension of the second module.
• First: An IModule<T, NumInputs, NumHidden>.
• Second: An IModule<T, NumHidden, NumOutputs>.

To chain more than two modules (e.g., modules A, B, and C), you must nest the ModuleChain type. You first create a chain of A and B, then create a second chain using the result of the first as the First module and C as the Second module.

struct ModuleChain<
    T : __BuiltinFloatingPointType,
    let NumInputs : int,
    let NumHidden : int,
    let NumOutputs : int,
    First : IModule<T, NumInputs, NumHidden>,
    Second : IModule<T, NumHidden, NumOutputs>
> : IModule<T, NumInputs, NumOutputs>
{
    First first;
    Second second;

    [BackwardDifferentiable]
    CoopVec<T, NumOutputs> forward(CoopVec<T, NumInputs> inputParams)
    {
        CoopVec<T, NumHidden> middle = first.forward(inputParams);
        return second.forward(middle);
    }
}

The training loop executes both a training pass and an optimizer pass multiple times per frame to accelerate convergence. In the provided example, BATCH_COUNT (defined in NetworkConfig.h, typically 100) determines how many iterations occur per frame.

1. Training Pass: Dispatches a compute shader to calculate forward/backward passes and gradients.
2. Optimizer Pass: Dispatches a compute shader (e.g., Adam) to update weights using the calculated gradients.

Example loop structure:

for (int i = 0; i < BATCH_COUNT; ++i)
{
    nvrhi::ComputeState state;
    ...
    // Training pass
    state.bindings = { m_trainingPass.bindingSet };
    state.pipeline = m_trainingPass.pipeline;
    ...
    m_commandList->setComputeState(state);
    m_commandList->dispatch(m_batchSize / THREADS_PER_GROUP_TRAIN, 1, 1);
    ...
    // Optimizer pass
    state.bindings = { m_optimizerPass.bindingSet };
    state.pipeline = m_optimizerPass.pipeline;
    ...
    m_commandList->setComputeState(state);
    m_commandList->dispatch(div_ceil(m_totalParameterCount, THREADS_PER_GROUP_OPTIMIZE), 1, 1);
    ...
}