slang Documentation

Making it easier to work with shaders

5,520 stars C++ C++ 1 file ยท ~49,358 tokens #cuda#d3d12#glsl#hlsl#shaders#vulkan
RAW Doc

Docker & Self-Hosted Deployment Verified OCI

Deploy slang on any server, homelab, or cloud container platform.

1-Click Terminal Command
bash (docker run)
docker run -d --name slang --restart always shader-slang/slang:latest

1. Project Overview & Quickstart (shader-slang/slang)

File: README.md

Slang

Slang is a shading language that makes it easier to build and maintain large shader codebases in a modular and extensible fashion, while also maintaining the highest possible performance on modern GPUs and graphics APIs.
Slang is based on years of collaboration between researchers at NVIDIA, Carnegie Mellon University, Stanford, MIT, UCSD and the University of Washington.

Why Slang?

The Slang shading language is designed to enable real-time graphics developers to work with large-scale, high-performance shader code.

Write Shaders Once, Run Anywhere

The Slang compiler can generate code for a wide variety of targets: D3D12, Vulkan, Metal, D3D11, CUDA, and even generate code to run on a CPU. For textual targets, such as Metal Shading Language (MSL) and CUDA, Slang produces readable code that preserves original identifier names, as well as the type and call structure, making it easier to debug.

Access the Latest GPU Features

Slang code is highly portable, but can still leverage unique platform capabilities, including the latest features in Direct3D and Vulkan. For example, developers can make full use of pointers when generating SPIR-V.
Slang's capability system helps applications manage feature set differences across target platforms by ensuring code only uses available features during the type-checking step, before generating final code. Additionally, Slang provides flexible interop features to enable directly embedding target code or SPIR-V into generated shaders.

Leverage Neural Graphics with Automatic Differentiation

Slang can automatically generate both forward and backward derivative propagation code for complex functions that involve arbitrary control flow and dynamic dispatch. This allows existing rendering codebases to easily become differentiable, or for Slang to serve as the kernel language in a PyTorch-driven machine learning framework via slangtorch.

Scalable Software Development with Modules

Slang provides a module system that enables logical organization of code for separate compilation. Slang modules can be independently compiled offline to a custom IR (with optional obfuscation) and then linked at runtime to generate code in formats such as DXIL or SPIR-V.

Code Specialization that Works with Modules

Slang supports generics and interfaces (a.k.a. type traits/protocols), allowing for clear expression of shader specialization without the need for preprocessor techniques or string-pasting. Unlike C++ templates, Slang's generics are pre-checked and don't produce cascading error messages that are difficult to diagnose. The same generic shader can be specialized for a variety of different types to produce specialized code ahead of time, or on the fly, entirely under application control.

Easy On-ramp for HLSL and GLSL Codebases

Slang's syntax is similar to HLSL, and most existing HLSL code can be compiled with the Slang compiler out-of-the-box, or with just minor modifications. This allows existing shader codebases to immediately benefit from Slang without requiring a complete rewrite or port.

Slang provides a compatibility module that enables the use of most GLSL intrinsic functions and GLSL's parameter binding syntax.

Comprehensive Tooling Support

Slang comes with full support of IntelliSense editing features in Visual Studio Code and Visual Studio through the Language Server Protocol.
Full debugging capabilities are also available through RenderDoc and SPIR-V based tools.

Getting Started

The fastest way to get started using Slang in your own development is to use a pre-built binary package, available through GitHub releases.
Slang binaries are also included in the Vulkan SDK since version 1.3.296.0.

There are packages built for x86_64 and aarch64 Windows, Linux and macOS.
Each binary release includes the command-line slangc compiler, a shared library for the compiler, and the slang.h header.

See the user-guide for info on using the slangc command-line tool: Slang Command Line Usage.

If you want to try out the Slang language without installing anything, a fast and simple way is to use the Slang Playground. The playground allows you to compile Slang code to a variety of targets, and even run some simple shaders directly within the browser. The playground loads Slang compiler to your browser and runs all compilation locally. No data will be sent to any servers.

If you would like to build Slang from source, please consult the build instructions.

Documentation

The Slang project provides a variety of different documentation, but most users would be well served starting with the User's Guide.

For developers writing Slang code, the Slang Core Module Reference provides detailed documentation on Slang's built-in types and functions.

We also provide a few examples of how to integrate Slang into a rendering application.

These examples use a graphics layer that we include with Slang called "GFX" which is an abstraction library of various graphics APIs (D3D11, D2D12, OpenGL, Vulkan, CUDA, and the CPU) to support cross-platform applications using GPU graphics and compute capabilities.
GFX is being deprecated in favor of slang-rhi.

Additionally, we recommend checking out Vulkan Mini Examples for more examples of using Slang's language features available on Vulkan, such as pointers and the ray tracing intrinsics.

Contributing

If you'd like to contribute to the project, we are excited to have your input.
The following guidelines should be observed by contributors:

  • Please follow the contributor Code of Conduct.
  • Bugs reports and feature requests should go through the GitHub issue tracker
  • Changes should ideally come in as small pull requests on top of master, coming from your own personal fork of the project
  • Large features that will involve multiple contributors or a long development time should be discussed in issues, and broken down into smaller pieces that can be implemented and checked in in stages

Contribution guide describes the workflow for contributors at more detail.

Limitations and Support

Platform support

The Slang compiler and libraries can be built on the following platforms:

Windows Linux MacOS WebAssembly
supported supported supported experimental

Both x86_64 and aarch64 architectures are supported on Windows, Linux and MacOS platforms.

Target support

Slang can compile shader code to the following targets:

Target Status Output Formats
Direct3D 11 supported HLSL
Direct3D 12 supported HLSL
Vulkan supported SPIRV, GLSL
Metal experimental* Metal Shading Language
WebGPU experimental** WGSL
CUDA supported C++ (compute only)
Optix experimental C++ (WIP)
CPU experimental C++ (kernel), C++ (host), standalone executable, dynamic library

*Slang currently supports generating vertex, fragment, compute, task and mesh
shaders for Metal.

**WGSL support is still work in-progress.

For greater detail, see the Supported Compilation
Targets
section of the
User Guide

The Slang project has been used for production applications and large shader
codebases, but it is still under active development. Support is currently
focused on the platforms (Windows, Linux) and target APIs (Direct3D 12, Vulkan)
where Slang is used most heavily. Users who are looking for support on other
platforms or APIs should coordinate with the development team via the issue
tracker to make sure that their use cases can be supported.

License

The Slang code itself is under the Apache 2.0 with LLVM Exception license (see LICENSE).

Builds of the core Slang tools depend on the following projects, either automatically or optionally, which may have their own licenses:

Slang releases may include LLVM under the license:

  • llvm (Apache 2.0 License with LLVM exceptions)

Some of the tests and example programs that build with Slang use the following projects, which may have their own licenses:

  • glm (MIT)
  • stb_image and stb_image_write from the stb collection of single-file libraries (Public Domain)
  • tinyobjloader (MIT)

File: docs/deprecated/a1-02-slangpy.md


Using Slang to Write PyTorch Kernels

Note

This documentation is about slang-torch, a way to use Slang with Python and PyTorch.
For new projects, we recommend exploring SlangPy as an alternative.
We plan to deprecate slang-torch in favor of SlangPy in the near future, and we will communicate any plans in advance.

If you are a PyTorch user seeking to write complex, high-performance, and automatically differentiated kernel functions using a per-thread programming model, we invite you to try Slang. Slang is a cutting-edge shading language that provides a straightforward way to define kernel functions that run incredibly fast in graphics applications. With the latest addition of automatic differentiation and PyTorch interop features, Slang offers an efficient solution for developing auto-differentiated kernels that run at lightning speed with a strongly typed, per-thread programming model.

One of the primary advantages of a per-thread programming model in kernel programming is the elimination of concerns regarding maintaining masks for branches. When developing a kernel in Slang, you can use all control flow statements, composite data types (structs, arrays, etc.), and function calls without additional effort. Code created with these language constructs can be automatically differentiated by the compiler without any restrictions. Additionally, Slang is a strongly typed language, which ensures that you will never encounter type errors at runtime. Most code errors can be identified as you type thanks to the compiler's coding assistance service, further streamlining the development process.

In addition, using a per-thread programming model also results in more optimized memory usage. When writing a kernel in Slang, most intermediate results do not need to be written out to global memory and then read back, reducing global memory bandwidth consumption and the delay caused by these memory operations. As a result, a Slang kernel can typically run at higher efficiency compared to the traditional bulk-synchronous programming model.

Getting Started with SlangTorch

In this tutorial, we will use a simple example to walk through the steps to use Slang in your PyTorch project.

Installation

slangtorch is available via PyPI, so you can install it simply through

sh
pip install slangtorch

Note that slangtorch requires torch with CUDA support. See the pytorch installation page to find the right version for your platform.

You can check that you have the right installation by running:

sh
python -c "import torch; print(f'cuda: {torch.cuda.is_available()}')"

Writing Slang kernels for `slangtorch` >= **v1.1.5**

From v2023.4.0, Slang supports auto-binding features that make it easier than ever to invoke Slang kernels from python, and interoperate seamlessly with pytorch tensors.

Here's a barebones example of a simple squaring kernel written in Slang (square.slang):

csharp
[AutoPyBindCUDA]
[CUDAKernel]
void square(TensorView<float> input, TensorView<float> output)
{
    // Get the 'global' index of this thread.
    uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();

    // If the thread index is beyond the input size, exit early.
    if (dispatchIdx.x >= input.size(0))
        return;

    output[dispatchIdx.x] = input[dispatchIdx.x] * input[dispatchIdx.x];
}

This code follows the standard pattern of a typical CUDA kernel function. It takes as input
two tensors, input and output.
It first obtains the global dispatch index of the current thread and performs range check to make sure we don't read or write out
of the bounds of input and output tensors, and then calls square() to compute the per-element result, and
store it at the corresponding location in output tensor.

slangtorch works by compiling kernels to CUDA and it identifies the functions to compile by checking for the [CUDAKernel] attribute.
The second attribute [AutoPyBindCUDA] allows us to call square directly from python without having to write any host code. If you would like to write the host code yourself for finer control, see the other version of this example here.

You can now simply invoke this kernel from python:

python
import torch
import slangtorch

m = slangtorch.loadModule('square.slang')

A = torch.randn((1024,), dtype=torch.float).cuda()

output = torch.zeros_like(A).cuda()

# Number of threads launched = blockSize * gridSize
m.square(input=A, output=output).launchRaw(blockSize=(32, 1, 1), gridSize=(64, 1, 1))

print(output)

The python script slangtorch.loadModule("square.slang") returns a scope that contains a handle to the square kernel.

The kernel can be invoked by

  1. calling square and binding torch tensors as arguments for the kernel, and then
  2. launching it using launchRaw() by specifying CUDA launch arguments to blockSize & gridSize. (Refer to the CUDA documentation for restrictions around blockSize)

Note that for semantic clarity reasons, calling a kernel requires the use of keyword arguments with names that are lifted from the .slang implementation.

Invoking derivatives of kernels using slangtorch

The [AutoPyBindCUDA] attribute can also be used on differentiable functions defined in Slang, and will automatically bind the derivatives. To do this, simply add the [Differentiable] attribute.

One key point is that the basic TensorView<T> objects are not differentiable. They can be used as buffers for data that does not require derivatives, or even as buffers for the manual accumulation of derivatives.

Instead, use the DiffTensorView type for when you need differentiable tensors. Currently, DiffTensorView only supports the float dtype variety.

Here's a barebones example of a differentiable version of square:

csharp
[AutoPyBindCUDA]
[CUDAKernel]
[Differentiable]
void square(DiffTensorView input, DiffTensorView output)
{
    uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim();

    if (dispatchIdx.x >= input.size(0))
        return;
    
    output[dispatchIdx.x] = input[dispatchIdx.x] * input[dispatchIdx.x];
}

Now, slangtorch.loadModule("square.slang") returns a scope with three callable handles square, square.fwd for the forward-mode derivative & square.bwd for the reverse-mode derivative.

You can invoke square() normally to get the same effect as the previous example, or invoke square.fwd() / square.bwd() by binding pairs of tensors to compute the derivatives.

python
import torch
import slangtorch

m = slangtorch.loadModule('square.slang')

input = torch.tensor((0, 1, 2, 3, 4, 5), dtype=torch.float).cuda()
output = torch.zeros_like(input).cuda()

# Invoke normally
m.square(input=input, output=output).launchRaw(blockSize=(6, 1, 1), gridSize=(1, 1, 1))

print(output)

# Invoke reverse-mode autodiff by first allocating tensors to hold the gradients
input = torch.tensor((0, 1, 2, 3, 4, 5), dtype=torch.float).cuda()
input_grad = torch.zeros_like(input).cuda()

output = torch.zeros_like(input)
# Pass in all 1s as the output derivative for our example
output_grad = torch.ones_like(output) 

m.square.bwd(
    input=(input, input_grad), output=(output, output_grad)
).launchRaw(
    blockSize=(6, 1, 1), gridSize=(1, 1, 1))

# Derivatives get propagated to input_grad
print(input_grad)

# Note that the derivatives in output_grad are 'consumed'.
# i.e. all zeros after the call.
print(output_grad)

slangtorch also binds the forward-mode version of your kernel (propagate derivatives of inputs to the output) which can be invoked the same way using module.square.fwd()

You can refer to this documentation for a detailed reference of Slang's automatic differentiation feature.

Wrapping your kernels as pytorch functions

pytorch offers an easy way to define a custom operation using torch.autograd.Function, and defining the .forward() and .backward() members.

This can be a very helpful way to wrap your Slang kernels as pytorch-compatible operations. Here's an example of the square kernel as a differentiable pytorch function.

python
import torch
import slangtorch

m = slangtorch.loadModule("square.slang")

class MySquareFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        output = torch.zeros_like(input)

        kernel_with_args = m.square(input=input, output=output)
        kernel_with_args.launchRaw(
            blockSize=(32, 32, 1),
            gridSize=((input.shape[0] + 31) // 32, (input.shape[1] + 31) // 32, 1))

        ctx.save_for_backward(input, output)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        (input, output) = ctx.saved_tensors

        input_grad = torch.zeros_like(input)
        
        # Note: When using DiffTensorView, grad_output gets 'consumed' during the reverse-mode.
        # If grad_output may be reused, consider calling grad_output = grad_output.clone()
        #
        kernel_with_args = m.square.bwd(input=(input, input_grad), output=(output, grad_output))
        kernel_with_args.launchRaw(
            blockSize=(32, 32, 1),
            gridSize=((input.shape[0] + 31) // 32, (input.shape[1] + 31) // 32, 1))
        
        return input_grad

Now we can use the autograd function MySquareFunc in our python script:

python
x = torch.tensor((3.0, 4.0), requires_grad=True, device='cuda')
print(f"X = {x}")
y_pred = MySquareFunc.apply(x)
loss = y_pred.sum()
loss.backward()
print(f"dX = {x.grad.cpu()}")

Output:

text
X = tensor([3., 4.],
           device='cuda:0', requires_grad=True)
dX = tensor([6., 8.])

And that's it! slangtorch.loadModule uses JIT compilation to compile your Slang source into CUDA binary.
It may take a little longer the first time you execute the script, but the compiled binaries will be cached and as long as the kernel code is not changed, future runs will not rebuild the CUDA kernel.

Because the PyTorch JIT system requires ninja, you need to make sure ninja is installed on your system
and is discoverable from the current environment, you also need to have a C++ compiler available on the system.
On Windows, this means that Visual Studio need to be installed.

Specializing shaders using slangtorch

slangtorch.loadModule allows specialization parameters to be specified since it might be easier to write shaders with placeholder definitions that can be substituted at load-time.
For instance, here's a sphere tracer that uses a compile-time specialization parameter for its maximum number of steps (N):

csharp
float sphereTrace<let N:int>(Ray ray, SDF sdf)
{
    var pt = ray.o;
    for (int i = 0; i < N; i++)
    {
        pt += sdf.eval(pt) * ray.d;
    }

    return pt;
}

float render(Ray ray)
{
    // Use N=20 for sphere tracing.
    float3 pt = sphereTrace<20>(ray, sdf);
    return shade(pt, sdf.normal());
}

However, instead of using a fixed 20 steps, the renderer can be configured to use an arbitrary compile-time constant.

csharp
// Compile-time constant. Expect "MAX_STEPS" to be set by the loadModule call.
static const uint kMaxSteps = MAX_STEPS;

float render(Ray ray)
{
    float3 pt = sphereTrace<kMaxSteps>(ray, sdf);
    return shade(pt, sdf.normal());
}

Then multiple versions of this shader can be compiled from Python using the defines argument:

python
import slangtorch

sdfRenderer20Steps = slangtorch.loadModule('sdf.slang', defines={"MAX_STEPS": 20})
sdfRenderer50Steps = slangtorch.loadModule('sdf.slang', defines={"MAX_STEPS": 50})
...

This is often helpful for code re-use, parameter sweeping, comparison/ablation studies, and more, from the convenience of Python.

Back-propagating Derivatives through Complex Access Patterns

In most common scenarios, a kernel function will access input tensors in a complex pattern instead of mapping
1:1 from an input element to an output element, like the square example shown above. When you have a kernel
function that access many different elements from the input tensors and use them to compute an output element,
the derivatives of each input element can't be represented directly as a function parameter, like the x in square(x).

Consider a 3x3 box filtering kernel that computes for each pixel in a 2D image, the average value of its
surrounding 3x3 pixel block. We can write a Slang function that computes the value of an output pixel:

csharp
float computeOutputPixel(TensorView<float> input, uint2 pixelLoc)
{
    int width = input.size(0);
    int height = input.size(1);

    // Track the sum of neighboring pixels and the number
    // of pixels currently accumulated.
    int count = 0;
    float sumValue = 0.0;

    // Iterate through the surrounding area.
    for (int offsetX = -1; offsetX <= 1; offsetX++)
    {
        // Skip out of bounds pixels.
        int x = pixelLoc.x + offsetX;
        if (x < 0 || x >= width) continue;

        for (int offsetY = -1; offsetY <= 1; offsetY++)
        {
            int y = pixelLoc.y + offsetY;
            if (y < 0 || y >= height) continue;
            sumValue += input[x, y];
            count++;
        }
    }

    // Compute the average value.
    sumValue /= count;

    return sumValue;
}

We can define our kernel function to compute the entire output image by calling computeOutputPixel:

csharp
[CudaKernel]
void boxFilter_fwd(TensorView<float> input, TensorView<float> output)
{
    uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy;
    int width = input.dim(0);
    int height = input.dim(1);
    if (pixelLoc.x >= width) return;
    if (pixelLoc.y >= height) return;

    float outputValueAtPixel = computeOutputPixel(input, pixelLoc)

    // Write to output tensor.
    output[pixelLoc] = outputValueAtPixel;
}

How do we define the backward derivative propagation kernel? Note that in this example, there
isn't a function like square that we can just mark as [Differentiable] and
call bwd_diff(square) to get back the derivative of an input parameter.

In this example, the input comes from multiple elements in a tensor. How do we propagate the
derivatives to those input elements?

The solution is to wrap tensor access with a custom function:

csharp
float getInputElement(
    TensorView<float> input,
    TensorView<float> inputGradToPropagateTo,
    uint2 loc)
{
    return input[loc];
}

Note that the getInputElement function simply returns input[loc] and is not using the
inputGradToPropagateTo parameter. That is intended. The inputGradToPropagateTo parameter
is used to hold the backward propagated derivatives of each input element, and is reserved for later use.

Now we can replace all direct accesses to input with a call to getInputElement. The
computeOutputPixel can be implemented as following:

csharp
[Differentiable]
float computeOutputPixel(
    TensorView<float> input,
    TensorView<float> inputGradToPropagateTo,
    uint2 pixelLoc)
{
    int width = input.dim(0);
    int height = input.dim(1);

    // Track the sum of neighboring pixels and the number
    // of pixels currently accumulated.
    int count = 0;
    float sumValue = 0.0;

    // Iterate through the surrounding area.
    for (int offsetX = -1; offsetX <= 1; offsetX++)
    {
        // Skip out of bounds pixels.
        int x = pixelLoc.x + offsetX;
        if (x < 0 || x >= width) continue;

        for (int offsetY = -1; offsetY <= 1; offsetY++)
        {
            int y = pixelLoc.y + offsetY;
            if (y < 0 || y >= height) continue;
            sumValue += getInputElement(input, inputGradToPropagateTo, uint2(x, y));
            count++;
        }
    }

    // Compute the average value.
    sumValue /= count;

    return sumValue;
}

The main changes compared to our original version of computeOutputPixel are:

  • Added a inputGradToPropagateTo parameter.
  • Modified input[x,y] with a call to getInputElement.
  • Added a [Differentiable] attribute to the function.

With that, we can define our backward kernel function:

csharp
[CudaKernel]
void boxFilter_bwd(
    TensorView<float> input,
    TensorView<float> resultGradToPropagateFrom,
    TensorView<float> inputGradToPropagateTo)
{
    uint2 pixelLoc = (cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx()).xy;
    int width = input.dim(0);
    int height = input.dim(1);
    if (pixelLoc.x >= width) return;
    if (pixelLoc.y >= height) return;

    bwd_diff(computeOutputPixel)(input, inputGradToPropagateTo, pixelLoc);
}

The kernel function simply calls bwd_diff(computeOutputPixel) without taking any return values from the call
and without writing to any elements in the final inputGradToPropagateTo tensor. But when exactly does the propagated
output get written to the output gradient tensor (inputGradToPropagateTo)?

And that logic is defined in our final piece of code:

csharp
[BackwardDerivativeOf(getInputElement)]
void getInputElement_bwd(
    TensorView<float> input,
    TensorView<float> inputGradToPropagateTo,
    uint2 loc,
    float derivative)
{
    float oldVal;
    inputGradToPropagateTo.InterlockedAdd(loc, derivative, oldVal);
}

Here, we are providing a custom defined backward propagation function for getInputElement.
In this function, we simply add derivative to the element in inputGradToPropagateTo tensor.

When we call bwd_diff(computeOutputPixel) in boxFilter_bwd, the Slang compiler will automatically
differentiate all operations and function calls in computeOutputPixel. By wrapping the tensor element access
with getInputElement and by providing a custom backward propagation function of getInputElement, we are effectively
telling the compiler what to do when a derivative propagates to an input tensor element. Inside the body
of getInputElement_bwd, we define what to do then: atomically adds the derivative propagated to the input element
in the inputGradToPropagateTo tensor. Therefore, after running boxFilter_bwd, the inputGradToPropagateTo tensor will contain all the
back propagated derivative values.

Again, to understand all the details of the automatic differentiation system, please refer to the
Automatic Differentiation chapter for a detailed explanation.

Manually binding kernels

[AutoPyBindCUDA] works for most use cases, but in certain situations, it may be necessary to write the host function by hand. The host function can also be written in Slang, and slangtorch handles its compilation to C++.

Here's the same square example from before:

csharp
// square.slang
float compute_square(float x)
{
    return x * x;
}

[CudaKernel]
void square_kernel(TensorView<float> input, TensorView<float> output)
{
    uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();

    if (globalIdx.x >= input.size(0))
        return;

    float result = compute_square(input[globalIdx.x]);

    output[globalIdx.x] = result;
}

To manually invoke this kernel, we then need to write a CPU(host) function that defines how this kernel is dispatched. This can be defined in the same Slang file:

csharp
[TorchEntryPoint]
TorchTensor<float> square(TorchTensor<float> input)
{
    var result = TorchTensor<float>.zerosLike(input);
    let blockCount = uint3(1);
    let groupSize = uint3(result.size(0), result.size(1), 1);
    __dispatch_kernel(square_kernel, blockCount, groupSize)(input, result);
    return result;
}

Here, we mark the function with the [TorchEntryPoint] attribute, so it will be compiled to C++ and exported as a python callable.
Since this is a host function, we can perform tensor allocations. For instance, square() calls TorchTensor<float>.zerosLike to allocate a 2D-tensor that has the same size as the input.
zerosLike returns a TorchTensor<float> object that represents a CPU handle of a PyTorch tensor.

Then we launch square_kernel with the __dispatch_kernel syntax. Note that we can directly pass
TorchTensor<float> arguments to a TensorView<float> parameter and the compiler will automatically convert the type and obtain a view into the tensor that can be accessed by the GPU kernel function.

Calling a `[TorchEntryPoint]` function from Python

You can use the following code to call square from Python:

python
import torch
import slangtorch

m = slangtorch.loadModule("square.slang")

x = torch.randn(2,2)
print(f"X = {x}")
y = m.square(x)
print(f"Y = {y.cpu()}")

Result output:

text
X = tensor([[ 0.1407,  0.6594],
        [-0.8978, -1.7230]])
Y = tensor([[0.0198, 0.4349],
        [0.8060, 2.9688]])

Manual binding for kernel derivatives

The above example demonstrates how to write a simple kernel function in Slang and call it from Python.
Another major benefit of using Slang is that the Slang compiler support generating backward derivative
propagation functions automatically.

In the following section, we walk through how to use Slang to generate a backward propagation function
for square, and expose it to PyTorch as an autograd function.

First we need to tell Slang compiler that we need the square function to be considered a differentiable function, so Slang compiler can generate a backward derivative propagation function for it:

csharp
[Differentiable]
float square(float x)
{
    return x * x;
}

This is done by simply adding a [Differentiable] attribute to our square function.

With that, we can now define square_bwd_kernel that performs backward propagation as:

csharp
[CudaKernel]
void square_bwd_kernel(TensorView<float> input, TensorView<float> grad_out, TensorView<float> grad_propagated)
{
    uint3 globalIdx = cudaBlockIdx() * cudaBlockDim() + cudaThreadIdx();

    if (globalIdx.x >= input.size(0) || globalIdx.y >= input.size(1))
        return;

    DifferentialPair<float> dpInput = diffPair(input[globalIdx.xy]);
    var gradInElem = grad_out[globalIdx.xy];
    bwd_diff(square)(dpInput, gradInElem);
    grad_propagated[globalIdx.xy] = dpInput.d;
}

Note that the function follows the same structure of square_fwd_kernel, with the only difference being that
instead of calling into square to compute the forward value for each tensor element, we are calling bwd_diff(square)
that represents the automatically generated backward propagation function of square.
bwd_diff(square) will have the following signature:

csharp
void bwd_diff_square(inout DifferentialPair<float> dpInput, float dOut);

Where the first parameter, dpInput represents a pair of original and derivative value for input, and the second parameter,
dOut, represents the initial derivative with regard to some latent variable that we wish to back-prop through. The resulting
derivative will be stored in dpInput.d. For example:

csharp
// construct a pair where the primal value is 3, and derivative value is 0.
var dp = diffPair(3.0);
bwd_diff(square)(dp, 1.0);
// dp.d is now 6.0

Similar to square_fwd, we can define the host side function square_bwd as:

csharp
[TorchEntryPoint]
TorchTensor<float> square_bwd(TorchTensor<float> input, TorchTensor<float> grad_out)
{
    var grad_propagated = TorchTensor<float>.zerosLike(input);
    let blockCount = uint3(1);
    let groupSize = uint3(input.size(0), input.size(1), 1);
    __dispatch_kernel(square_bwd_kernel, blockCount, groupSize)(input, grad_out, grad_propagated);
    return grad_propagated;
}

Builtin Library Support for PyTorch Interop

As shown in previous tutorial, Slang has defined the TorchTensor<T> and TensorView<T> type for interop with PyTorch
tensors. The TorchTensor<T> represents the CPU view of a tensor and provides methods to allocate a new tensor object.
The TensorView<T> represents the GPU view of a tensor and provides accessors to read write tensor data.

Following is a list of built-in methods and attributes for PyTorch interop.

`TorchTensor` methods

`static TorchTensor TorchTensor.alloc(uint x, uint y, ...)`

Allocates a new PyTorch tensor with the given dimensions. If T is a vector type, the length of the vector is implicitly included as the last dimension.
For example, TorchTensor<float3>.alloc(4, 4) allocates a 3D tensor of size (4,4,3).

`static TorchTensor TorchTensor.emptyLike(TorchTensor other)`

Allocates a new PyTorch tensor that has the same dimensions as other without initializing it.

`static TorchTensor TorchTensor.zerosLike(TorchTensor other)`

Allocates a new PyTorch tensor that has the same dimensions as other and initialize it to zero.

`uint TorchTensor.dims()`

Returns the tensor's dimension count.

`uint TorchTensor.size(int dim)`

Returns the tensor's size (in number of elements) at dim.

`uint TorchTensor.stride(int dim)`

Returns the tensor's stride (in bytes) at dim.

`TensorView` methods

`TensorView.operator[uint x, uint y, ...]`

Provide an accessor to data content in a tensor.

`TensorView.operator[vector index]`

Provide an accessor to data content in a tensor, indexed by a uint vector.
tensor[uint3(1,2,3)] is equivalent to tensor[1,2,3].

`uint TensorView.dims()`

Returns the tensor's dimension count.

`uint TensorView.size(int dim)`

Returns the tensor's size (in number of elements) at dim.

`uint TensorView.stride(int dim)`

Returns the tensor's stride (in bytes) at dim.

`void TensorView.fillZero()`

Fills the tensor with zeros. Modifies the tensor in-place.

`void TensorView.fillValue(T value)`

Fills the tensor with the specified value, modifies the tensor in-place.

`T* TensorView.data_ptr_at(vector index)`

Returns a pointer to the element at index.

`void TensorView.InterlockedAdd(vector index, T val, out T oldVal)`

Atomically add val to element at index.

`void TensorView.InterlockedMin(vector index, T val, out T oldVal)`

Atomically computes the min of val and the element at index. Available for 32 and 64 bit integer types only.

`void TensorView.InterlockedMax(vector index, T val, out T oldVal)`

Atomically computes the max of val and the element at index. Available for 32 and 64 bit integer types only.

`void TensorView.InterlockedAnd(vector index, T val, out T oldVal)`

Atomically computes the bitwise and of val and the element at index. Available for 32 and 64 bit integer types only.

`void TensorView.InterlockedOr(vector index, T val, out T oldVal)`

Atomically computes the bitwise or of val and the element at index. Available for 32 and 64 bit integer types only.

`void TensorView.InterlockedXor(vector index, T val, out T oldVal)`

Atomically computes the bitwise xor of val and the element at index. Available for 32 and 64 bit integer types only.

`void TensorView.InterlockedExchange(vector index, T val, out T oldVal)`

Atomically swaps val into the element at index. Available for float and 32/64 bit integer types only.

`void TensorView.InterlockedCompareExchange(vector index, T compare, T val)`

Atomically swaps val into the element at index if the element equals to compare. Available for float and 32/64 bit integer types only.

`DiffTensorView` methods

`DiffTensorView.operator[uint x, uint y, ...]`

Provide an accessor to data content in a tensor. This method is differentiable, and has the same semantics as using a .load() to get data, and .store() to set data.

`DiffTensorView.operator[vector index]`

Provide an accessor to data content in a tensor, indexed by a uint vector.tensor[uint3(1,2,3)] is equivalent to tensor[1,2,3]. This method is differentiable, and has the same semantics as using a .load() to get data, and .store() to set data.

`float DiffTensorView.load(vector index)`

Loads the 32-bit floating point data at the specified multi-dimensional index. This method is differentiable, and in reverse-mode will perform an atomic-add.

`void DiffTensorView.store(vector index, float val)`

Stores the 32-bit floating point value val at the specified multi-dimensional index. This method is differentiable, and in reverse-mode will perform an atomic exchange to retrieve the derivative and replace with 0.

`float DiffTensorView.loadOnce(vector index)`

Loads the 32-bit floating point data at the specified multi-dimensional index. This method is differentiable, and uses a simple store for the reverse-mode for faster gradient aggregation, but loadOnce must be used at most once per index. loadOnce is ideal for situations where each thread loads data from a unique index, but will cause incorrect gradients when an index may be accessed multiple times.

`void DiffTensorView.storeOnce(vector index, float val)`

Stores the 32-bit floating point value val at the specified multi-dimensional index. This method is differentiable, and uses a simple load for the reverse-mode for faster gradient loading, but storeOnce must be used at most once per index. loadOnce is ideal for situations where each thread stores data to a unique index, but will cause incorrect gradient propagation when an index may be accessed multiple times.

`uint DiffTensorView.size(int dim)`

Returns the underlying primal tensor's size (in number of elements) at dim.

`uint DiffTensorView.dims()`

Returns the underlying primal tensor's dimension count.

`uint DiffTensorView.stride(uint dim)`

Returns the stride of the underlying primal tensor's dim dimension

CUDA Support Functions

`cudaThreadIdx()`

Returns the threadIdx variable in CUDA.

`cudaBlockIdx()`

Returns the blockIdx variable in CUDA.

`cudaBlockDim()`

Returns the blockDim variable in CUDA.

`syncTorchCudaStream()`

Waits for all pending CUDA kernel executions to complete on host.

Attributes for PyTorch Interop

`[CudaKernel]` attribute

Marks a function as a CUDA kernel (maps to a __global__ function)

`[TorchEntryPoint]` attribute

Marks a function for export to Python. Functions marked with [TorchEntryPoint] will be accessible from a loaded module returned by slangtorch.loadModule.

`[CudaDeviceExport]` attribute

Marks a function as a CUDA device function, and ensures the compiler to include it in the generated CUDA source.

`[AutoPyBindCUDA]` attribute

Marks a cuda kernel for automatic binding generation so that it may be invoked from python without having to hand-code the torch entry point. The marked function must also be marked with [CudaKernel]. If the marked function is also marked with [Differentiable], this will also generate bindings for the derivative methods.

Restriction: methods marked with [AutoPyBindCUDA] will not operate

Type Marshalling Between Slang and Python

Python-CUDA type marshalling for functions using `[AutoPyBindCUDA]`

When using auto-binding, aggregate types like structs are converted to Python namedtuples and are made available when using slangtorch.loadModule.

csharp
// mesh.slang
struct Mesh
{
    TensorView<float> vertices;
    TensorView indices;
};

[AutoPyBindCUDA]
[CUDAKernel]
void processMesh(Mesh mesh)
{
    /* ... */ 
}

Here, since Mesh is being used by renderMesh, the loaded module will provide Mesh as a python namedtuple with named fields.
While using the namedtuple is the best way to use structured arguments, they can also be passed as a python dict or tuple

python
m = slangtorch.loadModule('mesh.slang')

vertices = torch.tensor()
indices = torch.tensor()

# use namedtuple to provide structured input.
mesh = m.Mesh(vertices=vertices, indices=indices)
m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))

# use dict to provide input.
mesh = {'vertices': vertices, 'indices':indices}
m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))

# use tuple to provide input (warning: user responsible for right order)
mesh = (vertices, indices)
m.processMesh(mesh=mesh).launchRaw(blockSize=(32, 32, 1), gridSize=(1, 1, 1))

Python-CUDA type marshalling for functions using `[TorchEntryPoint]`

The return types and parameters types of an exported [TorchEntryPoint] function can be a basic type (e.g. float, int etc.), a vector type (e.g. float3), a TorchTensor<T> type, an array type, or a struct type.

When you use struct or array types in the function signature, it will be exposed as a Python tuple.
For example,

csharp
struct MyReturnType
{
    TorchTensor<T> tensors[3];
    float v;
}

[TorchEntryPoint]
MyReturnType myFunc()
{
    ...
}

Calling myFunc from python will result in a python tuple in the form of

text
[[tensor, tensor, tensor], float]

The same transform rules apply to parameter types.


File: docs/design/autodiff/basics.md

This documentation is intended for Slang contributors and is written from a compiler engineering point of view. For Slang users, see the user-guide at this link: https://shader-slang.com/slang/user-guide/autodiff.html

What is Automatic Differentiation?

Before diving into the design of the automatic differentiation (for brevity, we will call it 'auto-diff') passes, it is important to understand the end goal of what auto-diff tries to achieve.

The over-arching goal of Slang's auto-diff is to enable the user to compute derivatives of a given shader program or function's output w.r.t its input parameters. This critical compiler feature enables users to quickly use their shaders with gradient-based parameter optimization algorithms, which forms the backbone of modern machine learning systems. It enables users to train and deploy graphics systems that contain ML primitives (like multi-layer perceptron's or MLPs) or use their shader programs as differentiable primitives within larger ML pipelines.

More Resources

Here are some links to resources that talk more about differentiable programming from a more mathematical perspective:

  1. UCSD CSE 291 (Spring 2024): https://cseweb.ucsd.edu/~tzli/cse291/sp2024/
  2. UW CSE 5990 (Winter 2024): https://sites.google.com/cs.washington.edu/cse-599o-dppl

Definition of Derivatives

This section is based off of these slides: https://cseweb.ucsd.edu/~tzli/cse291/sp2024/lectures/03_forward_mode.pdf.

Here, we establish the mathematical definition of derivatives, starting with a simple 1D case (function with a single input and output), and extending to the general case of functions mapping multiple inputs to multiple outputs.

To avoid confusion, we will denote mathematical functions using LaTeX italic script ($f$, $g$, etc..) and programs that compute these functions with markdown code (f, g, etc..)

Derivatives of scalar (1D) functions

Consider the simplest case: a smooth scalar mathematical function that maps a real number to another real number:

$$f : \mathbb{R} \to \mathbb{R}$$

There are several definitions for a derivative, but we will use the definition that a derivative is the closest linear approximation of the output function at a given input location.
Concretely, given a specific input $x$, we can create a linear approximation of the function $f$ around $x$ as follows:

$$ f(x + dx) \approx f(x) + Df(x) \cdot dx $$

This can also be understood as a geometric 'tangent' to the function at $x$. $Df(x)$ is the slope of $f$ at $x$, i.e. $\frac{\partial f}{\partial x}$, and $dx$ is the perturbation away from $x$. Our approximation is linear as a function of the perturbation $dx$. Note that no matter how non-linear or complex the underlying function $f(x)$ is, the approximation is always linear (this property becomes very important later).

Forward-mode derivative functions

Now consider a concrete program f that computes some function.

c
// Computes square of x
float f(float x)
{
    return x * x;
}

What should its derivative program look like? We the need the output $f(x)$ and the product of derivative at $x$, $Df(x)$ with the differential $dx$.

In Slang, we put both of these together into a single function, called the forward-mode derivative function, which takes in a pair $(x, dx)$ returns a pair $(f(x), Df(x)\cdot dx)$ Note that in auto-diff literature, this is also often referred to as the total derivative function.

c
DifferentialPair<float> fwd_f(DifferentialPair<float> dpx)
{
    float x = dpx.getPrimal(); // Can also be accessed via property dpx.p
    float dx = dpx.getDifferential(); // Can also be accessed via property dpx.d
    return makePair(x * x, (2 * x) * dx);
}

Note that (2 * x) is the multiplier corresponding to $Df(x)$. We refer to $x$ and $f(x)$ as "primal" values and the perturbations $dx$ and $Df(x)\cdot dx$ as "differential" values. The reason for this separation is that the "differential" output values are always linear w.r.t their "differential" inputs.

As the name implies, DifferentialPair<T> is a special pair type used by Slang to hold values and their corresponding differentials.

Forward-mode derivatives for higher-dimensional functions

In practice, most functions tend to have multiple inputs and multiple outputs, i.e. $f: \mathbb{R}^N \to \mathbb{R}^M$

The definition above can be extended to higher dimensions, using the closest-linear-approximation idea. The main difference is that the derivative function represents a hyperplane rather than a line.

Effectively, we want our forward-mode derivative to compute the following:

$$ f(\mathbf{x} + \mathbf{dx}) \approx f(\mathbf{x}) + \langle Df(\mathbf{x}),\mathbf{dx}\rangle $$

Here, the input and its differential can be represented as a vector quantity $\mathbf{x}, \mathbf{dx} \in \mathbb{R}^N$ and the multiplier $Df(\mathbf{x})$ (also known as the Jacobian matrix) is a NxM matrix, and $\left< \cdot,\cdot \right>$ denotes the inner product (i.e. matrix-vector multiplication)

Here's an example of a Slang function taking in two inputs (N=2) and generating one output (M=1)

c
// Compute length of hypotenuse.
float f(float x, float y)
{
    return sqrt(x * x + y * y);
}

and its forward-mode derivative:

c
// Closest linear approximation at x, y
DifferentialPair<float> fwd_f(DifferentialPair<float> dpx, DifferentialPair<float> dpy)
{
    float x = dpx.p;
    float y = dpy.p;
    float dx = dpx.d;
    float dy = dpx.d;

    return DifferentialPair<float>(
        sqrt(x * x + y * y),                       // f(x, y)
        (x * dx + y * dy) / sqrt(x * x, y * y));   // <Df(x,y), dx>
}

Important note: the forward-mode function only needs to compute the inner product $\langle Df(\mathbf{x}),\mathbf{dx} \rangle$. The Jacobian matrix itself never needs to be fully materialized. This is a key design element of automatic differentiation, one which allows it to scale to huge input/output counts.

Building Blocks: Forward-mode derivatives compose in forward order of execution.

In practice, we compute forward-mode derivatives of a complex function by decomposing them into constituent functions (or in compiler-speak: instructions) and composing the forward-mode derivative of each piece in the same order.
This is because of each forward derivative is a 'right-side' product (or product of Jacobian matrix with a vector)

Here's an example of this in action (consider a complex function $h$ composed of $f$ and $g$):

$$ h(\mathbf{x}) = f(g(\mathbf{x})) $$

It's forward-mode derivative is then:

$$ \langle Dh(\mathbf{x}), \mathbf{dx}\rangle = \big\langle Df(\mathbf{x}), \langle Dg(\mathbf{x}), \mathbf{dx}\rangle\big\rangle $$

which is the forward-mode derivative of the outer function $f$ evaluated on the result of the forward-mode derivative of the inner function $g$.

An example of this in Slang code:

c
// Compute square.
float sqr(float x)
{
    return x * x;
}

// Compute length of hypotenuse.
float f(float x, float y)
{
    float x_sqr = sqr(x);
    float y_sqr = sqr(y)
    return sqrt(x_sqr + y_sqr);
}

The resulting derivative of f can be computed by composition:

c
// Forward-mode derivative of sqr()
DifferentialPair<float> fwd_sqr(DifferentialPair<float> dpx)
{
    float x = dpx.getPrimal();
    float dx = dpx.getDifferential();

    return DifferentialPair<float>(x * x, 2 * x * dx);
}

// Forward-mode derivative of f()
DifferentialPair<float> fwd_f(DifferentialPair<float> dpx, DifferentialPair<float> dpy)
{
    DifferentialPair<float> dp_x_sqr = fwd_sqr(dpx);
    DifferentialPair<float> dp_y_sqr = fwd_sqr(dpy);

    float x_sqr = dp_x_sqr.getPrimal();
    float y_sqr = dp_y_sqr.getPrimal();
    float x_sqr_d = dp_x_sqr.getDifferential();
    float y_sqr_d = dp_y_sqr.getDifferential();

    return DifferentialPair<float>(
        sqrt(x_sqr + y_sqr),
        (x_sqr_d + y_sqr_d) / sqrt(x_sqr + y_sqr));
}

Tip: Extracting partial derivatives from a forward-mode derivative (i.e. a 'total' derivative)

As we discussed above, forward-mode derivatives compute $\langle Df(\mathbf{x}),\mathbf{dx}\rangle$ rather than what you may be used to seeing in a calculus course (e.g. partial derivatives like $\frac{\partial f}{\partial x}$).

In fact, the forward-mode derivative is simply an product of the partial derivative w.r.t each input parameter multiplied by their differential perturbations $\frac{\partial f}{\partial x} * dx + \frac{\partial f}{\partial x} * dy$. This is the reason for the alternative name: total derivative.

Thus, partial derivative can be obtained by successively setting each input's differential to 1 (and 0 for everything else)
Example:

c
// Compute partial derivative w.r.t x (pass dx=1.0)
float df_dx = fwd_f(DifferentialPair<float>(x, 1.0), DifferentialPair<float>(y, 0.0)).d;

// Compute partial derivaive w.r.t y (pass dy=1.0)
float df_dy = fwd_f(DifferentialPair<float>(x, 0.0), DifferentialPair<float>(y, 1.0)).d;

Tip: Testing forward-mode derivatives using the first principles of calculus (i.e. the *finite difference* method)

In Calculus, partial derivatives of a function are often defined in a 'black box' manner using limits, by perturbing a single parameter by an infinitesimal amount:

$$ \frac{\partial f}{\partial x} = \lim_{dx\to 0} \frac{f(x + dx) - f(x - dx)}{2 * dx} $$

At the moment, we cannot leverage programming languages to compute true inifinitesimal limits, but we can replace $dx \to 0$ with a sufficiently small $\epsilon$ leading to the following 'test' to check if derivatives produced by automatic differentiation match with their true mathematical expected values.

Here's an example of using this idea to test functions (many autodiff tests were written this way)

c
// Compute partial derivative w.r.t x analytically
float df_dx_ad = fwd_f(DifferentialPair<float>(x, 1.0), DifferentialPair<float>(y, 0.0))

// Compute partial derivative w.r.t x through the finite difference (FD) method.
float eps = 1e-4
float df_dx_fd = (f(x + eps, y) - f(x - eps, y)) / (2 * eps);

// If computed correctly, df_dx_ad and df_dx_fd are very close.

Caveats:
Since the finite difference method only produces a biased estimate of the derivative, the result is only numerically close to the auto-diff-based result. Poorly behaved functions (those that rapidly change, or are discontinuous or otherwise non-differentiable) will result in a (expected) mismatch between FD and AD results.

Reverse-mode derivative functions

This section is based off of these slides: https://cseweb.ucsd.edu/~tzli/cse291/sp2024/lectures/05_reverse_mode.pdf.

Motivation: Challenges with scaling forward-mode derivatives

A big problem with forward-mode derivatives is their inability to scale to great parameter counts.

Machine learning pipelines often compute derivatives of a large complex pipeline with millions or even billions of input parameters, but a single output value, i.e. the loss or objective function, frequently denoted by $\mathcal{L}$.
Computing $\frac{\partial \mathcal{L}}{\partial x_i}$ for $N$ inputs $x_i$ using the one-hot vector approach will involve invoking the forward-mode derivative function $N$ times.

The reason for this limitation is that forward-mode derivatives pass derivatives from the inputs through to the outputs by computing the dot-product $\left< Df(\mathbf{x}),\mathbf{dx}\right>$.
Instead, we employ a different approach called the reverse-mode derivative, which propagates differentials backwards from outputs to inputs.

Key Idea: Generate code to compute $\langle \frac{\partial \mathcal{L}}{\partial f}, Df(\mathbf{x})\rangle$ rather than $\langle Df(\mathbf{x}),\mathbf{dx}\rangle$

The fundamental building blocks of reverse-mode derivatives are the left-side inner product. That is, the product of a vector of derivatives of w.r.t outputs $\frac{\partial \mathcal{L}}{\partial f}$ with the Jacobian matrix $Df(\mathbf{x})$.

An important thing to keep in mind is that it does not necessarily matter what the scalar quantity $\mathcal{L}$ is. The goal of this product is to propagate the derivatives of any scalar value $\mathcal{L}$ w.r.t output vector $f(\mathbf{x})$ (i.e., $\frac{\partial \mathcal{L}}{\partial f}$) into derivatives of that same scalar value $\mathcal{L}$ w.r.t the input vector $\mathbf{x}$ (i.e., $\frac{\partial \mathcal{L}}{\partial \mathbf{x}}$).

Here's an example of a Slang function computing the reverse-mode derivative.

c
// Compute length of hypotenuse
float f(float x, float y)
{
    return sqrt(x * x + y * y);
}

// Reverse-mode derivative of f. dOutput represents the derivative dL/dOutput of the output w.r.t scalar value.
void rev_f(inout DifferentialPair<float> dpx, inout DifferentialPair<float> dpy, float dOutput)
{
    float x = dpx.getPrimal();
    float y = dpy.getPrimal();

    float t = 1.0 / (sqrt(x * x + y * y));

    dpx = DifferentialPair<float>(
        x,                 // The primal part of the return value is *always* copied in from the input as-is.
        dOutput * x * t);  // The differential part for x is the derivative dL/dx computed as 
                           // (dL/dOutput) * (dOutput/dx), where dOutput/dx = x / sqrt(x*x+y*y).

    dpy = DifferentialPair<float>(
        y,                
        dOutput * y * t);  // The differential part for y is the derivative dL/dy computed as 
                           // (dL/dOutput) * (dOutput/dy), where dOutput/dy = y / sqrt(x*x+y*y).
}

Note that rev_f accepts derivatives w.r.t the output value as the input, and returns derivatives w.r.t inputs as its output (through inout parameters). rev_f still needs the primal values x and y to compute the derivatives, so those are still passed in as an input through the primal part of the differential pair.

Also note that the reverse-mode derivative function does not have to compute the primal result value (its return is void). The reason for this is a matter of convenience: reverse-mode derivatives are often invoked after all the primal functions, and there is typically no need for these values. We go into more detail on this topic in the checkpointing chapter.

The reverse mode function can be used to compute both dOutput/dx and dOutput/dy with a single invocation (unlike the forward-mode case where we had to invoke fwd_f once for each input)

c
DifferentialPair<float> dpx = makePair<float>(x, 0.f); // Initialize diff-value to 0 (not necessary)
DifferentialPair<float> dpx = makePair<float>(y, 0.f); // Initialize diff-value to 0 (not necessary)

rev_f(dpx, dpy, 1.0); // Pass 1.0 for dL/dOutput so that the results are (1.0 * dOutput/dx) and (1.0 * dOutput/dy)

float doutput_dx = dpx.getDifferential(); 
float doutput_dy = dpy.getDifferential();

Extension to multiple outputs

The extension to multiple outputs is fairly natural. Each output gets a separate input for its derivative.
Here is an example:

c
// Computation involving multiple inputs and outputs.
float2 f_multi_output(float x, float y)
{
    return float2(
        x * x,
        x + y);
}

// Reverse-mode derivative of 'f_multi_output'. The derivative of the outputs is also a vector quantity 
// (type follows from return type of f_multi_output)
void rev_f_multi_output(DifferentialPair<float> dpx, DifferentialPair<float> dpy, float2 dOut)
{
    float x = dpx.getPrimal();
    float y = dpy.getPrimal();

    dpx = DifferentialPair<float>(x, dOut[0] * 2 * x + dOut[1]);
    dpy = DifferentialPair<float>(x, dOut[1]);
}

Jacobian method: Generate forward- and reverse-mode derivatives from first principles.

A simple way to figure out what the generated reverse (or forward) derivative function is supposed to compute is to write down the entire Jacobian function. That is, write down the partial derivative of each input w.r.t each output

$$
D\mathbf{f}(\mathbf{x}) = \begin{bmatrix}
\partial f_0 / \partial x & \partial f_0 / \partial y \
\partial f_1 / \partial x & \partial f_1 / \partial y \
\end{bmatrix} =
\begin{bmatrix}
2x & 0.0 \
1.0 & 1.0 \
\end{bmatrix}
$$

The reverse-mode derivative's outputs should match the left-product of this matrix with the vector of derivatives w.r.t outputs:

$$ \left\langle \frac{\partial \mathcal{L}}{\partial \mathbf{f}}, D\mathbf{f}(\mathbf{x})\right\rangle =
\begin{bmatrix}
\frac{\partial \mathcal{L}}{\partial f_0} & \frac{\partial \mathcal{L}}{\partial f_1}
\end{bmatrix}
\begin{bmatrix}
2x & 0.0 \
1.0 & 1.0 \
\end{bmatrix} =
\begin{bmatrix} \left(\frac{\partial \mathcal{L}}{\partial f_0} \cdot 2x + \frac{\partial \mathcal{L}}{\partial f_1}\right) & \frac{\partial \mathcal{L}}{\partial f_1} \end{bmatrix}
$$

and the forward-mode derivative's outputs should match the right-product of this matrix with the vector of differentials of the inputs:

$$ \langle D\mathbf{f}(\mathbf{x}), d\mathbf{x}\rangle =
\begin{bmatrix}
2x & 0.0 \
1.0 & 1.0 \
\end{bmatrix}
\begin{bmatrix}
dx \ dy
\end{bmatrix} =
\begin{bmatrix} 2x \cdot dx & dx + dy \end{bmatrix}
$$

Note that when we generate derivative code in practice, we do not materialize the full Jacobian matrix, and instead use the composition property to chain together derivatives at the instruction level.
However, the resulting code is equivalent to the Jacobian method (mathematically), and it is a good, analytical way to confirm that the generated code is indeed correct (or when thinking about what the derivative of a particular instruction/set of instructions should be)

Building Blocks: Reverse-mode derivatives compose in reverse order of execution.

A consequence of using the 'left-side inner product' is that derivatives of a composite function must be computed in the reverse of the order of primal computation.

Here's an example of a composite function $h$ (similar to the example used in forward-mode building blocks):

$$ h(\mathbf{x}) = f(g(\mathbf{x})) $$

where (for brevity):

$$ \mathbf{y} = g(\mathbf{x}) $$

The reverse-mode derivative function for $h$ can be written as the composition of the reverse-mode derivatives of $f$ and $g$

$$ \left\langle \frac{\partial L}{\partial h}, Dh(\mathbf{x})\right\rangle = \left\langle \left\langle \frac{\partial L}{\partial h}, Df(\mathbf{y})\right\rangle , Dg(\mathbf{x})\right\rangle $$

Note the 'backward' order here. We must first pass the derivatives through the outer function $f$, and then pass the result through the inner function $g$ to compute derivatives w.r.t inner-most inputs $\mathbf{x}$. This process of passing derivatives backwards is often referred to as backpropagation.

A more concrete Slang example of the same:

c
// Compute square
float sqr(float x)
{
    return x * x;
}

// Compute length of hypotenuse
float f(float x, float y)
{
    return sqrt(sqr(x) + sqr(y));
}

The derivative functions are then:

c
void rev_sqr(DifferentialPair<float> dpx, float dOutput)
{
    float x = dpx.getPrimal();

    dpx = DifferentialPair<float>(x, dOutput * 2 * x);
}

void rev_f(DifferentialPair<float> dpx, DifferentialPair<float> dpy, float dOut)
{
    float t = 0.5f / sqrt(x * x + y * y);
    
    float d_xsqr = t * dOut; // Calculate derivatives w.r.t output of sqr(x)
    float d_ysqr = t * dOut; // Calculate derivatives w.r.t output of sqr(y)

    rev_sqr(dpx, d_xsqr); // Propagate to x
    rev_sqr(dpx, d_ysqr); // Propagate to y
}

When comparing rev_f's implementation to fwd_f, note the order of computing derivative w.r.t sqr (in rev_f, rev_sqr is called at the end, while in fwd_f it is called at the beginning)


File: docs/design/autodiff/decorators.md

This document details auto-diff-related decorations that are lowered in to the IR to help annotate methods with relevant information.

`[Differentiable]`

The [Differentiable] attribute is used to mark functions as being differentiable. The auto-diff process will only touch functions that are marked explicitly as [Differentiable]. All other functions are considered non-differentiable and calls to such functions from a differentiable function are simply copied as-is with no transformation.

Further, only [Differentiable] methods are checked during the derivative data-flow pass. This decorator is translated into BackwardDifferentiableAttribute (which implies both forward and backward differentiability), and then lowered into the IR OpBackwardDifferentiableDecoration

Note: [Differentiable] was previously implemented as two separate decorators [ForwardDifferentiable] and [BackwardDifferentiable] to denote differentiability with each type of auto-diff transformation. However, these are now deprecated. The preferred approach is to use only [Differentiable]

fwd_diff and bwd_diff cannot be directly called on methods that don't have the [Differentiable] tag (will result in an error). If non-[Differentiable] methods are called from within a [Differentiable] method, they must be wrapped in no_diff() operation (enforced by the derivative data-flow analysis pass )

`[Differentiable]` for `interface` Requirements

The [Differentiable] attribute can also be used to decorate interface requirements. In this case, the attribute is handled in a slightly different manner, since we do not have access to the concrete implementations.

The process is roughly as follows:

  1. During the semantic checking step, when checking a method that is an interface requirement (in checkCallableDeclCommon in slang-check-decl.cpp), we check if the method has a [Differentiable] attribute
  2. If yes, we construct create a set of new method declarations, one for the forward-mode derivative (ForwardDerivativeRequirementDecl) and one for the reverse-mode derivative (BackwardDerivativeRequirementDecl), with the appropriate translated function types and insert them into the same interface.
  3. Insert a new member into the original method to reference the new declarations (DerivativeRequirementReferenceDecl)
  4. When lowering to IR, the DerivativeRequirementReferenceDecl member is converted into a custom derivative reference by adding the OpBackwardDerivativeDecoration(deriv-fn-req-key) and OpForwardDerivativeDecoration(deriv-fn-req-key) decorations on the primal method's requirement key.

Here is an example of what this would look like:

c
interface IFoo
{
    [Differentiable]
    float bar(float);
};

// After checking & lowering
interface IFoo_after_checking_and_lowering
{
    [BackwardDerivative(bar_bwd)]
    [ForwardDerivative(bar_fwd)]
    float bar(float);

    void bar_bwd(inout DifferentialPair<float>, float);

    DifferentialPair<float> bar_fwd(DifferentialPair<float>);
};

Note: All conforming types must also declare their corresponding implementations as differentiable so that their derivative implementations are synthesized to match the interface signature. In this sense, the [Differentiable] attribute is part of the functions signature, so a [Differentiable] interface requirement can only be satisfied by a [Differentiable] function implementation

`[TreatAsDifferentiable]`

In large codebases where some interfaces may have several possible implementations, it may not be reasonable to have to mark all possible implementations with [Differentiable], especially if certain implementations use hacks or workarounds that need additional consideration before they can be marked [Differentiable]

In such cases, we provide the [TreatAsDifferentiable] decoration (AST node: TreatAsDifferentiableAttribute, IR: OpTreatAsDifferentiableDecoration), which instructs the auto-diff passes to construct an 'empty' function that returns a 0 (or 0-equivalent) for the derivative values. This allows the signature of a [TreatAsDifferentiable] function to match a [Differentiable] requirement without actually having to produce a derivative.

Custom derivative decorators

In many cases, it is desirable to manually specify the derivative code for a method rather than let the auto-diff pass synthesize it from the method body. This is usually desirable if:

  1. The body of the method is too complex, and there is a simpler, mathematically equivalent way to compute the same value (often the case for intrinsics like sin(x), arccos(x), etc..)
  2. The method involves global/shared memory accesses, and synthesized derivative code may cause race conditions or be very slow due to overuse of synchronization. For this reason Slang assumes global memory accesses are non-differentiable by default, and requires that the user (or the core module) define separate accessors with different derivative semantics.

The Slang front-end provides two sets of decorators to facilitate this:

  1. To reference a custom derivative function from a primal function: [ForwardDerivative(fn)] and [BackwardDerivative(fn)] (AST Nodes: ForwardDerivativeAttribute/BackwardDerivativeAttribute, IR: OpForwardDervativeDecoration/OpBackwardDerivativeDecoration), and
  2. To reference a primal function from its custom derivative function: [ForwardDerivativeOf(fn)] and [BackwardDerivativeOf(fn)] (AST Nodes: ForwardDerivativeAttributeOf/BackwardDerivativeAttributeOf). These attributes are useful to provide custom derivatives for existing methods in a different file without having to edit/change that module. For instance, we use diff.meta.slang to provide derivatives for the core module functions in hlsl.meta.slang. When lowering to IR, these references are placed on the target (primal function). That way both sets of decorations are lowered on the primal function.

These decorators also work on generically defined methods, as well as struct methods. Similar to how function calls work, these decorators also work on overloaded methods (and reuse the ResolveInoke infrastructure to perform resolution)

Checking custom derivative signatures

To ensure that the user-provided derivatives agree with the expected signature, as well as resolve the appropriate method when multiple overloads are available, we check the signature of the custom derivative function against the translated version of the primal function. This currently occurs in checkDerivativeAttribute()/checkDerivativeOfAttribute().

The checking process re-uses existing infrastructure from ResolveInvoke, by constructing a temporary invoke expr to call the user-provided derivative using a set of 'imaginary' arguments according to the translated type of the primal method. If ResolveInvoke is successful, the provided derivative signature is considered to be a match. This approach also automatically allows us to resolve overloaded methods, account for generic types and type coercion.

`[PrimalSubstitute(fn)]` and `[PrimalSubstituteOf(fn)]`

In some cases, we face the opposite problem that inspired custom derivatives. That is, we want the compiler to auto-synthesize the derivative from the function body, but there is no function body to translate.
This frequently occurs with hardware intrinsic operations that are lowered into special op-codes that map to hardware units, such as texture sampling & interpolation operations.
However, these operations do have reference 'software' implementations which can be used to produce the derivative.

To allow user code to use the fast hardware intrinsics for the primal pass, but use synthesized derivatives for the derivative pass, we provide decorators [PrimalSubstitute(ref-fn)] and [PrimalSubstituteOf(orig-fn)] (AST Node: PrimalSubstituteAttribute/PrimalSubstituteOfAttribute, IR: OpPrimalSubstituteDecoration), that can be used to provide a reference implementation for the auto-diff pass.

Example:

c
[PrimalSubstitute(sampleTexture_ref)]
float sampleTexture(TexHandle2D tex, float2 uv)
{
    // Hardware intrinsics
}

float sampleTexture_ref(TexHandle2D tex, float2 uv)
{
    // Reference SW implementation.
}

void sampleTexture_bwd(TexHandle2D tex, inout DifferentialPair<float2> dp_uv, float dOut)
{
    // Backward derivate code synthesized using the reference implementation.
}

The implementation of [PrimalSubstitute(fn)] is relatively straightforward. When the transcribers are asked to synthesize a derivative of a function, they check for a OpPrimalSubstituteDecoration, and swap the current function out for the substitute function before proceeding with derivative synthesis.


File: docs/design/autodiff/ir-overview.md

This documentation is intended for Slang contributors and is written from a compiler engineering point of view. For Slang users, see the user-guide at this link: https://shader-slang.com/slang/user-guide/autodiff.html

Overview of Automatic Differentiation's IR Passes

In this document we will detail how Slang's auto-diff passes generate valid forward-mode and reverse-mode derivative functions. Refer to Basics for a review of the two derivative propagation methods and their mathematical connotations & Types for a review of how types are handled under differentiation.

Auto-Diff Pass Invocation

Note that without an explicit auto-diff instruction (fwd_diff(fn) or bwd_diff(fn)) from the user present anywhere in the code, none of the auto-diff passes will do anything.

Auto-diff processing operates on a function-by-function basis. Most of the logic is contained in AutoDiffPass::processReferencedFunctions. Here is a high-level workflow:

  1. Scanning reachable insts in the module looking for IRForwardDifferentiate or IRBackwardDifferentiate operations. These instructions are added onto a work-list. The subject of a differentiate inst may be a plain function (IRFunc), a specialize inst (IRSpecialize(a : IRGeneric, ...)) in case of a generic method, or a lookup inst (IRLookupWitness(a : IRWitnessTableType)) when differentiating a method of an interface.

  2. Dispatch each differentiation request through the appropriate 'transcriber' class. A transcriber (implements AutodiffTranscriberBase) is responsible for accepting a differentiation request and resolving it by replacing it with a generated function or a call to an already existing function that computes its derivative.

  3. Once all currently available derivative insts have been dispatched, the follow-up work-list is checked for more transcription requests. This is a global list that all transcribers can add more follow-up work to. As an example, differentiating a function that calls another function will generate a follow-up task for this inner function, even though the latter never appears directly in a IRForwardDifferentiate or IRBackwardDifferentiate inst.
    At this step, there are 2 other variants that can appear IRBackwardDifferentiatePrimal and IRBackwardDifferentiatePropagate (though these can't be invoked by the user directly).

  4. This process from (1.) is run in a loop. This is because we can have nested differentiation requests such as IRForwardDifferentiate(IRBackwardDifferentiate(a : IRFuncType)). The inner request is processed in the first pass, and the outer request gets processed in the next pass.

Auto-Diff Passes for `IRForwardDifferentiate`

For forward-mode derivatives, we only require a single pass implemented wholly in ForwardDiffTranscriber. This implements the linearization algorithm, which roughly follows this logic:

  1. Create a clone of the original function

  2. Perform pre-autodiff transformations, the most
    a. Temp-Var-For-Mutable-Params Using IRVar to load from
    b. Linkage-Removal: This is simply so the cloned function can be eliminated by DCE after auto-diff is complete
    c. Force-Inline: Inline all __unsafeForceEarlyInline & [ForceInline] functions prior to auto-diff, so their contents can be included in the differentiation pass (even if they aren't actually marked as [Differentiable])

  3. Create a new blank function for the fwd-mode func (usually named s_fwd_<ORIGINAL-FUNC-NAME>) with the function type derived by transforming the original function type (See Types for more information).

  4. Create new blocks into the new fwd-func for each block in the orig-func.

  5. Go through instructions in each block and dispatch to the appropriate generator function to emit the derivative logic into the corresponding block in the fwd-func. Each generator method is responsible for cloning in the original instruction as well as emitting a corresponding derivative instruction. mapPrimalInst and mapDifferentialInst are used to keep track of prior results so that operands for new instructions can be looked up.

The generator for each instruction computes the forward-mode derivative of each instruction using the Jacobian method that is detailed in Basics#Jacobian-Method. Since forward-mode derivatives can be composed in the same order as the original instructions, our generation process goes through instructions in each block in the order that they appear, creating differential insts which act as operands for future insts.

Here's an example of this in IR-form

text
/* Detailed source-code truncated for AI context efficiency. */

Multiple Differential Insts:
In the above example, the derivative of each inst was a single inst. This is not always the case.
For instance, OpMul %a %b translates to three insts:

rust
%1 = OpMul %a_diff %b_primal : %float
%2 = OpMul %a_primal %b_diff : %float
%3 = OpAdd %1 %2 : %float

Combined Primal & Differential Insts:
In some cases, there is not need to clone in the primal inst since both the primal and differential can be computed in a single inst. An example is IRCall, whose deriavative only needs a single call (though it needs plenty of insts to pair and unpair arguments)

rust
// Original inst
%1 = OpCall %func %a %b : %float

//
// Upon differentiation:

// Pack args into pairs
%a_pair = OpDifferentialPairMakePair %a_primal %a_diff : %dpfloat
%b_pair = OpDifferentialPairMakePair %b_primal %b_diff : %dpfloat

// Call into fwd-mode deriv which computes *both* primal and differential
// values.
//
%func_fwd = OpForwardDifferentiate %func : %functype_fwd
%1_pair = OpCall %func_fwd %a_pair %b_pair : %float

// Split into primal and differential so they can be used for future insts.
%1_primal = OpDifferentialPairGetPrimal %1_pair : %float
%1_diff = OpDifferentialPairGetDifferential %1_pair : %float

Phi Arguments

Block arguments are handled the same way as function arguments (which in the Slang IR, are also simply block arguments of the first block), and are converted into pair type arguments, with OpDifferentialPairGetPrimal and OpDifferentialPairGetDifferential insts automatically added to extract the primal and differential parts of each argument.

Auto-Diff Passes for `IRBackwardDifferentiate`

For reverse-mode derivatives, we need several passes that also includes differentiating the forward-mode derivative. Most of this logic is contained in BackwardDiffTranscriberBase::transcribeFuncImpl. These passes are inspired by the paper "You Only Linearize Once: Tangents Transpose to Gradients", which describes this approach in a functional language setting. These passes extend these ideas to work for a general-purpose imperative language structure.

1. Preparation

The reverse-mode derivative generation involves a lot of large scale control-flow manipulation, including a CFG reversal step that aims to construct a method that flows from the end of the function to the beginning in order to compose reverse-mode derivatives.
To avoid having to deal with too many corner cases (and the maintainability issues that come with it), we bring the function to a 'normal form' before running our differentiation steps. This greatly simplifies the logic of the future passes.

Another high-level goal of these transformations is to bring the control-flow graph to a reversible form. That is, we can represent the reverse of control-flow graph using existing Slang constructs (IRIfElse, IRUnconditionalBranch, IRLoop and IRSwitch). This is not necessarily true of any valid Slang IR, so we perform additional transformations.

Note: These transformations are always applied onto a temporary clone of the original function. The original function is never touched so as to not affect its use in non-autodiff contexts.

Specifically we:

  1. Bring the function into single-return form: If there are multiple blocks with return statements (i.e. multiple exit points) in a function, we eliminate this by wrapping the complete function body in a trivial loop (i.e. a single-iteration loop) and replacing existing return statements with breaks (or multi-level breaks) into its break block, which serves as the unique exit point for the function. This pass is currently contained in convertFuncToSingleReturnForm()

  2. Eliminate continue statements: Loop continue statements introduce a reversibility problem. Since the forward loop can have multiple exit point, the reverse loop needs to have multiple entry points. Slang's loops do not support this. So, we eliminate these statements wrapping the body of the loop in another trivial loop (i.e. single-iteration loop) and turning the continue statements into break statements. This also involves writing break statements in the original loop into multi-level breaks.

    Here is an example:

    c
    // Original loop
    for (uint i = 0; i < N; i++)
    {
        if (i > 5)
            continue;
    
        if (i > 9)
            break;
        
        x = x + i;
    }
    
    // After continue-elimination
    outer_for:
    for (uint i = 0; i < N; i++)
    {
        inner_for:
        for (;;)
        {
            if (i > 5)
                break;
    
            if (i > 9)
                break outer_for; // multi-level break
            
            x = x + i;
    
            break;
        }
    }
  3. Eliminate multi-level breaks: Slang supports breaking out to an outer loop. Unfortunately, this operation is hard to reverse since Slang (and shading languages in general) do not support arbitrary goto statements. We eliminate multi-level breaks by assigning each nested loop a nesting index (a constant uint denoting the nesting level). All break statements are rewritten to break out to the immediate next level (i.e. a standard break) with a index parameter denoting the intended break level. This parameter is checked at each level and if the break index does not match the level index, we break again to the immediate upper level. This pass is currently contained in eliminateMultiLevelBreakForFunc

    Continuing the above example, here is the code after multi-level break elimination.

    c
    // After multi-level-break elimination
    uint level = -1;
    for (uint i = 0; i < N; i++)
    {
        for (;;)
        {
            if (i > 5)
            {
                level = 1;
                break;
            }
    
            if (i > 9)
            {
                level = 0;
                break;
            }
            
            x = x + i;
    
            level = 1;
            break;
        }
    
        if (level != 1) // Level check immediately after breaking out of each loop.
            break;
    }
  4. Eliminate break statements (enclosed in normalizeCFG()): Break statements also pose the same problem as continue statements (i.e. multiple exit points require the reverse loop to have multiple entry points, and Slang does not have a primitive for this). We eliminate break statements by introducing a boolean break flag which is set to false to indicate a break instead of using the break statement. Each region is enclosed in a if-else statement that checks the break flag and skips to the end if necessary.

    Break elimination proceeds with the following steps;

    Here is the above example code after break elimination.

    c
    // After break elimination
    
    uint level = -1;
    bool bflag_0 = true; // for outer loop (true => keep-going, false => break)
    
    for (uint i = 0; (i < N) && bflag_0; i++) // Insert flag into the loop condition (&& with the current condition)
    {
        bool bflag_1 = true; // for inner loop (true => keep-going, false => break)
    
        for (;bflag_1;) // Insert flag into the loop condition
        {
            if (i > 5)
            {
                level = 1;
                bflag_1 = false; // break
            }
    
            // Region after any break statement is enclosed in a 
            // if-else check.
            // 
            if (bflag_1)
            {
                if (i > 9)
                {
                    level = 0;
                    bflag_1 = false; // break
                }
    
                // Another if-else enclosure, this time for the second
                // break.
                // 
                if (bflag)
                {
                    x = x + i;
                    level = 1;
                }
    
                bflag_1 = false;
            }
        }
    
        if (level != 1)
        {
            bflag_0 = false;
        }
    }

    Extra evaluation of the condition block: The CFG normalization passes always attempt to preserve the equivalence of the original function while manipulating the control-flow constructs (i.e. ensure that the transformed code always computes the same thing). However, there is one corner-case exception: after break-elimination, the loop condition code can be evaluated 1 additional time, since we don't directly break out of the loop, but go through an extra loop condition check. This becomes important during the checkpointing step, when arrays are allocated to hold loop variables. The array bounds must account for an additional loop iteration to avoid correctness problems.

2. Linearization with Inst-Tagging

This is the same as generating the forward-derivative function, and is in-fact handled in the same way, by invoking ForwardDiffTranscriber. The inst-tagging part of this pass is not necessary for forward-mode auto-diff (simply discarded after the auto-diff pass), but is essential for reverse-mode.

Inst-Tagging: This pass also tags every instruction and block with either IRPrimalInstDecoration, IRDifferentialInstDecoration or IRMixedDifferentialnstDecoration, depending on whether an instruction contains/computes/reads/writes a primal value, a differential value or both.

This assignment is according to the following rules:

  1. The result of .getDifferential() from an inst of IRDifferentialPairType is a differential inst and .getPrimal() is a primal inst NOTE: This does not apply to IRDifferentialPairUserCodeType, all of whose operations yield a primal inst.
  2. Further, any inst which contains a differential inst as an operand AND whose output value may be affected by this operand is a differential inst (e.g. if isDifferentialInst(a) = true then isDifferentialInst( IRMul(a, b) ) = true)
  3. If an inst contains multiple outputs, some of which are differential and the others are primal, then these are mixed-differential insts. E.g. (a value of IRDifferentialPairType contains both a primal and differential value, and similarly a call of the form IRCall(IRForwardDifferentiate(inner_fn))(...) results in a mixed differential type since the primal part is not affect by differential inputs)
  4. All other insts are primal by default.
  5. Blocks are marked differential or primal if they contain ONLY differential or primal insts (respectively). Otherwise they are marked mixed-differential. The vast majority of blocks are mixed-differential.

Correct tag information is critical for the next steps to correctly transform the forward-mode derivative into the reverse-mode derivative function.

Here's the same forward-mode example, but with insts tagged accordingly

rust
OpModule
{
    // Generated function type
    ...

    // Generated function
    ...
        [OpMixedDifferentiaInstDecoration]
        %b_fwd = OpBlock
        {
            // Block params are mixed differentials since they carry both
            // primal and differential values
            // 
            [OpMixedDifferentialInstDecoration]
            %dpa = OpParam : %dpfloat 
            [OpMixedDifferentialInstDecoration]
            %dpb = OpParam : %dpfloat 

            [OpPrimalInstDecoration]
            %a = OpDifferentialPairGetPrimal %dpa : %float

            [OpDifferentialInstDecoration]
            %da = OpDifferentialPairGetDifferential %dpa : %float

            [OpPrimalInstDecoration]
            %b = OpDifferentialPairGetPrimal %dpb : %float

            [OpDifferentialInstDecoration]
            %db = OpDifferentialPairGetDifferential %dpb : %float

            [OpPrimalInstDecoration]
            %1_primal = OpAdd %a %b : %float

            [OpDifferentialInstDecoration]
            %1_diff = OpAdd %da %db : %float

            [OpPrimalInstDecoration]
            %2_primal = OpAdd %1_primal %1_primal : %float

            [OpDifferentialInstDecoration]
            %2_diff = OpAdd %1_diff %1_diff : %float

            // Return both the primal and differential
            [OpMixedDifferentialInstDecoration]
            %2_pair = OpDifferentialPairMakePair %2_primal %2_diff : %dpfloat

            [OpDifferentialInstDecoration]
            OpReturn %2_pair
        }
    ...
}

3. Unzipping

Implemented by DiffUnzipPass, this pass is responsible for separating primal instructions from differential instructions (as denoted by their decorations), by creating a full set of duplicate blocks that start after the last block, i.e. return block (the return statement is removed).

This separation is possible because the computation of a differential inst may include primal operands but a primal inst can never use a differential operand.

The unzipping pass uses the decorations from the linearization step to figure out which instructions need to be moved.

The separation process uses the following high-level logic:

  1. Create two clones of all the blocks in the provided function (one for primal insts, one for differential insts), and hold a mapping between each original (mixed) block to each primal and differential block. The return statement of the current final block is removed.
  2. Process each instruction of each block: instructions marked as primal are moved to the corresponding primal block, instructions marked differential are moved to the corresponding differential block.
  3. Instructions marked mixed need op-specific handling, and so are dispatched to the appropriate splitting function. For instance, block parameters that are holding differential-pair values are split into parameters for holding primal and differential values (the exception is function parameters, which are not affected). Similarly, IRVars, IRTerminatorInsts (control-flow) and IRCalls are all split into multiple insts.
  4. Except for IRReturn, all other control-flow insts are effectively duplicated so that the control-flow between the primal blocks and differential blocks both follow the original blocks' control-flow. The main difference is that PHI arguments are split (primal blocks carry primal values in their PHI arguments, and differential blocks carry diff values) between the two. Note that condition values (i.e. booleans) are used by both the primal and differential control-flow insts. However, since booleans are always primal values, they are always defined in the primal blocks.

Block-Tagging: Blocks are now tagged primal or differential depending on whether they are holding primal or differential insts. This is important for the next step (transposition) to figure out which blocks need to be transposed.

Out-of-Scope Accesses: After unzipping, the resulting IR is often not valid. If the control-flow is straight line (i.e. no branching or loops), the resulting IR is valid. However, if there is control-flow, then instructions can use operands whose definition does not dominate the use. This invalid IR is currently allowed to persist until the end of the auto-diff passes, when the checkpointing step occurs (i.e. Running IR validation will fail in between these steps)

Here is an example of unzipped code:

text
/* Detailed source-code truncated for AI context efficiency. */

4. Transposition

The next step involves converting each differential instruction into its transpose. Effectively, we are re-writing each forward-mode derivative into its reverse-mode equivalent.

Recall from auto-diff basics, that both the forward and reverse mode derivatives can be derived from the Jacobian matrix of any operation. The main difference is whether we multiply the derivatives of the inputs with the Jacobian or multiply the Jacobian with the derivatives w.r.t the outputs. These two operations are the transpose of each other, in that the reverse-mode derivative can be thought of as multiplying with the transpose of the Jacobian.

We perform this transposition on a per-instruction level.

Here is an example of a transposition of a multiplication operation:

rust
[OpPrimalInstDecoration]
%b = OpLoad %var_b // %b is a primal value

[OpDifferentialInstDecoration]
%da = OpLoad %var_da // %da is a differential value

// The operation we want to transpose
[OpDifferentialInstDecoration]
%1d = OpMul %da %b : %float

[OpDifferentialInstDecoration]
OpStore %1d %var_result

This multiplication can be represented as a tiny matrix multiplication between a singleton vector [%da] and singleton matrix [%b].
It's transpose will be the multiplication of the transpose of that matrix (which is the value itself [%b]) with a derivative w.r.t its output %1d, i.e. it becomes %da = OpMul %1d %b. Note that we now have to provide %1d as an input, and receive da was an output.

The resulting code is then:

rust
[OpPrimalInstDecoration]
%b = OpLoad %var_b : %float // primal values are unaffected (at this stage, they are in primal blocks)

// Reverse-mode code: (_rev) appended to all variables & insts to keep them distinct from the fwd-mode code.
[OpDifferentialInstDecoration]
%1d_rev = OpLoad %var_result_rev : %float

// The operation we want to transpose
[OpDifferentialInstDecoration]
%da_rev = OpMul %1d_rev %b : %float

[OpDifferentialInstDecoration]
OpStore %da_rev %var_da_rev

Notice that the three differential instructions are effectively run backwards and transposed. Loads become stores,
the OpMul is transposed into another OpMul, and stores become loads. This backwards transposition is because the differential outputs become differential inputs, and thus, we need to process the future instructions first so that the new operands are defined before bring used for the new instruction.

This reverse order of operations also applies to control-flow. The rule of thumb is: if the forward-mode pass takes a particular path through the code, for a given set of primal values, the reverse-mode must "re-trace" the same path through the code, but in reverse by starting at the end.

We synthesize a CFG that satisfies this property through the following steps:

  1. Clone the provided unzipped forward-mode function (and all blocks + instructions) to serve as the reverse-mode function.

  2. Remove all differential blocks and create a set of corresponding reverse-mode blocks for each differential block removed (primal blocks are simply left alone), while holding a map between corresponding blocks. Initially, they are empty.

  3. Using the provided unzipped forward-mode function as a reference, process each differential block by walking each instruction from the last (terminator) inst, and dispatching to the appropriate op-specific transposeXYZ() method to emit the appropriate transposed instructions into the corresponding reverse-mode block.

    There are several concerns that must be taken care of:

    1. Multiple Derivaive Outputs: Unlike forward-mode auto-diff, where an inst producing a single value, would only need a single derivative (corresponding to that value), reverse-mode auto-diff can produce multiple derivatives from an inst. For instance %dc = IRAdd(%da, %db) produces two derivatives: %da_rev = %dc_rev and %db_rev = %dc_rev. Thus, the transposeXYZ() implementation for any instruction can return a set of derivative insts for each relevant input differential value.

    2. Insts Used in Multiple Places (Derivative Accumulation): If an inst is used in multiple places, and receives a reverse-mode derivative from several of those places, these results need to be added up to get the correct derivative.

      Consider this forward-mode example

      rust
      [OpDifferentialInstDecoration]
      %db = OpAdd %da, %da : %float
      
      [OpDifferentialInstDecoration]
      %dc = OpAdd %db, %da : %float

      It's reverse-mode derivative will look like this:

      rust
      %db_rev = %dc_rev // %db only has one differential since it only consumed in one place.
      
      // reverse-mode differential for %da from trnaposing the first instruction
      [OpDifferentialInstDecoration]
      %da_rev_1 = OpAdd %db_rev, %db_rev : %float
      
      // reverse-mode differential for %da from transposing the second instruction
      [OpDifferentialInstDecoration]
      %da_rev_2 = %dc_rev
      
      // add them together to get the final derivative for %da
      [OpDifferentialInstDecoration]
      %da_rev = OpAdd %da_rev_1 %da_rev_2 : %float

      Derivative accumulation is achieved through two ways:

      Within a block, we keep a list all the reverse derivative insts for each inst and only materialize the total derivative when it is required as an operand. This is the most efficient way to do this, because we can apply certain optimizations for composite types (derivative of an array element, vector element, struct field, etc..).

      Across blocks, we use an accumulator variable that is inserted into a top-level block in the function, and add to this variable whenever a transposition operation generates a new inst. This can sometimes produce sub-optimal code for aggregate/large data types, but at the moment, the accumulator method is necessary because insts can receive derivatives from conditionally executed blocks.

      While this example uses OpAdd to demonstrate accumulation, in practice, we use the derivative type system (See Types for more) to look up the derivative addition function (dadd) to add two values of an arbitrary differential type. In practice, the OpAdd is replaced by OpCall %float_dadd %da_rev1 %da_rev_2. Similarly, for accumulator variables, we must initialize them to zero for the accumulation to work correctly, and we lookup the dzero interface method to initialize it in a type-specific way.

    3. Deferred Materialization for Derivatives of Composite Types:
      Non-primitive types, such as vectors, arrays, structs, etc. whose elements are used in several places in the forward-mode code, can result in sub-optimal reverse-mode code. Here is an example (in Slang source-style):

      c
      float f_fwd(DifferentialPair<float3> input)
      {
          float3 dinput = input.getDifferential();
          float a = dinput.x + dinput.y;
          float b = a + dinput.z;
      
          return b;
      }
      
      // Transposed code (naively, without deferred materialization)
      void f_rev(inout DifferentialPair<float3> input, float d_output)
      {
          // transpose of (return b;)
          float db_rev = d_output;
          
          // transpose of (float b = a + dinput.z)
          float da_rev = db_rev;
          float3 dinput_rev_1 = float3(0.f, 0.f, da_rev);
      
          // transpose of (float a = dinput.x + dinput.y)
          float3 dinput_rev_2 = float3(0.f, da_rev, 0.f);
          float3 dinput_rev_3 = float3(da_rev, 0.f, 0.f);
      
          // Accumulate [dinput_rev_1, dinput_rev_2, dinput_rev_3]
          float3 dinput = dinput_rev_1 + dinput_rev_2 + dinput_rev_3
      
          input = DifferentialPair<float3>(
              input.getPrimal(),
              dinput);
      }

      Note that, this approach to inst-by-inst transposition can use a lot more stack space than is necessary (dinput_rev_1, dinput_rev_2 and dinput_rev_3 all only have a single non-0 entry). This is a known complexity issue with naive inst-by-inst transposition: hypothetically, an size-$N$ vector/array would end up allocating $O(N^2)$ memory even if only $N$ elements are non-0.
      In our Slang implementation, we circumvent this (to an extent) by deferring materialization. Rather than create each component dinput_rev_i as soon as we see an inst use, we hold the derivative with a special flavor value for lookups (say Swizzle or GetElement). When the total value dinput_rev is necessary, we process components of each flavor type at once and create a single derivative from all the components.

      Here is the same example, with deferred materialization:

      c
      // Transposed code (naively, without deferred materialization)
      void f_rev(inout DifferentialPair<float3> input, float d_output)
      {
          // transpose of (return b;)
          float db_rev = d_output;
          
          // transpose of (float b = a + dinput.z), hold {flavor=Swizzle, component=.z, derivInst=db_rev} in list.
          float da_rev = db_rev;
      
          // transpose of (float a = dinput.x + dinput.y), 
          // hold {flavor=Swizzle, component=.x, derivInst=da_rev} and {flavor=Swizzle, component=.y, derivInst=da_rev} in list.
      
          // Materialize when required (for constructing return pair)
          float3 dinput = float3(db_rev, da_rev, da_rev);
      
          input = DifferentialPair<float3>(
              input.getPrimal(),
              dinput);
      }

      Note that this only really works for accumulation within a single block/control-flow region. For across regions, we still have to materialize when we exit a region, so this memory problem can still manifest for control-flow heavy functions, where each region must allocate enough space for its contribution to the full derivative, even if only a small subset is non-0.

c
float a[10] = /*...*/;
for (int i = 0; i < 10; i++)
{
    a[i] = f(a[i]);
}
c
// Entry block
%t = OpBlock
{
    IRLoop %c %br %c 0
}

// Condition
%c = OpBlock
{
    %i = OpParam : %float
    %a = OpParam : %Array(%float, 10)
    
    %2 = OpLesser(%i, 10) : %bool

    %OpIfElse(%2, %b, %br, %br)
}

// Loop body.
%b = OpBlock 
{
    %a_i = OpGetElement(%a, %i) : %float
    %f_a_i = OpCall(f, %a_i) : %float

    %a_next = OpUpdateElement(%a, %i, %f_a_i) : %Array(%float, 10)

    %i_next = OpAdd(%i, 1)

    OpUnconditionalBranch(%c, %i_next, %a_next)
}

// Break block
%br = OpBlock
{
    //...
}

After AD passes, this results in the following code:

text
/* Detailed source-code truncated for AI context efficiency. */
  1. Construct the reverse control-flow (reveseCFGRegion()) by going through the reference forward-mode blocks, and cloning the control-flow onto the reverse-mode blocks, but in reverse. This is achieved by running reverseCFGRegion() recursively on each sub-region, where a region is defined as a set of blocks with a single entry block and a single exit block. This definition of a region only works because we normalized the CFG into this form.

    The reversal logic follows these general rules:

    1. Unconditional Branch: For an unconditional branch from A->B we simply have to map the reverse version of B with that of A. i.e. rev[B] -> rev[A]
    2. If-Else: For an if-else of the form A->[true = T->...->T_last->M, false = F->...->F_last->M], we construct rev[M]->[true = rev[T_last]->...->rev[T_last]->rev[A], false = rev[F_last]->...->rev[F]->rev[A]]. That is, we reverse each sub-region, and start from the merge block and end at the split block.
      Note that we need to identify T_last and F_last i.e. the last two blocks in the true and false regions. We make the last block in the region an additional return value of reverseCFGRegion(), so that when reversing the true and false sub-regions, we also get the relevant last block as an additional output. Also note that additional empty blocks may be inserted to carry derivatives of the phi arguments, but this does not alter the control-flow.
    3. Switch-case: Proceeds in exactly the same way as if-else reversal, but with multiple cases instead of just 2.
    4. Loop: After normalization, all (non-trivial) loops are of the form: A->C->[true = T->...->T_last->C, false=B->...->M]. We reverse this loop into rev[M]->...rev[B]->rev[C]->[true=rev[T_last]->...->rev[T]->rev[C], false=rev[A]]. The actual reversal logic also handles some corner cases by inserting additional blank blocks to avoid situations where regions may share the same merge block.

    Finally, we process the first and last blocks (entry and return blocks) by inserting a void return (reverse-mode derivative functions are always of void result type)

At this stage, the reverse-mode generation is almost complete. The control-flow and the derivative logic is present, but we still have to resolve out-of-scope accesses from the new differential blocks into the primal block.

5. Checkpointing/Recomputation (also called 'primal-hoisting')

This step legalizes the out-of-scope accesses of primal insts from within differential blocks. This is to prepare us for the next step (i.e. extraction) that splits the function into two by moving the primal blocks into a separate primal-context-generator function, and the differential blocks into the backward-propagation function.

Before we can perform this extraction, we must find any primal values being used in differential blocks and handle them in one of two ways:
Store (put the values in a static struct) or Recompute (clone the necessary instructions to recompute when necessary). We first classify all necessary instructions into one of the two buckets before processing each use accordingly.

  1. Classify uses into each set: Note that rather than proceeding on an inst-by-inst basis, we classify uses of insts. The same inst can be used in several places, and we may decide to store one use and recompute another (in some cases, this could be the optimal result).
    The classification process uses a work-list approach that roughly looks like the following:

    1. Add all uses of primal insts in an inst within a differential block to the work list. This is our initial set of uses that require classification.
    2. Query the active policy object (which for now is hardcoded) to obtain the classification based on heuristics & user decorations (Specifically [PreferRecompute] and [PreferCheckpoint] decorations influence the classification policy)
    3. For uses that should be recomputed, we have to now make the same decision one their operands, in order to make them available for the recomputation insts. Thus, their operands are added to the work list.
    4. For uses that should be stored, there is no need to consider their operands, since the computed value will be explicitly stored and loaded later.
    5. Once the worklist is empty, go over all the uses and their classifications, and convert them into a list of insts that should be stored or recomputed. Note that if an inst has uses with both classifications, then it can appear in both lists.
  2. Process 'Store' (i.e. checkpoint) insts: Store them into a single variable (of a struct type that is synthesized as necessary), and then loaded from in the differential blocks. This allows us to simply turn this variable into an output parameter from the context function and an input parameter for the backprop function.
    When storing values this way, we must consider that instructions within loops can have different values each iteration. Thus, we must use an array to store each value, and this array's size must be statically known since we wish to synthesize a static struct type to hold all the stored values. Thus, we enforce the requirement of a [MaxIters(N)] decoration and attempt to infer a loop iteration limit if one is not provided.

    Here's an example of a case where we decide to checkpoint all relevant uses:

    text

/* Detailed source-code truncated for AI context efficiency. */

text
Another example with a function `g` that does contain loops:

    ```
/* Detailed source-code truncated for AI context efficiency. */
text
**Indexed Region Processing:** In order to be able to allocate the right array and use the right indices, we need information about which blocks are part of which loop (and loops can be nested, so blocks can be part of multiple loops). To do this, we run a pre-processing step that maps all blocks to all relevant loop regions, the corresponding index variables and the inferred iteration limits (maximum times a loop can run). Note that if an instruction appears in a nested block, we create a multi-dimensional array and use multiple indices.

**Loop State Variables:** Certain variables cannot be classified as recompute. Major examples are loop state variables which are defined as variables that are read from and written to within the loop. In practice, they appear as phi-variables on the first loop block after SSA simplification. Their uses _must_ be classified as 'store', because recomputing them requires duplicating the primal loop within the differential loop. This is because the differential loop runs backwards so the state of a primal variable at loop index $N$ cannot be recomputed when the loop is running backwards ($N+1 \to N \to N-1$), and involves running the primal loop up to $N$ times within the current iteration of the differential loop. In terms of complexity, this turns an $O(N)$ loop into an $O(N^2)$ loop, and so we disallow this.
It is possible that the resulting $O(N^2)$ loop may end up being faster in practice due to reduced memory requirements, but we currently lack the infrastructure to robustly allow such loop duplication while keeping the user informed of the potentially drastic complexity issues.
  1. Process 'Recompute' insts: Insert a copy of the primal instruction into a corresponding 'recomputation' block that is inserted into the differential control-flow so that it dominates the use-site.

    Insertion of Recompute Blocks: In order to accommodate recomputation, we first preprocess the function, by going through each breakable (i.e. loop) region in the differential blocks, looking up the corresponding primal region and cloning all the primal blocks into the beginning of the differential region. Note that this cloning process does not actually clone the instructions within each block, only the control-flow (i.e. terminator) insts. This way, there is a 1:1 mapping between the primal blocks and the newly created recompute blocks, This way, if we decide to 'recompute' an instruction, we can simply clone it into the corresponding recompute block, and we have a guarantee that the definition and use-site are within the same loop scope, and that the definition comes before the use.

    Legalizing Accesses from Branches: Our per-loop-region recompute blocks ensure that the recomputed inst is always within the same region as its uses, but it can still be out-of-scope if it is defined within a branch (i.e. if-else). We therefore still run a light-weight hoisting pass that detects these uses, inserts an IRVar at the immediate dominator of the def and use, and inserts loads and stores accordingly. Since they occur within the same loop region, there is no need to worry about arrays/indices (unlike the 'store' case).

    Marking Recompute Blocks: These blocks are marked with OpRecomputeBlockDecoration to identify them as containing primal instructions, even though they are within differential regions. This helps us remove any unused blocks if none of the instructions end up being recomputed.

    Here is an example of recomputation demonstrated in Slang source-style (although this takes place in IR-form)

    text

/* Detailed source-code truncated for AI context efficiency. */

text
For completeness, here is another example of a function `g` which contains a loop to demonstrate how recomputation works when there are
    multiple loop regions.

    ```
/* Detailed source-code truncated for AI context efficiency. */

6. Extraction

The final step involves splitting the function immediately after the primal block to create two functions: a primal context function that computes the primal value normally, but also outputs a context object with relevant intermediate values, and a backward propagation function that computes the backward derivative and consumes this context object for the required intermediate values.

The first 5 steps have set us up for this final step, so it is not particularly complex. We follow this high-level logic:

  1. Create an empty function for the primal context function. The type of this function is the same as the primal function, but with an additional out parameter for the intermediate context, whose type is undecided at this stage. We use a temporary function-specific type called OpBackwardDerivativeIntermediateContextType(func) as a placeholder.

  2. Move primal blocks to the primal context function. Re-create the return inst (the return value is temporarily remembered using a decoration during the rest of the AD process). Also, the first block (reserved for function parameters) is also duplicated and processed to have primal parameters in the primal function and pair parameters in the differential function.

  3. Lower all OpBackwardDerivativeIntermediateContextType types into concrete struct types by creating a field for each 'stored' inst from Step 5. This lowering process happens at the end of the current AD pass after all relevant methods have completed Step 5. We need Step 5 (hoisting) to be complete for all relevant methods because the context struct for a given function can include context structs of other functions that are called from it. Our context-type lowering therefore proceeds recursively by lowering the context for inner functions as necessary. The lowering process also removes the temporary vars that were created to hold the store insts, and replaces them with a stores and loads from the context struct.

    Recursive Functions are Disallowed: Since we lower all intermediate types into a static struct type, recursive calls cannot currently be supported from differentiable functions. The context struct for a method may include itself, creating an impossible scenario.

Here is one of the examples above (g) after checkpointing:

text
/* Detailed source-code truncated for AI context efficiency. */

Having separate methods for the primal and backward passes is necessary when reverse-mode differentiating a method that calls out to other differentiable functions.
Here is an example of differentiating a method that calls out to multiple methods, to get an idea for why we need the primal context method to be separate

c
float outer(float x)
{
    float y = f(x);
    float z = g(y);
    float w = h(z);

    return w;
}

// It's complete reverse mode derivative looks like the following:
void outer_rev(DifferentialPair<float> dpx, float d_output)
{
    // Compute the primal values in the forward direction, while producing relevant context. 
    f_Intermediates f_ctx;
    g_Intermediates g_ctx;
    h_Intermediates h_ctx;

    float y = s_primal_ctx_f(x, f_ctx);
    float z = s_primal_ctx_g(y, g_ctx);
    float w = s_primal_ctx_h(z, h_ctx);

    // Note that at this point, we are holding intermediate context variables for f, g and h.

    // Consume the context while evaluating the propagating the derivatives backwards.
    DifferentialPair<float> dpz = {z, 0.f};
    s_bwd_h(dpz, d_output, h_ctx);

    DifferentialPair<float> dpy = {y, 0.f};
    s_bwd_g(dpy, dpz.getDifferential(), g_ctx);

    DifferentialPair<float> _dpx = {x, 0.f};
    s_bwd_f(dpx, dpy.getDifferential(), f_ctx);

    dpx = _dpx;
}

File: docs/design/autodiff/types.md

This documentation is intended for Slang contributors and is written from a compiler engineering point of view. For Slang users, see the user-guide at this link: https://shader-slang.com/slang/user-guide/autodiff.html

Before diving into this document, please review the document on Basics for the fundamentals of automatic differentiation.

Components of the Type System

Here we detail the main components of the type system: the IDifferentiable interface to define differentiable types, the DifferentialPair<T> type to carry a primal and corresponding differential in a single type.
We also detail how auto-diff operators are type-checked (the higher-order function checking system), how the no_diff decoration can be used to avoid differentiation through attributed types, and the derivative data flow analysis that warns the the user of unintentionally stopping derivatives.

`interface IDifferentiable`

Defined in core.meta.slang, IDifferentiable forms the basis for denoting differentiable types, both within the core module, and otherwise.
The definition of IDifferentiable is designed to encapsulate the following 4 items:

  1. Differential: The type of the differential value of the conforming type. This allows custom data-structures to be defined to carry the differential values, which may be optimized for space instead of relying solely on compiler synthesis/

Since the computation of derivatives is inherently linear, we only need access to a few operations. These are:

  1. dadd(Differential, Differential) -> Differential: Addition of two values of the differential type. It's implementation must be associative and commutative, or the resulting derivative code may be incorrect.
  2. dzero() -> Differential: Additive identity (i.e. the zero or empty value) that can be used to initialize variables during gradient aggregation
  3. dmul<S:__BuiltinRealType>(S, Differential): Scalar multiplication of a real number with the differential type. It's implementation must be distributive over differential addition (dadd).

Points 2, 3 & 4 are derived from the concept of vector spaces. The derivative values of any Slang function always form a vector space (https://en.wikipedia.org/wiki/Vector_space).

Derivative member associations

In certain scenarios, the compiler needs information on how the fields in the original type map to the differential type. Particularly, this is a problem when differentiate the implicit construction of a struct through braces (i.e. {}), represented by kIROp_MakeStruct. We provide the decorator [DerivativeMember(DifferentialTypeName.fieldName)] (ASTNode: DerivativeMemberAttribute, IR: kIROp_DerivativeMemberDecoration) to explicitly mark these associations.
Example

c
struct MyType : IDifferentiable
{
    typealias Differential = MyDiffType;
    float a;

    [DerivativeMember(MyDiffType.db)]
    float b;

    /* ... */
};

struct MyDiffType
{
    float db;
};

Automatic Synthesis of `IDifferentible` Conformances for Aggregate Types

It can be tedious to expect users to hand-write the associated Differential type, the corresponding mappings and interface methods for every user-defined struct type. For aggregate types, these are trivial to construct by analysing which of their components conform to IDifferentiable.
The synthesis proceeds in roughly the following fashion:

  1. IDifferentiable's components are tagged with special decorator __builtin_requirement(unique_integer_id) which carries an enum value from BuiltinRequirementKind.
  2. When checking that types conform to their interfaces, if a user-provided definition does not satisfy a requirement with a built-in tag, we perform synthesis by dispatching to trySynthesizeRequirementWitness.
  3. For user-defined types, Differential types are synthesized during conformance-checking through trySynthesizeDifferentialAssociatedTypeRequirementWitness by checking if each constituent type conforms to IDifferentiable, looking up the corresponding Differential type, and constructing a new aggregate type from these differential types. Note that since it is possible that a Differential type of a constituent member has not yet been synthesized, we have additional logic in the lookup system (trySynthesizeRequirementWitness) that synthesizes a temporary empty type with a ToBeSynthesizedModifier, so that the fields can be filled in later, when the member type undergoes conformance checking.
  4. For user-defined types, Differential methods (dadd, dzero and dmul) are synthesized in trySynthesizeDifferentialMethodRequirementWitness by utilizing the Differential member and its [DifferentialMember] decorations to determine which fields need to be considered and the base type to use for each field. There are two synthesis patterns. The fully-inductive pattern is used for dadd and dzero which works by calling dadd and dzero respectively on the individual fields of the Differential type under consideration.
    Example:
c
// Synthesized from "struct T {FT1 field1; FT2 field2;}"
T.Differential dadd(T.Differential a, T.Differential b)
{
    return Differential(
        FT1.dadd(a.field1, b.field1),
        FT2.dadd(a.field2, b.field2),
    )
}

On the other hand, dmul uses the fixed-first arg pattern since the first argument is a common scalar, and proceeds inductively on all the other args.
Example:

c
// Synthesized from "struct T {FT1 field1; FT2 field2;}"
T.Differential dmul<S:__BuiltinRealType>(S s, T.Differential a)
{
    return Differential(
        FT1<S>.dmul(s, a.field1),
        FT2<S>.dmul(s, a.field2),
    )
}
  1. During auto-diff, the compiler can sometimes synthesize new aggregate types. The most common case is the intermediate context type (kIROp_BackwardDerivativeIntermediateContextType), which is lowered into a standard struct once the auto-diff pass is complete. It is important to synthesize the IDifferentiable conformance for such types since they may be further differentiated (through higher-order differentiation). This implementation is contained in fillDifferentialTypeImplementationForStruct(...) and is roughly analogous to the AST-side synthesis.

Differentiable Type Dictionaries

During auto-diff, the IR passes frequently need to perform lookups to check if an IRType is differentiable, and retrieve references to the corresponding IDifferentiable methods. These lookups also need to work on generic parameters (that are defined inside generic containers), and existential types that are interface-typed parameters.

To accommodate this range of different type systems, Slang uses a type dictionary system that associates a dictionary of relevant types with each function. This works in the following way:

  1. When CheckTerm() is called on an expression within a function that is marked differentiable ([Differentiable]), we check if the resolved type conforms to IDifferentiable. If so, we add this type to the dictionary along with the witness to its differentiability. The dictionary is currently located on DifferentiableAttribute that corresponds to the [Differentiable] modifier.

  2. When lowering to IR, we create a DifferentiableTypeDictionaryDecoration which holds the IR versions of all the types in the dictionary as well as a reference to their IDifferentiable witness tables.

  3. When synthesizing the derivative code, all the transcriber passes use DifferentiableTypeConformanceContext::setFunc() to load the type dictionary. DifferentiableTypeConformanceContext then provides convenience functions to lookup differentiable types, appropriate IDifferentiable methods, and construct appropriate DifferentialPair<T>s.

Looking up Differential Info on _Generic_ types

Generically defined types are also lowered into the differentiable type dictionary, but rather than having a concrete witness table, the witness table is itself a parameter. When auto-diff passes need to find the differential type or place a call to the IDifferentiable methods, this is turned into a lookup on the witness table parameter (i.e. Lookup(, <WitnessTableParameter>)). Note that these lookups instructions are inserted into the generic parent container rather than the inner most function.
Example:

c
T myFunc<T:IDifferentiable>(T a)
{
    return a * a;
}

// Reverse-mode differentiated version
void bwd_myFunc<T:IDifferentiable>(
    inout DifferentialPair<T> dpa,
    T.Differential dOut) // T.Differential is Lookup('Differential', T_Witness_Table)
{
    T.Differential da = T.dzero(); // T.dzero is Lookup('dzero', T_Witness_Table)

    da = T.dadd(dpa.p * dOut, da); // T.dadd is Lookup('dadd', T_Witness_Table)
    da = T.dadd(dpa.p * dOut, da);

    dpa = diffPair(dpa.p, da);
}

Looking up Differential Info on _Existential_ types

Existential types are interface-typed values, where there are multiple possible implementations at run-time. The existential type carries information about the concrete type at run-time and is effectively a 'tagged union' of all possible types.

Differential type of an Existential

The differential type of an existential type is tricky to define since our type system's only restriction on the .Differential type is that it also conforms to IDifferentiable. The differential type of any interface IInterface : IDifferentiable is therefore the interface type IDifferentiable. This is problematic since Slang generally requires a static anyValueSize that must be a strict upper bound on the sizes of all conforming types (since this size is used to allocate space for the union). Since IDifferentiable is defined in the core module core.meta.slang and can be used by the user, it is impossible to define a reliable bound.
We instead provide a new any-value-size inference pass (slang-ir-any-value-inference.h/slang-ir-any-value-inference.cpp) that assembles a list of types that conform to each interface in the final linked IR and determines a relevant upper bound. This allows us to ignore types that conform to IDifferentiable but aren't used in the final IR, and generate a tighter upper bound.

Future work:
This approach, while functional, creates a locality problem since the size of IDifferentiable is the max of all types that conform to IDifferentiable in visible modules, even though we only care about the subset of types that appear as T.Differential for T : IInterface. The reason for this problem is that upon performing an associated type lookup, the Slang IR drops all information about the base interface that the lookup starts from and only considers the constraint interface (in this case Differential : IDifferentiable).
There are several ways to resolve this issue, including (i) a static analysis pass that determines the possible set of types at each use location and propagates them to determine a narrower set of types, or (ii) generic (or 'parameterized') interfaces, such as IDifferentiable<T> where each version can have a different set of conforming types.

Example:

c
interface IInterface : IDifferentiable
{
    [Differentiable]
    This foo(float val);

    [Differentiable]
    float bar();
};

float myFunc(IInterface obj, float a)
{
    IInterface k = obj.foo(a);
    return k.bar();
}

// Reverse-mode differentiated version (in pseudo-code corresponding to IR, some of these will get lowered further)
void bwd_myFunc(
    inout DifferentialPair dpobj,
    inout DifferentialPair<float> dpa,
    float.Differential dOut) // T.Differential is Lookup('Differential', T_Witness_Table)
{
    // Primal pass..
    IInterface obj = dpobj.p;
    IInterface k = obj.foo(a);
    // .....

    // Backward pass
    DifferentialPair dpk = diffPair(k);
    bwd_bar(dpk, dOut);
    IDifferentiable dk = dpk.d; // Differential of `IInterface` is `IDifferentiable`

    DifferentialPair dp = diffPair(dpobj.p);
    bwd_foo(dpobj, dpa, dk);
}

Looking up `dadd()` and `dzero()` on Existential Types

There are two distinct cases for lookup on an existential type. The more common case is the closed-box existential type represented simply by an interface. Every value of this type contains a type identifier & a witness table identifier along with the value itself. The less common case is when the function calls are performed directly on the value after being cast to the concrete type.

dzero() for "closed" Existential type: The NullDifferential Type
For concrete and even generic types, we can initialize a derivative accumulator variable by calling the appropriate Type.dzero() method. This is unfortunately not possible when initializing an existential differential (which is currently of type IDifferentiable), since we must also initialize the type-id of this existential to one of the implementations, but we do not know which one yet since that is a run-time value that only becomes known after the first differential value is generated.

To get around this issue, we declare a special type called NullDifferential that acts as a "none type" for any IDifferentiable existential object.

dadd() for "closed" Existential types: __existential_dadd
We cannot directly use dadd() on two existential differentials of type IDifferentiable because we must handle the case where one of them is of type NullDifferential and dadd() is only defined for differentials of the same type.
We handle this currently by synthesizing a special method called __existential_dadd (getOrCreateExistentialDAddMethod in slang-ir-autodiff.cpp) that performs a run-time type-id check to see if one of the operand is of type NullDifferential and returns the other operand if so. If both are non-null, we dispatch to the appropriate dadd for the concrete type.

dadd() and dzero() for "open" Existential types
If we are dealing with values of the concrete type (i.e. the opened value obtained through ExtractExistentialValue(ExistentialParam)). Then we can perform lookups in the same way we do for generic type. All existential parameters come with a witness table. We insert instructions to extract this witness table and perform lookups accordingly. That is, for dadd(), we use Lookup('dadd', ExtractExistentialWitnessTable(ExistentialParam)) and place a call to the result.

`struct DifferentialPair`

The second major component is DifferentialPair<T:IDifferentiable> that represents a pair of a primal value and its corresponding differential value.
The differential pair is primarily used for passing & receiving derivatives from the synthesized derivative methods, as well as for block parameters on the IR-side.
Both fwd_diff(fn) and bwd_diff(fn) act as function-to-function transformations, and so the Slang front-end translates the type of fn to its derivative version so the arguments can be type checked.

Pair type lowering.

The differential pair type is a special type throughout the AST and IR passes (AST Node: DifferentialPairType, IR: kIROp_DifferentialPairType) because of its use in front-end semantic checking and when synthesizing the derivative code for the functions. Once the auto-diff passes are complete, the pair types are lowering into simple structs so they can be easily emitted (DiffPairLoweringPass in slang-ir-autodiff-pairs.cpp).
We also define additional instructions for pair construction (kIROp_MakeDifferentialPair) and extraction (kIROp_DifferentialPairGetDifferential & kIROp_DifferentialPairGetPrimal) which are lowered into struct construction and field accessors, respectively.

"User-code" Differential Pairs

Just as we use special IR codes for differential pairs because they have special handling in the IR passes, sometimes differential pairs should be treated as regular struct types during the auto-diff passes.
This happens primarily during higher-order differentiation when the user wishes to differentiate the same code multiple times.
Slang's auto-diff approaches this by rewriting all the relevant differential pairs into 'irrelevant' differential pairs (kIROp_DifferentialPairUserCode) and 'irrelevant' accessors (kIROp_DifferentialPairGetDifferentialUserCode, kIROp_DifferentialPairGetPrimalUserCode) at the end of each auto-diff iteration so that the next iteration treats these as regular differentiable types.
The user-code versions are also lowered into structs in the same way.

Type Checking of Auto-Diff Calls (and other _higher-order_ functions)

Since fwd_diff and bwd_diff are represented as higher order functions that take a function as an input and return the derivative function, the front-end semantic checking needs some notion of higher-order functions to be able to check and lower the calls into appropriate IR.

Higher-order Invocation Base: `HigherOrderInvokeExpr`

All higher order transformations derive from HigherOrderInvokeExpr. For auto-diff there are two possible expression classes ForwardDifferentiateExpr and BackwardDifferentiateExpr, both of which derive from this parent expression.

Higher-order Function Call Checking: `HigherOrderInvokeExprCheckingActions`

Resolving the concrete method is not a trivial issue in Slang, given its support for overloading, type coercion and more. This becomes more complex with the presence of a function transformation in the chain.
For example, if we have fwd_diff(f)(DiffPair<float>(...), DiffPair<double>(...)), we would need to find the correct match for f based on its post-transform argument types.

To facilitate this we use the following workflow:

  1. The HigherOrderInvokeExprCheckingActions base class provides a mechanism for different higher-order expressions to implement their type translation (i.e. what is the type of the transformed function).
  2. The checking mechanism passes all detected overloads for f through the type translation and assembles a new group out of the results (the new functions are 'temporary')
  3. This new group is used by ResolveInvoke when performing overload resolution and type coercion using the user-provided argument list.
  4. The resolved signature (if there is one) is then replaced with the corresponding function reference and wrapped in the appropriate higher-order invoke.

Example:

Let's say we have two functions with the same name f: (int -> float, double, double -> float)
and we want to resolve fwd_diff(f)(DiffPair<float>(1.0, 0.0), DiffPair<float>(0.0, 1.0)).

The higher-order checking actions will synthesize the 'temporary' group of translated signatures (int -> DiffPair<float>, DiffPair<double>, DiffPair<double> -> DiffPair<float>).
Invoke resolution will then narrow this down to a single match (DiffPair<double>, DiffPair<double> -> DiffPair<float>) by automatically casting the floats to doubles. Once the resolution is complete,
we return InvokeExpr(ForwardDifferentiateExpr(f : double, double -> float), casted_args) by wrapping the corresponding function in the corresponding higher-order expr

Attributed Types (`no_diff` parameters)

Often, it will be necessary to prevent gradients from propagating through certain parameters, for correctness reasons. For example, values representing random samples are often not differentiated since the result may be mathematically incorrect.

Slang provides the no_diff operator to mark parameters as non-differentiable, even if they use a type that conforms to IDifferentiable

c
float myFunc(float a, no_diff float b)
{
    return a * b;
}

// Resulting fwd-mode derivative:
DiffPair<float> myFunc(DiffPair<float> dpa, float b)
{
    return diffPair(dpa.p * b, dpa.d * b);
}

Slang uses OpAttributedType to denote the IR type of such parameters. For example, the lowered type of b in the above example is OpAttributedType(OpFloat, OpNoDiffAttr). In the front-end, this is represented through the ModifiedType AST node.

Sometimes, this additional layer can get in the way of things like type equality checks and other mechanisms where the no_diff is irrelevant. Thus, we provide the unwrapAttributedType helper to remove attributed type layers for such cases.

Derivative Data-Flow Analysis

Slang has a derivative data-flow analysis pass that is performed on a per-function basis immediately after lowering to IR and before the linking step (slang-ir-check-differentiability.h/slang-ir-check-differentiability.cpp).

The job of this pass is to enforce that instructions that are of a differentiable type will propagate a derivatives, unless explicitly dropped by the user through detach() or no_diff. The reason for this is that Slang requires functions to be decorated with [Differentiable] to allow it to propagate derivatives. Otherwise, the function is considered non-differentiable, and effectively produces a 0 derivative. This can lead to frustrating situations where a function may be dropping non-differentiable on purpose. Example:

c
float nonDiffFunc(float x)
{
    /* ... */
}

float differentiableFunc(float x) // Forgot to annotate with [Differentiable]
{
    /* ... */
}

float main(float x)
{
    // User doesn't realise that the function that is supposed to be differentiable is not 
    // getting differentiated, because the types here are all 'float'.
    // 
    return nonDiffFunc(x) * differentiableFunc(x);
}

The data-flow analysis step enforces that non-differentiable functions used in a differentiable context should get their derivative dropped explicitly. That way, it is clear to the user whether a call is getting differentiated or dropped.

Same example with no_diff enforcement:

c
float nonDiffFunc(float x)
{
    /* ... */
}

[Differentiable]
float differentiableFunc(float x)
{
    /* ... */
}

float main(float x)
{
    return no_diff(nonDiffFunc(x)) * differentiableFunc(x);
}

A no_diff can only be used directly on a function call, and turns into a TreatAsDifferentiableDecoration that indicates that the function will not produce a derivative.

The derivative data-flow analysis pass works similar to a standard data-flow pass:

  1. We start by assembling a set of instructions that 'produce' derivatives by starting with the parameters of differentiable types (and without an explicit no_diff), and propagating them through each instruction in the block. An inst carries a derivative if there one of its operands carries a derivative, and the result type is differentiable.
  2. We then assemble a set of instructions that expect a derivative. These are differentiable operands of differentiable functions (unless they have been marked by no_diff). We then reverse-propagate this set by adding in all differentiable operands (and repeating this process).
  3. During this reverse-propagation, if there is any OpCall in the 'expect' set that is not also in the 'produce' set, then we have a situation where the gradient hasn't been explicitly dropped, and we create a user diagnostic.

File: docs/design/autodiff.md

Reverse Mode Autodiff (Out of Date)

This document serves as a design reference for reverse-mode auto-diff in the Slang compiler.

Reverse-Mode Passes

Rather than implementing reverse-mode as a separate pass, Slang implements this as a series of independent passes:

If a function needs a reverse-mode version generated:

  • Linearize the function, and all dependencies.
  • Propagate differential types through the linearized code.
  • Unzip by moving primal insts to before differential insts.
  • Transpose the differential insts.

Linearization (Forward-mode)

Overview

(This is a incomplete section. More details coming soon)

Consider an arbitrary function float f(float a, float b, float c, ..., z) which takes in N inputs and generates one output y. Linearization aims to generate the first-order Taylor expansion of f about all of it's inputs.

Mathematically, the forward derivative fwd_f represents df/da * (a_0 - a) + df/db * (b_0 - b) + ..., where a_0 is the value at which the Taylor expansion was produced. The quantity a_0 - a is known as the 'differential' (for brevity we'll denote them da, db, dc, etc..), and there is at-most one differential per input.

Thus, the new function's signature should be fwd_f(float a, float da, float b, float db, float c, float dc, ...). For simplicity, we'll use pairs instead of interleaving the original and differential parameters. We use the intrinsic DifferentialPair<T> (or for short: DP<T>) to denote this.

The signature we use is then fwd_f(DP<float> a, DP<float> b, DP<float> c)

An example of linearization:

c
float f(float a, float b)
{
    if (a > 0)
    {
        return a + b + 2.0 * a * b;
    }
    else
    {
        return sqrt(a);
    }
}

We'll write out the SSA form of this function.

c
float f_SSA(float a, float b)
{
    bool _b1 = a > 0;
    if (_b1)
    {
        float _t1 = a + b;
        float _t2 = 2.0 * a;
        float _t3 = _t2 * b;
        float _t4 = _t1 + _t3;

        return _t4;
    }
    else
    {
        float _t1 = sqrt(a);
        return _t1;
    }
}

DP<float> f_SSA(DP<float> dpa, DP<float> dpb)
{

    bool _b1 = dpa.p > 0;
    if (_b1)
    {
        float _t1 = dpa.p + dpb.p;
        float _t1_d = dpa.d + dpb.d;

        float _t2 = 2.0 * dpa.p;
        float _t2_d = 0.0 * dpa.p + 2.0 * dpa.d;

        float _t3 = _t2 * dpb.p;
        float _t3_d = _t2_d * dpb.p + _t2 * dpb.d;

        float _t4 = _t1 + _t3;
        float _t4_d = _t1_d + _t3_d;

        return DP<float>(_t4, _t4_d);
    }
    else
    {
        DP<float> _t1_dp = sqrt_fwd(dpa);
        return DP<float>(_t1_dp.p, _t1_dp.d);
    }
}

In the result, the primal part of the pair holds the original computation, while the differential part computes the dot product of the differentials with the derivatives of the function's output w.r.t each input.

Propagation

This step takes a linearized function and propagates information about which instructions are computing a differential and which ones are part of the primal (original) computation.

Assuming first-order differentiation only:
The approach will be to mark any instructions that extract the differential from the differential pair as a differential. Then any instruction that uses the differential is itself marked as a differential and so on. The only exception is the call instruction which is either non-differentiable (do nothing) or differentiable and returns a pair (follow the same process)

Here's the above example with propagated type information (we use float.D to denote intermediaries that have been marked as differential, and also expand everything so that each line has a single operation)

c
DP<float> f_SSA_Proped(DP<float> dpa, DP<float> dpb)
{
    bool _b1 = dpa.p > 0;
    if (_b1)
    {
        float _t1 = dpa.p + dpb.p;
        
        float.D _q1_d = dpa.d;
        float.D _q2_d = dpb.d;

        float.D _t1_d = _q1_d + _q2_d;

        float _t2 = 2.0 * dpa.p;
        
        float.D _q2_d = dpa.d;
        float.D _q3_d = 2.0 * dpa.d;

        float _q4 = dpa.p;
        float.D _q4_d = 0.0 * dpa.p;

        float.D _t2_d = _q4_d + _q3_d;

        float _t3 = _t2 * dpb.p;

        float _q5 = dpb.p;
        float.D _q6_d = _q5 * _t2_d;

        float.D _q7_d = dpb.d;
        float.D _q8_d = _t2 * _q7_d

        float _t3_d = _q6_d + _q8_d;

        float _t4 = _t1 + _t3;

        float.D _t4_d = _t1_d + _t3_d;

        return DP<float>(_t4, _t4_d);
    }
    else
    {
        DP<float> _t1_dp = sqrt_fwd(dpa);

        float _q1 = _t1_dp.p;
        float.D _q1_d = _t1_dp.d;

        return DP<float>(_q1, _q1_d);
    }
}

Unzipping

This is a fairly simple process when there is no control flow. We simply move all non-differential instructions to before the first differential instruction.

When there is control flow, we need to be a bit more careful: the key is to replicate the control flow graph once for primal and once for the differential.

Here's the previous example unzipped:

c
DP<float> f_SSA_Proped(DP<float> dpa, DP<float> dpb)
{
    bool _b1 = dpa.p > 0;

    float _t1, _t2, _q4, _t3, _q5, _t3_d, _t4, _q1;

    if (_b1)
    {
        _t1 = dpa.p + dpb.p;
        
        _t2 = 2.0 * dpa.p;
        
        _q4 = dpa.p;
        
        _t3 = _t2 * dpb.p;

        _q5 = dpb.p;

        _t4 = _t1 + _t3;

    }
    else
    {

        _q1 = sqrt_fwd(DP<float>(dpa.p, 0.0));
    }

    // Note here that we have to 'store' all the intermediaries 
    // _t1, _t2, _q4, _t3, _q5, _t3_d, _t4 and _q1. This is fundamentally
    // the tradeoff between fwd_mode and rev_mode

    if (_b1)
    {
        float.D _q1_d = dpa.d;
        float.D _q2_d = dpb.d;

        float.D _t1_d = _q1_d + _q2_d;

        float.D _q2_d = dpa.d;
        float.D _q3_d = 2.0 * dpa.d;

        float.D _q4_d = 0.0 * dpa.p;

        float.D _t2_d = _q4_d + _q3_d;

        float.D _q6_d = _q5 * _t2_d;

        float.D _q7_d = dpb.d;
        float.D _q8_d = _t2 * _q7_d

        float.D _t3_d = _q6_d + _q8_d;

        float.D _t4_d = _t1_d + _t3_d;

        return DP<float>(_t4, _t4_d);
    }
    else
    {
        DP<float> _t1_dp = sqrt_fwd(dpa);

        float.D _q1_d = _t1_dp.d;

        return DP<float>(_q1, _q1_d);
    }
}

Transposition

Overview

This transposition pass assumes that provided function is linear in it's differentials.
It is out of scope of this project to attempt to enforce that constraint for user-defined differential code.

For transposition we walk all differential instructions in reverse starting from the return statement, and apply the following rules:

We'll have an accumulator dictionary Dictionary accMap holding assignments for
intermediaries which don't have concrete variables. When we add a pair (A, C) and (A, B) already exists, this will form the pair (A, ADD(C, B)) in the dictionary. (ADD will be replaced with a call to T.dadd for a generic type T)

  • If inst is a RETURN(A), add pair (A, d_out) to accMap
  • If an instruction is MUL(P, D) where D is the differential, add pair (D, MUL(P, accMap[this_inst])) to accMap
  • If an instruction is ADD(D1, D2), where both D1 and D2 are differentials (this is the only config that should occur), then add pair (D1, accMap[this_inst]) to accMap
  • If an instruction is CALL(f_fwd, (P1, D1), (P2, D2), ...), create variables D1v, D2v, ... for D1, D2, ..., then replace with CALL(f_rev, (P1, D1v), (P2, D2v), ..., accMap[this_inst]), and finally add pairs (D1, LOAD[D1v]), (D2, LOAD[D2v]), ... to accMap
c
void f_SSA_Rev(inout DP<float> dpa, inout DP<float> dpb, float dout)
{
   bool _b1 = dpa.p > 0;

   float _t1, _t2, _q4, _t3, _q5, _t3_d, _t4, _q1;

   if (_b1)
   {
       _t1 = dpa.p + dpb.p;
       
       _t2 = 2.0 * dpa.p;
       
       _q4 = dpa.p;
       
       _t3 = _t2 * dpb.p;

       _q5 = dpb.p;

       _t4 = _t1 + _t3;

   }
   else
   {

       _q1 = sqrt_fwd(DP<float>(dpa.p, 0.0));
   }

   // Note here that we have to 'store' all the intermediaries 
   // _t1, _t2, _q4, _t3, _q5, _t3_d, _t4 and _q1. This is fundamentally
   // the tradeoff between fwd_mode and rev_mode

   if (_b1)
   {

       float.D _t4_rev = d_out;

       float.D _t1_rev = _t4_rev;
       float.D _t3_rev = _t4_rev;

       float.D _q8_rev = _t3_rev;
       float.D _q6_rev = _t3_rev;

       float.D _q7_rev = _t2 * _q8_rev;

       dpb.d += _q7_rev;

       float.D _t2_rev = _q5 * _q6_rev;

       float.D _q4_rev = _t2_rev;
       float.D _q3_rev = _t2_rev;

       dpa.d += 2.0 * _q3_rev;

       float.D _q1_rev = _t1_rev;
       float.D _q2_rev = _t1_rev;

       dpb.d += _q2_rev;
       dpa.d += _q1_rev;
   }
   else
   {
       _q1_rev = d_out;

       DP<float> dpa_copy;
       sqrt_rev(dpa_copy, _q1_rev);

       dpa.d += dpa_copy.d;
   }
}

File: docs/design/capabilities.md

Capabilities (Out of Date)

Slang aims to be a portable language for shader programming, which introduces two complementary problems:

  1. We need a way to indicate that certain constructs (types, functions, etc.) are only allowed on certain targets, so that a user gets a meaningful error if they try to do something that won't work on one or more of the APIs or platforms they want to target. Similarly, the user expects to get an error if they call a fragment-shader-specific function inside of, say, compute shader code, or vice versa.

  2. If the same feature can be implemented across multiple platforms, but the best (or only) implementation path differs across platforms, then we need a way to express the platform specific code and pick the right implementation per-target.

Item (2) is traditionally handled with preprocessor techniques (e.g., #ifdefing the body of a function based on target platform), but that of course requires that the user invoke the Slang front end once for each target platform, and target-specific coding in a library will then "infect" code that uses that library, forcing them to invoke the front-end once per target as well.

We are especially sensitive to this problem in the compiler itself, because we have to author and maintain the Slang standard modules, which needs to (1) expose the capabilities of many platforms and (2) work across all those platforms. It would be very unfortunate if we had to build different copies of our standard modules per-target.

The intention in Slang is to solve both of these problems with a system of capabilities.

What is a capability?

For our purposes a capability is a discrete feature that a compilation target either does or does not support.
We could imagine defining a capability for the presence of texture sampling operations with implicit gradients; this capability would be supported when generating fragment shader kernel code, but not when generating code for other stages.

Let's imagine a language syntax that the standard modules could use to define some atomic capabilities:

text
capability implicit_gradient_texture_fetches;

We can then imagine using attributes to indicate that a function requires a certain capability:

text
struct Texture2D
{
	...

	// Implicit-gradient sampling operation.
	[availableFor(implicit_gradient_texture_fetches)]
	float4 Sample(SamplerState s, float2 uv);
}

(Note that the [availableFor(...)] syntax is just a straw-man to write up examples, and a better name would be desirable if/when we implement this stuff.)

Given those declarations, we could then check when compiling code if the user is trying to call Texture2D.Sample in code compiled for a target that doesn't support implicit-gradient texture fetches, and issue an appropriate error.
The details on how to sequence this all in the compiler will be covered later.

Derived Capabilities

Once we can define atomic capabilities, the next step is to be able to define derived capabilities.
Let's imagine that we extend our capability syntax so that we can define a new capability that automatically implies one or more other capabilities:

text
capability fragment : implicit_gradient_texture_fetches;

Here we've said that whenever the fragment capability is available, we can safely assume that the implicit_gradient_texture_fetches capability is available (but not vice versa).

Given even a rudimentary tool like that, we can start to build up capabilities that relate closely to the "profiles" in things like D3D:

text
capability d3d;
capability sm_5_0 : d3d;
capability sm_5_1 : sm_5_0;
capability sm_6_0 : sm_5_1;
...

capability d3d11 : d3d, sm_5_0;
capability d3d12 : d3d, sm_6_0;

capability khronos;
capability glsl_400 : khronos;
capability glsl_410 : glsl_400;
...

capability vulkan : khronos, glsl_450;
capability opengl : khronos;

Here we are saying that sm_5_1 supports everything sm_5_0 supports, and potentially more. We are saying that d3d12 supports sm_6_0 but maybe not, e.g., sm_6_3.
We are expressing that fact that having a glsl_* capability means you are on some Khronos API target, but that it doesn't specify which one.
(The exact details of these declarations obviously aren't the point; getting a good hierarchy of capabilities will take time.)

Capability Composition

Sometimes we'll want to give a distinct name to a specific combination of capabilities, but not say that it supports anything new:

text
capability ps_5_1 = sm_5_1 & fragment;

Here we are saying that the ps_5_1 capability is equivalent to the combination of sm_5_1 and fragment (that is, if you support both sm_5_1 and fragment then you support ps_5_1 and vice versa).

Compositions should be allowed in [availableFor(...)] attributes (e.g., [availableFor(vulkan & glsl_450)]), but pre-defined compositions should be favored when possible.

When composing things with & it is safe for the compiler to filter out redundancies based on what it knows so that, e.g., ps_5_0 & fragment resolves to just ps_5_0.

Once we have an & operator for capabilities, it is easy to see that "derived" capabilities are really syntax sugar, so that a derived capability like:

text
capability A : B, C

could have been written instead as :

text
capability A_atomic
capability A = A_atomic & B & C

Where the A_atomic capability guarantees that A implies B and C but not vice versa.

It is also useful to think of an | operator on capabilities.
In particular if a function has multiple [availableFor(...)] attributes:

text
[availableFor(vulkan & fragment)]
[availableFor(d3d12 & fragment)]
void myFunc();

This function should be equivalent to one with just a single [availableFor((vulkan & fragment) | (d3d12 & fragment))] which is equivalent to [availableFor((vulkan | d3d12) & fragment)].
Simplification should generally push toward "disjunctive normal form," though, rather than pursue simplifications like that.
Note that we do not include negation, so that capabilities are not general Boolean expressions.

Validation

For a given function definition F, the front end will scan its body and see what it calls, and compose the capabilities required by the called functions using & (simplifying along the way). Call the resulting capability (in disjunctive normal form) R.

If F doesn't have an [availableFor(...)] attribute, then we can derive its effective [availableFor(...)] capability as R (this probably needs to be expressed as an iterative dataflow problem over the call graph, to handle cycles).

If F does have one or more [availableFor(...)] clauses that amount to a declared capability C (again in disjunctive normal form), then we can check that C implies R and error out if it is not the case.
A reasonable implementation would track which calls introduced which requirements, and be able to explain why C does not capture the stated requirements.

For a shader entry point, we should check it as if it had an [availableFor(...)] that is the OR of all the specified target profiles (e.g., sm_5_0 | glsl_450 | ...) ANDed with the specified stage (e.g., fragment).
Any error here should be reported to the user.
If an entry point has an explicit [availableFor(...)] then we should AND that onto the profile computed above, so that the user can restrict certain entry points to certain profiles.

In order to support separate compilation, the functions that are exported from a module should probably either have explicit availability attributes, or else they will be compiled against a kind of "default capability" used for the whole module.
Downstream code that consumes such a module would see declarations with explicit capabilities only.
Picking an appropriate "default capability" to use when compiling modules is an important challenge; it would in practice define the "min spec" to use when compiling.

Capability Overriding

It should be possible to define multiple versions of a function, having different [availableFor(...)] attributes:

text
[availableFor(vulkan)] void myFunc() { ... }

[availableFor(d3d12)] void myFunc() { ... }

For front-end checking, these should be treated as if they were a single definition of myFunc with an ORed capability (e.g., vulkan | d3d12).
Overload resolution will pick the "best" candidate at a call site based only on the signatures of the function (note that this differs greatly from how profile-specific function overloading works in Cg).

The front-end will then generate initial IR code for each definition of myFunc.
Each of the IR functions will have the same mangled name, but different bodies, and each will have appropriate IR decorations to indicate the capabilities it requires.

The choice of which definition to use is then put off until IR linking for a particular target.
At that point we can look at all the IR functions matching a given mangled name, filter them according to the capabilities of the target, and then select the "best" one.

In general a definition A of an IR symbol is better than another definition B if the capabilities on A imply those on B but not versa.
(In practice this probably needs to be "the capabilities on A intersected with those of the target," and similarly for B)

This approach allows us to defer profile-based choices of functions to very late in the process. The one big "gotcha" to be aware of is when functions are overloaded based on pipeline stage, where we would then have to be careful when generating DXIL or SPIR-V modules with multiple entry points (as a single function f might need to be specialized twice if it calls a stage-overloaded function g).

Capabilities in Other Places

So far I've talked about capabilities on functions, but they should also be allowed on other declarations including:

  • Types, to indicate that code using that type needs the given capability
  • Interface conformances, to indicate that a type only conforms to the interface when the capabilities are available
  • Struct fields, to indicate that the field is only present in the type when the capabilities are present
  • Extension declarations, to indicate that everything in them requires the specified capabilities

We should also provide a way to specify that a register or other layout modifier is only applicable for specific targets/stages. Such a capability nominally exists in HLSL today, but it would be much more useful if it could be applied to specify target-API-specific bindings.

Only functions should support overloading based on capability. In all other cases there can only be one definition of an entity, and capabilities just decide when it is available.

API Extensions as Capabilities

One clear use case for capabilities is to represent optional extensions, including cases where a feature is "built-in" in D3D but requires an extension in Vulkan:

text
capability KHR_secret_sauce : vulkan;

[available_for(sm_7_0)] // always available for D3D Shader Model 7.0
[available_for(KHR_secret_sauce)] // Need the "secret sauce" extension for Vulkan
void improveShadows();

When generating code for Vulkan, we should be able to tell the user that the improveShadows() function requires the given extension. The user should be able to express compositions of capabilities in their -profile option (and similarly for the API):

text
slangc code.slang -profile vulkan+KHR_secret_sauce

(Note that for the command line, it is beneficial to use + instead of & to avoid conflicts with shell interpreters)

An important question is whether the compiler should automatically infer required extensions without them being specified, so that it produces SPIR-V that requires extensions the user didn't ask for.
The argument against such inference is that users should opt in to non-standard capabilities they are using, but it would be unfortunate if this in turn requires verbose command lines when invoking the compiler.
It should be possible to indicate the capabilities that a module or entry point should be compiled to use without command-line complications.

(A related challenge is when a capability can be provided by two different extensions: how should the compiler select the "right" one to use?)

Disjoint Capabilities

Certain compositions of capabilities make no sense. If a user declared a function as needing vulkan & d3d12 they should probably get an error message.

Knowing that certain capabilities are disjoint can also help improve the overall user experience.
If a function requires (vulkan & extensionA) | (d3d12 & featureb) and we know we are compiling for vulkan we should be able to give the user a pointed error message saying they need to ask for extensionA, because adding featureB isn't going to do any good.

As a first-pass model we could have a notion of abstract capabilities that are used to model the root of hierarchies of disjoint capabilities:

text
abstract capability api;

abstract capability d3d : api;
capability d3d11 : d3d;
capability d3d12 : d3d;

abstract capability khronos : api;
capability vulkan : khronos;
capability opengl : khronos;

As a straw man: we could have a rule that to decide if non-abstract capabilities A and B are disjoint, we look for their common ancestor in the tree of capabilities.
If the common ancestor is abstract, they are disjoint, and if not they not disjoint.
We'd also know that if the user tries to compile for a profile that includes an abstract capability but not some concrete capability derived from it, then that is an error (we can't generate code for just d3d).

The above is an over-simplification because we don't have a tree of capabilities, but a full graph, so we'd need an approach that works for the full case.

Interaction with Generics/Interfaces

It should be possible for an interface requirement to have a capability requirement attached to it.
This would mean that users of the interface can only use the method/type/whatever when the capability is present (just like for any other function):

text
interface ITexture
{
	float4 sampleLevel(float2 uv, float lod);

	[availableFor(fragment)]
	float4 sample(float2 uv); // can only call this from fragment code
}

When implementing an interface, any capability constraints we put on a member that satisfies an interface requirement would need to guarantee that either:

  • the capabilities on our method are implied by those on the requirement (we don't require more), or

  • the capabilities on the method are implied by those on the type itself, or its conformance to the interface (you can't use the conformance without the capabilities), or

  • the capabilities are already implied by those the whole module is being compiled for

In each case, you need to be sure that YourType can't be passed as a generic argument to some function that uses just the ITexture interface above and have them call a method on your type from a profile that doesn't have the required capabilities.

Interaction with Heterogeneity

If Slang eventually supports generating CPU code as well as shaders, it should use capabilities to handle the CPU/GPU split similar to how they can be used to separate out vertex- and fragment-shader functionality.
Something like a cpu profile that works as a catch-all for typical host CPU capabilities would be nice, and could be used as a convenient way to mark "host" functions in a file that is otherwise compiled for a "default profile" that assumes GPU capabilities.

Conclusion

Overall, the hope is that in many cases developers will be able to use capability-based partitioning and overloading of APIs to build code that only has to pass through the Slang front-end once, but that can then go through back-end code generation for each target.
In cases where this can't be achieved, the way that capability-based overloading is built into the Slang IR design means that we should be able to merge multiple target-specific definitions into one IR module, so that a module can employ target-specific specializations while still presenting a single API to consumers.


File: docs/design/casting.md

Casting in the Slang Compiler

The following discussion is about casting within the C++ implementation of the slang compiler.

C++'s built in mechanisms for casting (principally dynamic_cast) is problematic within the slang compiler codebase. Code using 'dynamic_cast' requires RTTI information is available, and that a type that uses it must have a vtbl (have at least one virtual member). Some problems with this...

  • There are types which we want to 'dynamic_cast' that do not have, and we do not want to have a Vtbl (for example Slang::IRInst).
  • There are types which a 'dynamic_cast' doesn't do quite what we want (for example casting on Type* derived types typically wants to work on their canonical type)
  • We may want to replace use of dynamic_cast in the future for speed/space or other reasons
  • It is common in the code base when using a 'smart pointer' type to cast it, but still return a smart pointer

To deal with these issues we need casting within Slang to follow it's own methodology. In summary it is as follows...

  • Use 'as' free function to do a typical 'dynamic like' cast.
    • 'as' doesn't guarantee the returned pointer points to the same object.
    • For example with Type* it actually does the cast on the canonical type which is often a different object.
  • If you want to literally do a dynamic cast use 'dynamicCast' free function.
    • This guarantees the returned pointer points to the same object (like normal dynamic_cast)
  • If you want to return a smart pointer from a cast from a smart pointer use the .as or .dynamicCast methods
  • If you want to determine if an 'as' cast is possible on a smart pointer use the .is method
    • Doing so will produce more efficient code because a new smart pointer does not need to be constructed

These functions will also work with types that do not have Vtbl - like IRInst derived types.

Both 'as' and 'dynamicCast' handle the case if the pointer is a nullptr, by returning a nullptr. If the cast succeeds the cast pointer is returned otherwise nullptr is returned. If a cast is performed with a free function it always returns a raw pointer.

So why have 'as' and 'dynamicCast' - they seem sort of similar? The primary difference is dynamicCast must always return a pointer to the same object, whilst 'as' can return a pointer to a different object if that is the desired 'normal' casting behavior for the type. This is the case for Type* when using 'as' it may return a different object - the 'canonical type' for the Type*. For a concrete example take 'NamedExpressionType', its canonical type is the type the name relates to. If you use 'as' on it - it will produce a pointer to a different object, an object that will not be castable back into a NamedExpressionType.

Also keep in mind that 'as' behavior is based on the pointer type being cast from. For any pointer to a type derived from Type it will cast the canonical type. BUT if the pointer is pointing to a Type derived object, but the pointer type is not derived from Type (like say RefObject*), then 'as' will behave like dynamicCast.

All this being said 'as' in usage is seen as the 'default' way to do a 'dynamic like' cast with these special behaviour appropriate for the type when necessary.

By having the free function and method versions of 'as' and 'dynamicCast', you can choose if you want a 'raw' or 'smart' pointer type returned from the cast. If you just want to test if something is a certain type, then using as/dynamicCast free functions is the faster way to do it. If you know that a raw pointer is ok, because the object will remain in scope, then again using the free function is better because it does less work. But as the examples following show, care is needed because if you get it wrong the object might go out of scope and leave the raw pointer pointing to a deleted object. When in doubt the safe choice is to typically use .as (or .dynamicCast if appropriate) methods.

Following example shows the different types of casting...

c++
void someFunction(Decl* decl, Type* type)
{
    RefPtr<Decl> declRefPtr(decl);
    RefPtr<Type> typeRefPtr(type);

    // Use of as
    {
        // Casting with as on a free function returns a raw pointer
        GenericDecl* genericDeclRaw0 = as<GenericDecl>(decl);
        // Free function again returns a raw pointer
        GenericDecl* genericDeclRaw1 = as<GenericDecl>(declRefPtr);

        // Using the as *method* returns a smart pointer holding the cast result
        RefPtr<GenericDecl> genericDeclRefPtr0 = declRefPtr.as<GenericDecl>();
        
        // Of course you can use auto with either
        auto genericDeclRefPtr1 = declRefPtr.as<GenericDecl>();
        
        auto genericDeclRaw2 = as<GenericDecl>(declRefPtr);
    }
    
    // Currently using as on anything not cast *from* Type is the same as dynamicCast.
    // But on Type* sometimes you may want to control the cast
    {
        // With a NamedExpressionType sometimes you don't want 'as' behaviour - if we want to see the information about the name (not the thing 
        // it relates to (the canonical type)
        NamedExpressionType* namedExpressionRawPtr = dynamicCast<NamedExpressionType>(type);
        
        
        // Returns the smart pointer 
        auto namedExpressionRefPtr = typeRefPtr.as<NamedExpressionType>();
    }

It is important to be aware of what style of cast you use where. Take for example the following function ...

c++
RefPtr<Expr> substitute(RefPtr<Expr> expr) const
    {
        return DeclRefBase::Substitute(expr);
    }

If you want to do a cast on it, you need to be careful especially about scope, for example...

c++
RefPtr<Expr> expr = ...;
    
    {
        // Whoops! This is a problem. When using the free function, the cast is to a *raw* pointer, so obj 
        // receives a raw pointer. When the RefPtr returned from Substitute goes out of scope (when the statement is left)
        // the ref will be removed and if the ref count was 1 destroyed. Now obj points to a freed object and so a crash is
        // likely to follow in the future! 
        
        auto obj = as<RefObject>(substitute(expr));
    }
    // So how do we avoid this? Well it depends what the function is returning and the scope. If it's returning a smart pointer, 
    // you could use the .as method
    {
        // This can only compile if it is a smart pointer (raw pointers don't have an as method)
        auto obj = substitute(expr).as<RefObject>();
    }

    // Another option is to put the created thing in a smart pointer so you know it's in scope
    {
        RefPtr<Expr> sub = substitute(expr);
        // Ok as long as sub is in scope
        auto obj = as<RefObject>(sub);
       
    }
 
    // More awkwardly you could use free function, but assign to a smart pointer, thus maintaining scope
    {
        RefPtr<RefObject> obj = as<RefObject>(substitute(expr));
    }

The following code shows the change in behavior of 'as' is based on the source pointer type NOT the object type..

c++
// Derives from Type
    NamedExpressionType* exprType = ...;

    
    // Will be the Type* of the *canonical* type, because the pointer is Type derived and we are using as!
    Type* type0 = as<Type>(exprType);
    // It' going to be pointing to a different object, because type0 is the cast of the *canonical* type, because exprType derives from Type
    SLANG_ASSERT(type0 != exprType);
    
    // If I do a dynamicCast the result is either nullptr or a pointer that *must* point to the same object
    Type* type1 = dynamicCast<Type>(exprType);
    SLANG_ASSERT(type1 == exprType);
    
    
    // Here, the pointer is pointing to a NamedExpressionType derived object. Which derives from Type. BUT our pointer here does *not* derive from type.
    RefObject* refObj = exprType;
    
    // 'as' just looks at the from type, and it doesn't derive from Type (it's just RefObject), so it does regular as, which is dynamicCast
    Type* type2 = as<Type>(refObject);
    
    SLANG_ASSERT(type2 == exprType);
    
    // Finally... 
    
    // Is true even though exprType is a NamedExpression, because the cast is on the canonical type
    SLANG_ASSERT(as<NamedExpression>(exprType) == nullptr);
    
    // dynamicCast is always the same object returned, so must match
    SLANG_ASSERT(dynamicCast<NamedExpression>(exprType) == exprType);

File: docs/design/coding-conventions.md

Slang Project Coding Conventions

Principles

This document attempts to establish conventions to be used in the Slang codebase.
We have two goals for this convention.

The first goal is to make the code look relatively consistent so that it is easy to navigate and understand for contributors.
Having varying styles across different modules, files, functions, or lines of code makes the overall design and intention of the codebase harder to follow.

The second goal is to minimize the scope complexity of diffs when multiple maintainers work together on the codebase.
In the absence of an enforced style, developers tend to "clean up" code they encounter to match their personal preferences, and in so doing create additional diffs that increase the chances of merge conflicts and pain down the line.

Because the Slang codebase has passed through many hands and evolved without a pre-existing convention, these two goals can come into conflict.
We encourage developers to err on the side of leaving well enough alone (favoring the second goal).
Don't rewrite or refactor code to match these conventions unless you were already going to have to touch all of those lines of code anyway.

Note that external code that is incorporated into the project is excluded from all of these conventions.

Languages

C++

Most code in the Slang project is implemented in C++.
We currently assume support for some C++11 idioms, but have explicitly avoided adding dependencies on later versions.

As a general rule, be skeptical of "modern C++" ideas unless they are clearly better to simpler alternatives.
We are not quite in the realm of "Orthodox C++", but some of the same guidelines apply:

  • Don't use exceptions for non-fatal errors (and even then support a build flag to opt out of exceptions)
  • Don't use the built-in C++ RTTI system (home-grown is okay)
  • Don't use the C++ variants of C headers (e.g., <cstdio> instead of <stdio.h>)
  • Don't use the STL containers
  • Don't use iostreams

The compiler implementation does not follow some of these guidelines at present; that should not be taken as an excuse to further the proliferation of stuff like dynamic_cast.
Do as we say, not as we do.

Some relatively recent C++ features that are okay to use:

  • Rvalue references for "move semantics," but only if you are implementing performance-critical containers or other code where this really matters.

  • auto on local variables, if the expected type is clear in context

  • Lambdas are allowed, but think carefully about whether just declaring a subroutine would also work.

  • Using >> to close multiple levels of templates, instead of > > (but did you really need all those templates?)

  • nullptr

  • enum class

  • Range-based for loops

  • override

  • Default member initializers in class/struct bodies

Templates are suitable in cases where they improve clarity and type safety.
As a general rule, it is best when templated code is kept minimal, and forwards to a non-templated function that does the real work, to avoid code bloat.

Any use of template metaprogramming would need to prove itself exceptionally useful to pay for the increase in cognitive complexity.
We don't want to be in the business of maintaining "clever" code.

As a general rule, const should be used sparingly and only with things that are logically "value types."
If you find yourself having to const-qualify a lot of member function in type that you expect to be used as a heap-allocated object, then something has probably gone wrong.

As a general rule, default to making the implementation of a type public, and only encapsulate state or operations with private when you find that there are complex semantics or invariants that can't be provided without a heavier hand.

Slang

The Slang project codebase also includes .slang files implementing the Slang core module, as well as various test cases and examples.
The conventions described here are thus the "official" recommendations for how users should format Slang code.

To the extent possible, we will try to apply the same basic conventions to both C++ and Slang.
In places where we decide that the two languages merit different rules, we will point it out.

Files and Includes

File Names

All files and directories that are added to codebase should have names that contain only ASCII lower-case letters, digits, dots (.) and dashes (-).
Operating systems still vary greatly in their handling of case sensitivity for file names, and non-ASCII code points are handled with even less consistency; sticking to a restricted subset of ASCII helps avoids some messy interactions between case-insensitive file systems and case-sensitive source-control systems like Git.
As with all these conventions, files from external projects are exempted from these restrictions.

Naming of Source and Header Files

In general the C++ codebase should be organized around logical features/modules/subsystem, each of which has a single .h file and zero or more .cpp files to implement it.

If there is a single .cpp file, its name should match the header: e.g., parser.h and parser.cpp.

If there is more than one .cpp file, their names should start with the header name: e.g., parser.h and parser-decls.cpp and parser-exprs.cpp.
If there are declarations that need to be shared by the .cpp files, but shouldn't appear in the public interface, then can go in a *-impl.h header (e.g., parser-impl.h).

Use best judgement when deciding what counts as a "feature." One class per file is almost always overkill, but the codebase currently leans too far in the other direction, with some oversized source files.

Headers

Every header file should have an include guard.
Within the implementation we can use #pragma once, but exported API headers (slang.h) should use traditional #ifdef style guards (and they should be consumable as both C and C++).

A header should include or forward-declare everything it needs in order to compile.
It is not up to the programmer who #includes a header to sort out the dependencies.

Avoid umbrella or "catch-all" headers.

Source Files

Every source file should start by including the header for its feature/module, before any other includes (this helps ensure that the header correctly includes its dependencies).

Functions that are only needed within that one source file can be marked static, but we should avoid using the same name for functions in different files (in order to support lumped/unified builds).

Includes

Cross-directory project-internal includes should name a header from an include root and should not
use ../ to traverse directories. The source/ directory is an include root, so cross-module
includes should use paths such as:

c++
#include "compiler-core/slang-source-loc.h"
#include "core/slang-string.h"

Targets that consume headers under another project directory, such as tools/ or examples/,
should expose that directory through their CMake include directories and use paths relative to that
root. Prefer a private include directory unless it is part of the target's public interface.

In general, includes should be grouped as follows:

  • First, the correspodning feature/module header, if we are in a source file
  • Next, any <>-enlosed includes for system/OS headers
  • Next, any ""-enclosed includes for external/third-part code that is stored in the project repository
  • Finally, any includes for other features in the project

Within each group, includes should be sorted alphabetically.
If this breaks because of ordering issues for system/OS/third-party headers (e.g., <windows.h> must be included before <GL/GL.h>), then ideally those includes should be mediated by a Slang-project-internal header that features can include.

Namespaces

Favor fewer namespaces when possible.
Small programs may not need any.

All standard module code that a Slang user might link against should go in the Slang namespace for now, to avoid any possibility of clashes in a static linking scenario.
The public C API is obviously an exception to this.

Code Formatting

  • For C++ files, please format using clang-format; .clang-format files in
    the source tree define the style.
  • For CMake files, please format using gersemi
  • For shell scripts, please format using shfmt
  • For YAML files, please use prettier

The formatting for the codebase is overall specified by the
extras/formatting.sh script.

If you open a pull request and the formatting is incorrect, you can comment
/format and a bot will format your code for you.

Naming

Casing

Types should in general use UpperCamelCase. This includes structs, classes, enums and typedefs.

Values should in general use lowerCamelCase. This includes functions, methods, local variables, global variables, parameters, fields, etc.

Macros should in general use SCREAMING_SNAKE_CASE.
It is important to prefix all macros (e.g., with SLANG_) to avoid collisions, since namespaces don't affect macros).

In names using camel case, acronyms and initialisms should appear eniterly in either upper or lower case (e.g., D3DThing d3dThing) and not be capitalized as if they were ordinary words (e.g., D3dThing d3dThing).
Note that this also applies to uses of "ID" as an abbreviation for "identifier" (e.g., use nodeID instead of nodeId).

Prefixes

Prefixes based on types (e.g., p for pointers) should never be used.

Global variables should have a g prefix, e.g. gCounter.
Non-const static class members can have an s prefix if that suits your fancy.
Of course, both of these should be avoided, so this shouldn't come up often.

Constant data (in the sense of static const) should have a k prefix.

In contexts where "information hiding" is relevant/important, such as when a type has both public and private members, or just has certain operations/fields that are considered "implementation details" that most clients should not be using, an m_ prefix on member variables and a _ prefix on member functions is allowed (but not required).

In function parameter lists, an in, out, or io prefix can be added to a parameter name to indicate whether a pointer/reference/buffer is intended to be used for input, output, or both input and output.
For example:

c++
void copyData(void* outBuffer, void const* inBuffer, size_t size);

Result lookupThing(Key k, Thing& outThing);

void maybeAppendExtraNames(std::vector<Name>& ioNames);

Public C APIs will prefix all symbol names while following the casing convention (e.g. SlangModule, slangLoadModule, etc.).

Internal `spirv_asm` result registers

Result registers named inside a spirv_asm block in the core module (the *.meta.slang files) must be given a __ prefix, e.g. %__result rather than %result.
The SPIR-V emitter emits an OpName for every named spirv_asm register, so an un-prefixed internal name such as %result surfaces in the generated module's debug names where it can be mistaken for a user symbol.
The __ prefix marks these names as compiler-internal, matching the convention above for other not-for-normal-use core-module identifiers.
This is enforced at parse time: the parser asserts that every named spirv_asm register in core-module code begins with __.
User-authored spirv_asm blocks are unaffected โ€” they may name registers however they like.

Enums

C-style enum should use the following convention:

c++
enum Color
{
    kColor_Red,
    kColor_Green,
    kColor_Blue,

    kColorCount,
};

When using enum class, drop the k and type name as prefix, but retain the UpperCamelCase tag names:

c++
enum class Color
{
    Red,
    Green,
    Blue,

    Count,
};

When defining a set of flags, separate the type definition from the enum:

c++
typedef unsigned int Axes;
enum
{
    kAxes_None = 0,

    kAxis_X = 1 << 0,
    kAxis_Y = 1 << 1,
    kAxis_Z = 1 << 2,

    kAxes_All = kAxis_X | kAxis_Y | kAxis_Z,
};

Note that the type name reflects the plural case, while the cases that represent individual bits are named with a singular prefix.

In public APIs, all enums should use the style of separating the type definition from the enum, and all cases should use SCREAMING_SNAKE_CASE:

c++
typedef unsigned int SlangAxes;
enum
{
    SLANG_AXES_NONE = 0,

    SLANG_AXIS_X = 1 << 0,
    SLANG_AXIS_Y = 1 << 1,
    SLANG_AXIS_Z = 1 << 2,

    SLANG_AXES_ALL = SLANG_AXIS_X | SLANG_AXIS_Y | SLANG_AXIS_Z,
};

General

Names should default to the English language and US spellings, to match the dominant conventions of contemporary open-source projects.

Function names should either be named with action verbs (get, set, create, emit, parse, etc.) or read as questions (isEnabled, shouldEmit, etc.).

Whenever possible, compiler concepts should be named using the most widely-understood term available: e.g., we use Token over Lexeme, and Lexer over Scanner simply because they appear to be the more common names.

Avoid abbreviations and initialisms unless they are already widely established across the codebase; a longer name may be cumbersome to write in the moment, but the code will probably be read many more times than it is written, so clarity should be preferred.
An important exception to this is common compiler concepts or techniques which may have laboriously long names: e.g., Static Single Assignment (SSA), Sparse Conditional Copy Propagation (SCCP), etc.

One gotcha particular to compiler front-ends is that almost every synonym for "type" has some kind of established technical meaning; most notably the term "kind" has a precise meaning that is relevant in our domain.
It is common practice in C and C++ to define tagged union types with a selector field called a "type" or "kind," which does not usually match this technical definition.
If a developer wants to avoid confusion, they are encouraged to use the term "flavor" instead of "type" or "kind" since this term (while a bit silly) is less commonly used in the literature.

Comments and Documentation

You probably know the drill: comments are good, but an out-of-date comment can be worse than no comment at all.
Try to write comments that explain the "why" of your code more than the "what."
When implementing a textbook algorithm or technique, it may help to imagine giving the reviewer of your code a brief tutorial on the topic.

In cases where comments would benefit from formatting, use Markdown syntax.
We do not currently have a setup for extracting documentation from comments, but if we add one we will ensure that it works with Markdown.

When writing comments, please be aware that your words could be read by many people, from a variety of cultures and backgrounds.
Default to a plain-spoken and professional tone and avoid using slang, idiom, profanity, etc.

2. Official Technical Reference & Guides (shader-slang/shader-slang.github.io)

File: README.md

Shader-Slang.com website

This theme is losely based on the Serif theme for Jekyll by Zerostatic Themes.

Installation

Installing Ruby & Jekyll

If this is your first time using Jekyll, please follow the Jekyll docs and make sure your local environment (including Ruby) is setup correctly.

Installing Theme

Download or clone the theme.

To run the theme locally, navigate to the theme directory and run:

text
bundle install

To start the Jekyll local development server.

text
bundle exec jekyll serve

To build the theme.

text
bundle exec jekyll build

Github Pages

This theme has been tested to work with Github Pages (and Github Project Pages). When using Github Pages you will need to update the baseurl in the _config.yml otherwise all the css, images and paths will be broken.

For example the site https://shader-slang.com would have baseurl: "/"


File: docs/index.md


text
<h3>Overview
                <hr>
            </h3>
            {% assign overview_items = site.data.documentation.overview %}
            
            {% for item in overview_items %}
            {% if item.title and item.title != '' %}
                
                    <h4>{{ item.title }}</h4>
                    {{ item.description }}
                    {% if item.link_url and item.link_url != '' %}
                    [{{ item.link_label }}]({{ item.link_url }})
                    {% else %}
                    {{ item.link_label }}
                    {% endif %}
                
                {% endif %}
            {% endfor %}
            
        
    



    
        
            <h3>Articles
                <hr>
            </h3>
            {% assign articles_items = site.data.documentation.articles %}
                
                {% for item in articles_items %}
                {% if item.title and item.title != '' %}
                    
                        <h4>{{ item.title }}</h4>
                        {{ item.description }}
                    {% if item.link_url and item.link_url != '' %}
                    [{{ item.link_label }}]({{ item.link_url }})
                    {% else %}
                    {{ item.link_label }}
                    {% endif %}
                    
                {% endif %}
                {% endfor %}
                
        
    


    
        
            <h3>Tutorials
                <hr>
            </h3>
            {% assign tutorials_items = site.data.documentation.tutorials %}
                
                {% for item in tutorials_items %}
                {% if item.title and item.title != '' %}
                    
                        <h4>{{ item.title }}</h4>
                        {{ item.description }}
                    {% if item.link_url and item.link_url != '' %}
                    [{{ item.link_label }}]({{ item.link_url }})
                    {% else %}
                    {{ item.link_label }}
                    {% endif %}
                    
                {% endif %}
                    {% endfor %}
                
        
    




    
        
            <h3>Contributions
                <hr>
            </h3>
            If youโ€™d like to contribute to the project, we are excited to have your input. Our community pages provides information
            on our structure and our process for accepting contributions.
            <a class="btn btn-primary" href="/community/">Making Community Contributions</a>

File: docs/getting-started.md


text
<h3>Try Slang
                <hr>
            </h3>
            Try slang locally in your browser without downloading or installing anything with the Slang
                Playground.
            <a class="btn btn-primary " href="https://try.shader-slang.org/">Try the Slang Playground</a>
        
        
            
        
    






    
        <h3>Machine Learning with Slang
            <hr>
        </h3>
        The SlangPy package seamlessly bridges the graphics and machine learning ecosystems in Python. You can write Slang code to implement complex data structures and algorithms, or to use rasterization and ray-tracing. Slang code can be integrated directly into your Python/PyTorch training loop with SlangPy. No boilerplate, no C++.
        <a class="btn btn-primary " href="https://slangpy.shader-slang.org/">Try SlangPy</a>
    







    
        <h3>Learn How to Become a Slang Expert
            <hr>
        </h3>
        Slang code is highly portable, but can still leverage unique platform capabilities, including the latest
            features in
            Direct3D and Vulkan. For example, developers can make full use of [pointers](https://docs.shader-slang.org/en/latest/external/slang/docs/user-guide/03-convenience-features.html#pointers-limited) when generating
            SPIR-V. Slang's <a href="https://docs.shader-slang.org/en/latest/external/slang/docs/user-guide/05-capabilities.html">capability
                system</a> helps applications manage feature set differences across target platforms by ensuring
            code only uses available
            features during the type-checking step, before generating final code. Additionally, Slang provides <a
                href="https://docs.shader-slang.org/en/latest/external/slang/docs/user-guide/a1-04-interop.html">flexible
                interop</a> features to enable directly embedding target code or SPIR-V into generated shaders.
    







    

        {% assign gs_items = site.data.getting-started.items %}
        
            {% for item in gs_items %}
            {% if item.title and item.title != '' %}
            
                {{ item.icon }}
                <h4>{{ item.title }}</h4>
                {{ item.description }}
                {% if item.link_url and item.link_url != '' %}
                [{{ item.link_label }} &gt;]({{ item.link_url }})
                {% else %}
                {{ item.link_label }}
                {% endif %}
            
            {% endif %}
            {% endfor %}