Automatic Differentiation

To support differentiable graphics systems such as Gaussian splatters, neural radiance fields, differentiable path tracers, and more, Slang provides first class support for differentiable programming. An overiew:

• Slang supports the fwd_diff and bwd_diff operators that can generate the forward and backward-mode derivative propagation functions for any valid Slang function annotated with the [Differentiable] attribute.
• The DifferentialPair<T> built-in generic type is used to pass derivatives associated with each function input.
• The IDifferentiable, and the experimental IDifferentiablePtrType, interfaces denote differentiable value and pointer types respectively, and allow finer control over how types behave under differentiation.
• Futher, Slang allows for user-defined derivative functions through the [ForwardDerivative(custom_fn)] and [BackwardDerivative(custom_fn)]
• All Slang features, such as control-flow, generics, interfaces, extensions, and more are compatible with automatic differentiation, though the bottom of this chapter documents some sharp edges & known issues.

Auto-diff operations fwd_diff and bwd_diff

In Slang, fwd_diff and bwd_diff are higher-order functions used to transform Slang functions into their forward or backward derivative methods. To better understand what these methods do, here is a small refresher on differentiable calculus:

Mathematical overview: Jacobian and its vector products

Forward and backward derivative methods are two different ways of computing a dot product with the Jacobian of a given function. Parts of this overview are based on JAX’s excellent auto-diff cookbook here. The relevant wikipedia article is also a great resource for understanding auto-diff.

The Jacobian (also called the total derivative) of a function $\mathbf{f}(\mathbf{x})$ is represented by $D\mathbf{f}(\mathbf{x})$.

For a general function with multiple scalar inputs and multiple scalar outputs, the Jacobian is a matrix where $D\mathbf{f}_{ij}$ represents the partial derivative of the $i^{th}$ output element w.r.t the $j^{th}$ input element $\frac{\partial f_i}{\partial x_j}$

As an example, consider a polynomial function \(f(x, y) = x^3 + x^2 - y\) Here, $f$ here has 1 output and 2 inputs. $Df$ is therefore the row matrix: \(Df(x, y) = [\frac{\partial f}{\partial x}, \frac{\partial f}{\partial y}] = [3x^2 + 2x, -1]\)

Another, more complex example with a function that has multiple outputs (for clarity, denoted by $f_1$, $f_2$, etc..) \(\mathbf{f}(x, y) = \begin{bmatrix} f_0(x, y) & f_1(x, y) & f_2(x, y) \end{bmatrix} = \begin{bmatrix} x^3 & y^2x & y^3 \end{bmatrix}\) Here, $Df$ is a 3x2 matrix with each element containing a partial derivative: \(D\mathbf{f}(x, y) = \begin{bmatrix} \partial f_0 / \partial x & \partial f_0 / \partial y \\ \partial f_1 / \partial x & \partial f_1 / \partial y \\ \partial f_2 / \partial x & \partial f_2 / \partial y \end{bmatrix} = \begin{bmatrix} 3x^2 & 0 \\ y^2 & 2yx \\ 0 & 3y^2 \end{bmatrix}\)

Computing full Jacobians is often unnecessary and expensive. Instead, auto-diff offers ways to compute products of the Jacobian with a vector, which is a much faster operation. There are two basic ways to compute this product:

1. the Jacobian-vector product $\langle D\mathbf{f}(\mathbf{x}), \mathbf{v} \rangle$, also called forward-mode autodiff, and can be computed using fwd_diff operator in Slang, and
2. the vector-Jacobian product $\langle \mathbf{v}^T, D\mathbf{f}(\mathbf{x}) \rangle$, also called reverse-mode autodiff, and can be computed using bwd_diff operator in Slang. From a linear algebra perspective, this is the transpose of the forward-mode operator.

Propagating derivatives with forward-mode auto-diff

The products described above allow the propagation of derivatives forward and backward through the function $f$

The forward-mode derivative (Jacobian-vector product) can convert a derivative of the inputs to a derivative of the outputs. For example, lets say inputs $\mathbf{x}$ depend on some scalar $\theta$, and $\frac{\partial \mathbf{x}}{\partial \theta}$ is a vector of partial derivatives describing that dependency.

Invoking forward-mode auto-diff with $\mathbf{v} = \frac{\partial \mathbf{x}}{\partial \theta}$ converts this into a derivative of the outputs w.r.t the same scalar $\theta$. This can be verified by expanding the Jacobian and applying the chain rule of derivatives: \(\langle D\mathbf{f}(\mathbf{x}), \frac{\partial \mathbf{x}}{\partial \theta} \rangle = \langle \begin{bmatrix} \frac{\partial f_0}{\partial x_0} & \frac{\partial f_0}{\partial x_1} & \cdots \\ \frac{\partial f_1}{\partial x_0} & \frac{\partial f_1}{\partial x_1} & \cdots \\ \cdots & \cdots & \cdots \end{bmatrix}, \begin{bmatrix} \frac{\partial x_0}{\partial \theta} \\ \frac{\partial x_1}{\partial \theta} \\ \cdots \end{bmatrix} \rangle = \begin{bmatrix} \frac{\partial f_0}{\partial \theta} \\ \frac{\partial f_1}{\partial \theta} \\ \cdots \end{bmatrix} = \frac{\partial \mathbf{f}}{\partial \theta}\)

Propagating derivatives with reverse-mode auto-diff

The reverse-mode derivative (vector-Jacobian product) can convert a derivative w.r.t outputs into a derivative w.r.t inputs. For example, lets say we have some scalar $\mathcal{L}$ that depends on the outputs $\mathbf{f}$, and $\frac{\partial \mathcal{L}}{\partial \mathbf{f}}$ is a vector of partial derivatives describing that dependency.

Invoking forward-mode auto-diff with $\mathbf{v} = \frac{\partial \mathcal{L}}{\partial \mathbf{f}}$ converts this into a derivative of the same scalar $\mathcal{L}$ w.r.t the inputs $\mathbf{x}$. To provide more intuition for this, we can expand the Jacobian in a same way we did above: \(\langle \frac{\partial \mathcal{L}}{\partial \mathbf{f}}^T, D\mathbf{f}(\mathbf{x}) \rangle = \langle \begin{bmatrix}\frac{\partial \mathcal{L}}{\partial f_0} & \frac{\partial \mathcal{L}}{\partial f_1} & \cdots \end{bmatrix}, \begin{bmatrix} \frac{\partial f_0}{\partial x_0} & \frac{\partial f_0}{\partial x_1} & \cdots \\ \frac{\partial f_1}{\partial x_0} & \frac{\partial f_1}{\partial x_1} & \cdots \\ \cdots & \cdots & \cdots \end{bmatrix} \rangle = \begin{bmatrix} \frac{\partial \mathcal{L}}{\partial x_0} & \frac{\partial \mathcal{L}}{\partial x_1} & \cdots \end{bmatrix} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}}^T\)

This mode is the most popular, since machine learning systems often construct their differentiable pipeline with multiple inputs (which can number in the millions or billions), and a single scalar output often referred to as the ‘loss’ denoted by $\mathcal{L}$. The desired derivative can be constructed with a single reverse-mode invocation.

Invoking auto-diff in Slang

With the mathematical foundations established, we can describe concretely how to compute derivatives using Slang.

In Slang derivatives are computed using fwd_diff/bwd_diff which each correspond to Jacobian-vector and vector-Jacobian products. For forward-diff, to pass the vector $\mathbf{v}$ and receive the outputs, we use the DifferentialPair<T> type. We use pairs of inputs because every input element $x_i$ has a corresponding element $v_i$ in the vector, and each original output element has a corresponding output element in the product.

Example of fwd_diff:

[Differentiable] // Auto-diff requires that functions are marked differentiable
float2 foo(float a, float b) 
{ 
    return float2(a * b * b, a * a);
}

void main()
{
    DifferentialPair<float> dp_a = diffPair(
        1.0, // input 'a'
        1.0  // vector 'v' for vector-Jacobian product input (for 'a')
    );

    DifferentialPair<float> dp_b = diffPair(2.4, 0.0);

    // fwd_diff to compute output and d_output w.r.t 'a'.
    // Our output is also a differential pair.
    //
    DifferentialPair<float2> dp_output = fwd_diff(foo)(dp_a, dp_b);

    // Extract output's primal part, which is just the standard output when foo is called normally.
    // Can also use `.getPrimal()`
    //
    float2 output_p = dp_output.p;

    // Extract output's derivative part. Can also use `.getDifferential()`
    float2 output_d = dp_output.d;

    printf("foo(1.0, 2.4) = (%f %f)\n", output_p.x, output_p.y);
    printf("d(foo)/d(a) at (1.0, 2.4) = (%f, %f)\n", output_d.x, output_d.y);
}

Note that all the inputs and outputs to our function become ‘paired’. This only applies to differentiable types, such as float, float2, etc. See the section on differentiable types for more info.

diffPair<T>(primal_val, diff_val) is a built-in utility function that constructs the pair from the primal and differential values.

Additionally, invoking forward-mode also computes the regular (or ‘primal’) output value (can be obtained from output.getPrimal() or output.p). The same is not true for reverse-mode.

For reverse-mode, the example proceeds in a similar way, and we still use DifferentialPair<T> type. However, note that each input gets a corresponding output and each output gets a corresponding input. Thus, all inputs become inout differential pairs, to allow the function to write into the derivative part (the primal part is still accepted as an input in the same pair data-structure). The one extra rule is that the derivative corresponding to the return value of the function is accepted as the last argument (an extra input). This value does not need to be a pair.

Example:

[Differentiable] // Auto-diff requires that functions are marked differentiable
float2 foo(float a, float b) 
{ 
    return float2(a * b * b, a * a);
}

void main()
{
    DifferentialPair<float> dp_a = diffPair(
        1.0 // input 'a'
    ); // Calling diffPair without a derivative part initializes to 0.

    DifferentialPair<float> dp_b = diffPair(2.4);

    // Derivatives of scalar L w.r.t output.
    float2 dL_doutput = float2(1.0, 0.0);

    // bwd_diff to compute dL_da and dL_db
    // The derivative of the output is provided as an additional _input_ to the call
    // Derivatives w.r.t inputs are written into dp_a.d and dp_b.d
    //
    bwd_diff(foo)(dp_a, dp_b, dL_doutput);

    // Extract the derivatives of L w.r.t input
    float dL_da = dp_a.d;
    float dL_db = dp_b.d;

    printf("If dL/dOutput = (1.0, 0.0), then (dL/da, dL/db) at (1.0, 2.4) = (%f, %f)", dL_da, dL_db);
}

Differentiable Type System

Slang will only generate differentiation code for values that has a differentiable type. Differentiable types are defining through conformance to one of two built-in interfaces:

1. IDifferentiable: For value types (e.g. float, structs of value types, etc..)
2. IDifferentiablePtrType: For buffer, pointer & reference types that represent locations rather than values.

Differentiable Value Types

All basic types (float, int, double, etc..) and all aggregate types (i.e. struct) that use any combination of these are considered value types in Slang.

Slang uses the IDifferentiable interface to define differentiable types. Basic types that describe a continuous value (float, double and half) and their vector/matrix versions (float3, half2x2, etc..) are defined as differentiable by the standard library. For all basic types, the type used for the differential (can be obtained with T.Differential) is the same as the primal.

Builtin Differentiable Value Types

The following built-in types are differentiable:

• Scalars: float, double and half.
• Vector/Matrix: vector and matrix of float, double and half types.
• Arrays: T[n] is differentiable if T is differentiable.
• Tuples: Tuple<each T> is differentiable if T is differentiable.

User-defined Differentiable Value Types

However, it is easy to define your own differentiable types. Typically, all you need is to implement the IDifferentiable interface.

struct MyType : IDifferentiable
{
    float x;
    float y;
};

The main requirement of a type implementing IDifferentiable is the Differential associated type that the compiler uses to carry the corresponding derivative. In most cases the Differential of a type can be itself, though it can be different if necessary. You can access the differential of any differentiable type through Type.Differential

Example:

MyType obj;
obj.x = 1.f;

MyType.Differential d_obj;
// Differentiable fields will have a corresponding field in the diff type
d_obj.x = 1.f;

Slang can automatically derive the Differential type in the majority of cases. For instance, for MyType, Slang can infer the differential trivially:

struct MyType : IDifferentiable
{
    // Automatically inserted by Slang from the fact that 
    // MyType has 2 floats which are both differentiable
    //
    typealias Differential = MyType;
    // ...
}

For more complex types that aren’t fully differentiable, a new type is synthesized automatically:

struct MyPartialDiffType : IDifferentiable
{
    // Automatically inserted by Slang based on which fields are differentiable.
    typealias MyPartialDiffType = syn_MyPartialDiffType_Differential;
    
    float x;
    uint y;
};

// Synthesized
struct syn_MyPartialDiffType_Differential
{
    // Only one field since 'y' does not conform to IDifferentiable
    float x;
};

You can make existing types differentiable through Slang’s extension mechanism. For instance, extension MyType : IDifferentiable { } will make MyType differentiable retroactively.

See the IDifferentiable reference documentation for more information on how to override the default behavior.

DifferentialPair: Pairs of differentiable value types

The DifferentialPair<T> type is used to pass derivatives to a derivative call by representing a pair of values of type T and T.Differential. Note that T must conform to IDifferentiable.

DifferentialPair<T> can either be created via constructor calls or the diffPair utility method.

Example:

MyType obj = {1.f, 2.f};

MyType.Differential d_obj = {0.4f, 3.f};

// The differential part of a differentiable-pair is of the diff type.
DifferentialPair<MyType> dp_obj = diffPair(obj, d_obj);

// Use .p to extract the primal part
MyType new_p_obj = dp_obj.p;

// Use .d to extract the differential part
MyType.Differential new_d_obj = dp_obj.d;

Differentiable Ptr types

Pointer types are any type that represents a location or reference to a value rather than the value itself. Examples include resource types (RWStructuredBuffer, Texture2D), pointer types (Ptr<float>) and references.

The IDifferentiablePtrType interface can be used to denote types that need to transform into pairs during auto-diff. However, unlike an IDifferentiable type whose derivative portion is an output under bwd_diff, the derivative part of IDifferentiablePtrType remains an input. This is because only the value is returned as an output, while the location where it needs to be written to, is still effectively an input to the derivative methods.

Note

Support for IDifferentiablePtrType is still experimental. There are no built-in types conforming to this interface, though we plan to add stdlib support in the near future.

IDifferentiablePtrType only requires a Differential associated type to be specified.

DifferentialPtrPair: Pairs of differentiable ptr types

For types conforming to IDifferentiablePtrType, the corresponding pair to use for passing the derivative counterpart is DifferentialPtrPair<T>, which represents a pair of T and T.Differential. Objects of this type can be created using a constructor.

Example of defining and using an IDifferentiablePtrType object.

Here is an example of create a differentiable buffer pointer type, and using it within a differentiable function. You can find an interactive sample on the Slang playground here.

struct MyBufferPointer : IDifferentiablePtrType
{
    // The differential part is another instance of MyBufferPointer.
    typealias Differential = MyBufferPointer;

    RWStructuredBuffer<float> buf;
    uint offset;
};

// Link a custom derivative
[BackwardDerivative(load_bwd)]
float load(MyBufferPointer p, uint index)
{
    return p.buf[p.offset + index];
}

// Note that the backward derivative signature is still an 'in' differential pair.
void load_bwd(DifferentialPtrPair<MyBufferPointer> p, uint index, float dOut)
{
    MyBufferPointer diff_ptr = p.d;
    diff_ptr.buf[diff_ptr.offset + index] += dOut;
}

[Differentiable]
float sumOfSquares<let N : int>(MyBufferPointer p)
{
    float sos = 0.f;

    [MaxIters(N)]
    for (uint i = 0; i < N; i++)
    {
        float val_i = load(p, i);
        sos += val_i * val_i;
    }

    return sos;
}

RWStructuredBuffer<float> inputs;
RWStructuredBuffer<float> derivs;

void main()
{
    MyBufferPointer ptr = {inputs, 0};
    print("Sum of squares of first 10 values: ", sumOfSquares<10>(ptr));

    MyBufferPointer deriv_ptr = {derivs, 0};

    // Pass a pair of pointers as input.
    bwd_diff(sumOfSquares<10>)(
        DifferentialPtrPair<MyBufferPointer>(ptr, deriv_ptr),
        1.0);
    
    print("Derivative of result w.r.t the 10 values: \n");
    for (uint i = 0; i < 10; i++)
        print("%d: %f\n", i, load(deriv_ptr, i));
}

User-Defined Derivative Functions

As an alternative to compiler-generated derivatives, you can choose to provide an implementation for the derivative, which the compiler will use instead of attempting to generate one.

This can be performed on a per-function basis by using the decorators [ForwardDerivative(fwd_deriv_func)] and [BackwardDerivative(bwd_deriv_func)] to reference the derivative from the primal function.

For instance, it often makes little sense to differentiate the body of a sin(x) implementation, when we know that the derivative is cos(x) * dx. In Slang, this can be represented in the following way:

DifferentialPair<float> sin_fwd(DifferentialPair<float> dpx)
{
    float x = dpx.p;
    float dx = dpx.d;
    return DifferentialPair<float>(dpx.p, cos(x) * dx);
}

// sin() is now considered differentiable (atleast for forward-mode) since it provides
// a derivative implementation.
//
[ForwardDerivative(sin_fwd)]
float sin(float x)
{
    // Calc sin(X) using Taylor series..
}

// Any uses of sin() in a `[Differentiable]` will automaticaly use the sin_fwd implementation when differentiated.

A similar example for a backward derivative.

void sin_bwd(inout DifferentialPair<float> dpx, float dresult)
{
    float x = dpx.p;

    // Write-back the derivative to each input (the primal part must be copied over as-is)
    dpx = DifferentialPair<float>(x, cos(x) * dresult);
}

[BackwardDerivative(sin_bwd)]
float sin(float x)
{
    // Calc sin(X) using Taylor series..
}

Note that the signature of the provided forward or backward derivative function must match the expected signature from invoking fwd_diff(fn)/bwd_diff(fn) For a full list of signature rules, see the reference section for the auto-diff operators.

Back-referencing User Derivative Attributes.

Sometimes, the original function’s definition might be inaccessible, so it can be tricky to add an attribute to create the association.

For such cases, Slang provides the [ForwardDerivativeOf(primal_fn)] and [BackwardDerivativeOf(primal_fn)] attributes that can be used on the derivative function and contain a reference to the function for which they are providing a derivative implementation. As long as both the derivative function is in scope, the primal function will be considered differentiable.

Example:

// Module A
float sin(float x) { /* ... */ } 

// Module B
import A;
[BackwardDerivativeOf(sin)] // Add a derivative implementation for sin() in module A.
void sin_bwd(inout DifferentialPair<float> dpx, float dresult) { /* ... */ }

User-defined derivatives also work for generic functions, member functions, accessors, and more. See the reference section for the [ForwardDerivative(fn)] and [BackwardDerivative(fn)] attributes for more.

Using Auto-diff with Generics

Automatic differentiation works seamlessly with generically-defined types and methods. For generic methods, differentiability of a type is defined either through an explicit IDifferentiable constraint or any other interface that extends IDifferentiable.

Example for generic methods:

[Differentiable]
T calcFoo<T : IDifferentiable>(T x) { /* ... */ }

[Differentiable]
T calcBar<T : __BuiltinFloatingPointType>(T x) { /* ... */ }

[Differentiable]
void main()
{
    DifferentialPair<float4> dpa = /* ... */;

    // Can call with any type that is IDifferentiable. Generic parameters
    // are inferred like any other call.
    //
    bwd_diff(calcFoo)(dpa, float4(1.f));

    // But you can also be explicit with < >
    bwd_diff(calcFoo<float4>)(dpa, float4(1.f));

    // x is differentiable for calcBar because 
    // __BuiltinFloatingPointType : IDifferentiable
    //
    DifferentialPair<double> dpb = /* .. */;
    bwd_diff(calcBar)(dpb, 1.0);
}

You can implement IDifferentiable on a generic type. Automatic synthesis still applies and will use generic constraints to resolve whether a field is differentiable or not.

struct Foo<T : IDifferentiable, U> : IDifferentiable
{
    T t;
    U u;
};

// The synthesized Foo<T, U>.Differential will contain a field for
// 't' but not 'U'
//

Using Auto-diff with Interface Requirements and Interface Types

For interface requirements, using [Differentiable] attribute enforces that any implementation of that method must also be differentiable. You can, of course, provide a manual derivative implementation to satisfy the requirement.

The following is a sample snippet. You can run the full sample on the playground here.

interface IFoo
{
    [Differentiable]
    float calc(float x);
}

struct FooImpl : IFoo
{
    // Implementation via automatic differentiation.
    [Differentiable]
    float calc(float x)
    { /* ... */ }
}

struct FooImpl2 : IFoo
{
    // Implementation via manually providing derivative methods.
    [ForwardDerivative(calc_fwd)]
    [BackwardDerivative(calc_bwd)]
    float calc(float x)
    { /* ... */ }

    DifferentialPair<float> calc_fwd(DifferentialPair<float> x)
    { /* ... */ }

    void calc_bwd(inout DifferentialPair<float> x, float dresult)
    { /* ... */ }
}

[Differentiable]
float compute(float x, uint obj_id)
{
    // Create an instance of either FooImpl1 or FooImpl2
    IFoo foo = createDynamicObject<IFoo>(obj_id); 
    
    // Dynamic dispatch to appropriate 'calc'.
    //
    // Note that foo itself is non-differentiable, and 
    // has no differential data, but 'x' and 'result'
    // will carry derivatives.s
    //
    var result = foo.calc(x);
    return result;
}

Differentiable Interface (and Associated) Types

Note: This is an advanced use-case and support is currently experimental.

You can have an interface or an interface associated type extend IDifferentiable and use that in differentiable interface requirement functions. This is often important in large code-bases with modular components that are all differentiable (one example is the material system in large production renderers)

Here is a snippet of how to make an interface and associated type (and by consequence all its implementations) differentiable. For a full working sample, check out the Slang playground here

interface IFoo : IDifferentiable
{
    associatedtype BaseType : IDifferentiable;

    [Differentiable]
    BaseType foo(BaseType x);
};

[Differentiable]
float calc(float x)
{
    // Note that since IFoo is differentiable, 
    // any data in the IFoo implementation is differentiable
    // and will carry derivatives.
    //
    IFoo obj = makeObj(/* ... */);
    return obj.foo(x);
}

Under the hood, Slang will automatically construct an anonymous abstract type to represent the differentials. However, on targets that don’t support true dynamic dispatch, these are lowered into tagged unions. While we are working to improve the implementation, this union can currently include all active differential types, rather than just the relevant ones. This can lead to increased memory use.

Primal Substitute Functions

Sometimes it is desirable to replace a function with another when generating derivative code. Most often, this is because a lot of shader operations may just not have a function body, such hardware intrinsics for texture sampling. In such cases, Slang provides a [PrimalSubstitute(fn)] attribute that can be used to provide a reference implementation that Slang can differentiate to generate the derivative function.

The following is a small snippet with bilinear texture sampling. For a full example application that uses this concept, see the texture differentiation sample in the Slang repository.

[PrimalSubstitute(sampleTextureBiliear_reference)]
float4 sampleTextureBilinear(Texture2D<float4> x, float2 loc) 
{ 
    // HW-accelerated sampling intrinsics. 
    // Slang does not have access to body, so cannot differentiate.
    //
    x.Sample(/*...*/)
}

// Since the substitute is differentiable, so is `sampleTextureBilinear`.
[Differentiable]
float4 sampleTextureBilinear_reference(Texture2D<float4> x, float2 loc)
{
    // Reference SW interpolation, that is differentiable.
}

[Differentiable]
float computePixel(Texture2D<float> x, float a, float b)
{
    // Slang will use HW-accelerated sampleTextureBilinear for standard function
    // call, but differentiate the SW reference interpolation during backprop.
    // 
    float4 sample1 = sampleTextureBilinear(x, float2(a, 1));
}

Similar to [ForwardDerivativeOf(fn)] and [BackwardDerivativeOf(fn)] attributes, Slang provides a [PrimalSubstituteOf(fn)] attribute that can be used on the substitute function to reference the primal one.

Working with Mixed Differentiable and Non-Differentiable Code

Introducing differentiability to an existing system often involves dealing with code that mixes differentiable and non-differentiable logic. Slang provides type checking and code analysis features to allow users to clarify the intention and guard against unexpected behaviors involving when to propagate derivatives through operations.

Excluding Parameters from Differentiation

Sometimes we do not wish a parameter to be considered differentiable despite it has a differentiable type. We can use the no_diff modifier on the parameter to inform the compiler to treat the parameter as non-differentiable and skip generating differentiation code for the parameter. The syntax is:

// Only differentiate this function with regard to `x`.
float myFunc(no_diff float a, float x);

The forward derivative and backward propagation functions of myFunc should have the following signature:

DifferentialPair<float> fwd_derivative(float a, DifferentialPair<float> x);
void back_prop(float a, inout DifferentialPair<float> x, float dResult);

In addition, the no_diff modifier can also be used on the return type to indicate the return value should be considered non-differentiable. For example, the function

no_diff float myFunc(no_diff float a, float x, out float y);

Will have the following forward derivative and backward propagation function signatures:

float fwd_derivative(float a, DifferentialPair<float> x);
void back_prop(float a, inout DifferentialPair<float> x, float d_y);

By default, the implicit this parameter will be treated as differentiable if the enclosing type of the member method is differentiable. If you wish to exclude this parameter from differentiation, use [NoDiffThis] attribute on the method:

struct MyDifferentiableType : IDifferentiable
{
    [NoDiffThis]   // Make `this` parameter `no_diff`.
    float compute(float x) { ... }
}

Excluding Struct Members from Differentiation

When using automatic IDifferentiable conformance synthesis for a struct type, Slang will by-default treat all struct members that have a differentiable type as differentiable, and thus include a corresponding field in the generated Differential type for the struct. For example, given the following definition

struct MyType : IDifferentiable
{
    float member1;
    float2 member2;
}

Slang will generate:

struct MyType.Differential : IDifferentiable
{
    float member1;  // derivative for MyType.member1
    float2 member2; // derivative for MyType.member2
}

If the user does not want a certain member to be treated as differentiable despite it has a differentiable type, a no_diff modifier can be used on the struct member to exclude it from differentiation. For example, the following code excludes member1 from differentiation:

struct MyType : IDifferentiable
{
    no_diff float member1;  // excluded from differentiation
    float2 member2;
}

The generated Differential in this case will be:

struct MyType.Differential : IDifferentiable
{
    float2 member2;
}

Assigning Differentiable Values into a Non-Differentiable Location

When a value with derivatives is being assigned to a location that is not differentiable, such as a struct member that is marked as no_diff, the derivative info is discarded and any derivative propagation is stopped at the assignment site. This may lead to unexpected results. For example:

struct MyType : IDifferentiable
{
    no_diff float member;
    float someOtherMember;
}
[Differentiable]
float f(float x)
{
    MyType t;
    t.member = x * x; // Error: assigning value with derivative into a non-differentiable location.
    return t.member;
}
...
let result = fwd_diff(f)(diffPair(3.0, 1.0)).d; // result == 0.0

In this case, we are assigning the value x*x, which carries a derivative, into a non-differentiable location MyType.member, thus throwing away any derivative info. When f returns t.member, there will be no derivative associated with it, so the function will not propagate the derivative through. This code is most likely not intending to discard the derivative through the assignment. To help avoid this kind of unintentional behavior, Slang will treat any assignments of a value with derivative info into a non-differentiable location as a compile-time error. To eliminate this error, the user should either make t.member differentiable, or to force the assignment by clarifying the intention to discard any derivatives using the built-in detach method. The following code will compile, and the derivatives will be discarded:

[Differentiable]
float f(float x)
{
    MyType t;
    // OK: the code has expressed clearly the intention to discard the derivative and perform the assignment.
    t.member = detach(x * x);
    return t.member;
}

Calling Non-Differentiable Functions from a Differentiable Function

Calling non-differentiable function from a differentiable function is allowed. However, derivatives will not be propagated through the call. The user is required to clarify the intention by prefixing the call with the no_diff keyword. An un-clarified call to non-differentiable function will result in a compile-time error.

For example, consider the following code:

float g(float x)
{
    return 2*x;
}

[Differentiable]
float f(float x)
{
    // Error: implicit call to non-differentiable function g.
    return g(x) + x * x;
}

The derivative will not propagate through the call to g in f. As a result, fwd_diff(f)(diffPair(1.0, 1.0)) will return {3.0, 2.0} instead of {3.0, 4.0} as the derivative from 2*x is lost through the non-differentiable call. To prevent unintended error, it is treated as a compile-time error to call g from f. If such a non-differentiable call is intended, a no_diff prefix is required in the call:

[Differentiable]
float f(float x)
{
    // OK. The intention to call a non-differentiable function is clarified.
    return no_diff g(x) + x * x;
}

However, the no_diff keyword is not required in a call if a non-differentiable function does not take any differentiable parameters, or if the result of the differentiable function is not dependent on the derivative being propagated through the call.

Treat Non-Differentiable Functions as Differentiable

Slang allows functions to be marked with a [TreatAsDifferentiable] attribute for them to be considered as differentiable functions by the type-system. When a function is marked as [TreatAsDifferentiable], the compiler will not generate derivative propagation code from the original function body or perform any additional checking on the function definition. Instead, it will generate trivial forward and backward propagation functions that returns 0.

This feature can be useful if the user marked an interface method as forward or backward differentiable, but only wish to provide non-trivial derivative propagation functions for a subset of types that implement the interface. For other types that does not actually need differentiation, the user can simply put [TreatAsDifferentiable] on the method implementations for them to satisfy the interface requirement.

See the following code for an example of [TreatAsDifferentiable]:

interface IFoo
{
    [Differentiable]
    float f(float v);
}

struct B : IFoo
{
    [TreatAsDifferentiable]
    float f(float v)
    {
        return v * v;
    }
}

[Differentiable]
float use(IFoo o, float x)
{
    return o.f(x);
}

// Test:
B obj;
float result = fwd_diff(use)(obj, diffPair(2.0, 1.0)).d;
// result == 0.0, since `[TreatAsDifferentiable]` causes a trivial derivative implementation
// being generated regardless of the original code.

Higher-Order Differentiation

Slang supports generating higher order forward and backward derivative propagation functions. It is allowed to use fwd_diff and bwd_diff operators inside a forward or backward differentiable function, or to nest fwd_diff and bwd_diff operators. For example, fwd_diff(fwd_diff(sin)) will have the following signature:

DifferentialPair<DifferentialPair<float>> sin_diff2(DifferentialPair<DifferentialPair<float>> x);

The input parameter x contains four fields: x.p.p, x.p.d,, x.d.p, x.d.d, where x.p.p specifies the original input value, both x.p.d and x.d.p store the first order derivative if x, and x.d.d stores the second order derivative of x. Calling fwd_diff(fwd_diff(sin)) with diffPair(diffPair(pi/2, 1.0), DiffPair(1.0, 0.0)) will result { { 1.0, 0.0 }, { 0.0, -1.0 } }.

User defined higher-order derivative functions can be specified by using [ForwardDerivative] or [BackwardDerivative] attribute on the derivative function, or by using [ForwardDerivativeOf] or [BackwardDerivativeOf] attribute on the higher-order derivative function.

Restrictions and Known Issues

The compiler can generate forward derivative and backward propagation implementations for most uses of array and struct types, including arbitrary read and write access at dynamic array indices, and supports uses of all types of control flows, mutable parameters, generics and interfaces. This covers the set of operations that is sufficient for a lot of functions. However, the user needs to be aware of the following restrictions when using automatic differentiation:

• All operations to global resources, global variables and shader parameters, including texture reads or atomic writes, are treated as a non-differentiable operation. Slang provides support for special data-structures (such as Tensor) through libraries such as SlangPy, which come with custom derivative implementations
• If a differentiable function contains calls that cause side-effects such as updates to global memory, there is currently no guarantee on how many times side-effects will occur during the resulting derivative function or back-propagation function.
• Loops: Loops must have a bounded number of iterations. If this cannot be inferred statically from the loop structure, the attribute [MaxIters(<count>)] can be used specify a maximum number of iterations. This will be used by compiler to allocate space to store intermediate data. If the actual number of iterations exceeds the provided maximum, the behavior is undefined. You can always mark a loop with the [ForceUnroll] attribute to instruct the Slang compiler to unroll the loop before generating derivative propagation functions. Unrolled loops will be treated the same way as ordinary code and are not subject to any additional restrictions.
• Double backward derivatives (higher-order differentiation): The compiler does not currently support multiple backward derivative calls such as bwd_diff(bwd_diff(fn)). The vast majority of higher-order derivative applications can be acheived more efficiently via multiple forward-derivative calls or a single layer of bwd_diff on functions that use one or more fwd_diff passes.

The above restrictions do not apply if a user-defined derivative or backward propagation function is provided.

Reference

This section contains some additional information for operators that are not currently included in the standard library reference

fwd_diff(f : slang_function) -> slang_function

The fwd_diff operator can be used on a differentiable function to obtain the forward derivative propagation function.

A forward derivative propagation function computes the derivative of the result value with regard to a specific set of input parameters. Given an original function, the signature of its forward propagation function is determined using the following rules:

• If the return type R implements IDifferentiable the forward propagation function will return a corresponding DifferentialPair<R> that consists of both the computed original result value and the (partial) derivative of the result value. Otherwise, the return type is kept unmodified as R.
• If a parameter has type T that implements IDifferentiable, it will be translated into a DifferentialPair<T> parameter in the derivative function, where the differential component of the DifferentialPair holds the initial derivatives of each parameter with regard to their upstream parameters.
• If a parameter has type T that implements IDifferentiablePtrType, it will be translated into a DifferentialPtrPair<T> parameter where the differential component references the differential component.
• All parameter directions are unchanged. For example, an out parameter in the original function will remain an out parameter in the derivative function.
• Differentiable methods cannot have a type implementing IDifferentiablePtrType as an out or inout parameter, or a return type. Types implementing IDifferentiablePtrType can only be used for input parameters to a differentiable method. Marking such a method as [Differentiable] will result in a compile-time diagnostic error.

For example, given original function:

[Differentiable]
R original(T0 p0, inout T1 p1, T2 p2, T3 p3);

Where R, T0, T1 : IDifferentiable, T2 is non-differentiable, and T3 : IDifferentiablePtrType, the forward derivative function will have the following signature:

DifferentialPair<R> derivative(DifferentialPair<T0> p0, inout DifferentialPair<T1> p1, T2 p2, DifferentialPtrPair<T3> p3);

This forward propagation function takes the initial primal value of p0 in p0.p, and the partial derivative of p0 with regard to some upstream parameter in p0.d. It takes the initial primal and derivative values of p1 and updates p1 to hold the newly computed value and propagated derivative. Since p2 is not differentiable, it remains unchanged.

bwd_diff(f : slang_function) -> slang_function

A backward derivative propagation function propagates the derivative of the function output to all the input parameters simultaneously.

Given an original function f, the general rule for determining the signature of its backward propagation function is that a differentiable output o becomes an input parameter holding the partial derivative of a downstream output with regard to the differentiable output, i.e. $\partial y/\partial o$; an input differentiable parameter i in the original function will become an output in the backward propagation function, holding the propagated partial derivative $\partial y/\partial i$; and any non-differentiable outputs are dropped from the backward propagation function. This means that the backward propagation function never returns any values computed in the original function.

More specifically, the signature of its backward propagation function is determined using the following rules:

• A backward propagation function always returns void.
• A differentiable in parameter of type T : IDifferentiable will become an inout DifferentialPair<T> parameter, where the original value part of the differential pair contains the original value of the parameter to pass into the back-prop function. The original value will not be overwritten by the backward propagation function. The propagated derivative will be written to the derivative part of the differential pair after the backward propagation function returns. The initial derivative value of the pair is ignored as input.
• A differentiable out parameter of type T : IDifferentiable will become an in T.Differential parameter, carrying the partial derivative of some downstream term with regard to the return value.
• A differentiable inout parameter of type T : IDifferentiable will become an inout DifferentialPair<T> parameter, where the original value of the argument, along with the downstream partial derivative with regard to the argument is passed as input to the backward propagation function as the original and derivative part of the pair. The propagated derivative with regard to this input parameter will be written back and replace the derivative part of the pair. The primal value part of the parameter will not be updated.
• A differentiable return value of type R will become an additional in R.Differential parameter at the end of the backward propagation function parameter list, carrying the result derivative of a downstream term with regard to the return value of the original function.
• A non-differentiable return value of type NDR will be dropped.
• A non-differentiable in parameter of type ND will remain unchanged in the backward propagation function.
• A non-differentiable out parameter of type ND will be removed from the parameter list of the backward propagation function.
• A non-differentiable inout parameter of type ND will become an in ND parameter.
• Types implemented IDifferentiablePtrType work the same was as the forward-mode case. They can only be used with in parameters, and are converted into DifferentialPtrPair types. Their directions are not affected.

For example consider the following original function:

struct T : IDifferentiable {...}
struct R : IDifferentiable {...}
struct P : IDifferentiablePtrType {...}
struct ND {} // Non differentiable

[Differentiable]
R original(T p0, out T p1, inout T p2, ND p3, out ND p4, inout ND p5, P p6);

The signature of its backward propagation function is:

void back_prop(
    inout DifferentialPair<T> p0,
    T.Differential p1,
    inout DifferentialPair<T> p2,
    ND p3,
    ND p5,
    DifferentialPtrPair<P> p6,
    R.Differential dResult);

Note that although p2 is still inout in the backward propagation function, the backward propagation function will only write propagated derivative to p2.d and will not modify p2.p.

Built-in Differentiable Functions

The following built-in functions are differentiable and both their forward and backward derivative functions are already defined in the standard library’s core module:

• Arithmetic functions: abs, max, min, sqrt, rcp, rsqrt, fma, mad, fmod, frac, radians, degrees
• Interpolation and clamping functions: lerp, smoothstep, clamp, saturate
• Trigonometric functions: sin, cos, sincos, tan, asin, acos, atan, atan2
• Hyperbolic functions: sinh, cosh, tanh
• Exponential and logarithmic functions: exp, exp2, pow, log, log2, log10
• Vector functions: dot, cross, length, distance, normalize, reflect, refract
• Matrix transforms: mul(matrix, vector), mul(vector, matrix), mul(matrix, matrix)
• Matrix operations: transpose, determinant
• Legacy blending and lighting intrinsics: dst, lit