Differentiable Programming with Slang

[saipraveenb25]

Differentiable Programming with Slang

This page is a living documentation of the currently available auto-diff features in Slang, and how to use them. We will be updating this as new features are added/removed

Forward-mode auto-diff with __fwd_diff(fn)

You can invoke the forward-mode derivative version of an arbitrary function by wrapping it with __fwd_diff().

As an example, lets take a simple 1-input, 1-output function that computes the square

float sqr(float x)
{
    return x*x;
}

The forward-mode derivative computes the differential of the output, given the differentials of it's inputs.

To obtain the forward-mode derivative of sqr, we first mark sqr as differentiable by adding a [ForwardDifferentiable] attribute to the signature, and then invoke __fwd_diff(sqr).

The new function requires differentials in addition to the original inputs. In Slang, we enable this by transforming each input into a pair of itself and its differential.

We achieve this through the generic struct DifferentialPair<T> which holds both the orginal value and the differential. The former can be accessed using .p() or .getPrimal(), and the latter can be accessed through .d() or .getDifferential()

The synthesized function looks like the following:

// Synthesized automatically by the compiler.
DifferentialPair<float> d_sqr(DifferentialPair<float> x)
{
    return 2 * x.p() * x.d(); // 2 * x * dx
}

Here's a complete example:

// Use an alias for brevity.
typedef DifferentialPair<float> dpfloat;

[ForwardDifferentiable]
float sqr(float x)
{
    return x*x;
}

void main()
{
        float x = 2.0;
        float dx = 1.0;
        
        // Call `sqr` normally. We get 4.0 here.
        float sqr_of_x = sqr(2.0);

        // Form a pair of x and dx to provide as input to the
        // differential function.
        dpfloat pair_of_x_and_dx = dpfloat(x, dx);

        // Invoke __fwd_diff(sqr), this will translate internally to the pair (x*x,  2 * x * dx), 
        // For x=2.0, dx=1.0, we get the pair 4.0, 2.0
        dpfloat dp_of_x = __fwd_diff(sqr)(pair_of_x_and_dx); 
        
       // Extract just the differential part (2.0)
       float d_sqr_of_x = dp_of_x.d();
}

Basic Differentiable Types

Currently, the following types will be detected as being automatically differentiable:

1. Floating-point type: float
2. Built-in Floating-point vector types float2, float3, and float4 (but not for general N-sized vectors)

More type support is being actively added.

Custom Differentiable Types

Define a type as differentiable by extending the IDifferentiable interface. In most cases, it is sufficient to just write:

struct MyType : IDifferentiable
{
    // Rest stays the same
}

And the compiler will automatically synthesis all the code required to implement IDifferentiable interface.

The definition of IDifferentiable interface has four requirements:

interface IDifferentiable
{
      associatedtype Differential : IDifferentiable;
      Differential dzero();
      Differential dadd(Differential, Differential);
      Differential dmul(This, Differential);
}

The Differential associatedtype specifies the type that represents the result of differentiation of value of this type.
The dzero method defines the 0 value of the differential type.
The dadd method defines the add operation on two differential values.
The dmul method defines the inner product of a primal value and differential value.

In order to use a member of a differential struct type, that member must be marked with a [DerivativeMember] attribute. Members without [DerivativeMember] attribute will be treated as they are non-differentiable. However if the user does not specify the Differential associatedtype, the compiler will automatically synthesize one that includes all the differentiable members, and automatically insert [DerivativeMember] attribute.

The user can make use of automatic IDifferentiable conformance synthesis by writing the following:

struct A : IDifferentiable
{
      // Non-differentiable field.
      int index;
      // A field whose type is differentiable will be included in the generated `Differential` type.
      float3 value;
};

And the compiler will automatically synthesize necessary definitions within the struct and the result is equivalent to the user writing the following code:

struct A : IDifferentiable
{
      struct Differential : IDifferentiable      // Compiler synthesized
      {
          float3.Differential value;   // float3.Differential is just float3.
      };
      
      // Non-differentiable field.
      int index;
      
      // Define differentiable members using a property
      [DerivativeMember(Differential.value)]    // Compiler synthesized
      float3 value;
      
      Differential dzero() { return {0.0;}}   // Compiler synthesized
      
      Differential dadd(Differential a, Differential b) { return {a.value + b.value}; }  // Compiler synthesized
      
      Differential dmul(This a, Differential b) {return {a.value * b.value}; } // Compiler synthesized
};

Custom Forward-mode functions

Sometimes, it is desirable to have a custom hand-written derivative for a function, rather than synthesizing the derivative
The attribute [ForwardDerivative(<derivative-fn-name>)] can be used to acheive this.

A complete example of this:

// Use an alias for brevity.
typedef DifferentialPair<float> dpfloat;

dpfloat d_sqr(dpfloat dpx)
{
     float x = dpx.p();
     float dx = dpx.d();

     return dpfloat(x * x, 2 * x * dx);
}

[ForwardDerivative(d_sqr)]
float sqr(float x)
{
    return x*x;
}

void main()
{
        float x = 2.0;
        float dx = 1.0;
        
        // Call `sqr` normally. We get 4.0 here.
        float sqr_of_x = sqr(2.0);

        // Form a pair of x and dx to provide as input to the
        // differential function.
        dpfloat pair_of_x_and_dx = dpfloat(x, dx);

        // Invoke __fwd_diff(sqr), this will translate internally to the pair (x*x,  2 * x * dx), 
        // For x=2.0, dx=1.0, we get the pair 4.0, 2.0
        dpfloat dp_of_x = __fwd_diff(sqr)(pair_of_x_and_dx); 
        
       // Extract just the differential part (2.0)
       float d_sqr_of_x = dp_of_x.d();
}

Custom derivative functions can also be provided with the [ForwardDerivativeOf] attribute, which works the opposite way as [ForwardDerivative]. Instead of decorating the original function, [ForwardDerivativeOf] attribute decorates the derivative function and associate it with an original function. The previous example can also be written as following:

float sqr(float x)
{
    return x*x;
}

[ForwardDerivativeOf(sqr)]
dpfloat d_sqr(dpfloat dpx)
{
     float x = dpx.p();
     float dx = dpx.d();

     return dpfloat(x * x, 2 * x * dx);
}

Calling Non-Differentiable Functions from a Differnetiable Function

Calling non-differentiable functions from a differentiable function is allowed, as long as the result of the non-differentiable function call does not directly contribute to any differentiable result (control-flow dependence doesn't count as contributing).
It is a compile-time error if the result of a non-differentiable function call is used to produce the a differentiable result value.
For example:

float nonDiff(float x)
{
    return x;
}

[ForwardDifferentiable]
float f(float x)
{
    float val = 0;
    if (x > 5)
        val = x + 1;
    else
        val = nonDiff(x * 2.0f); // call to non-differentiable function causes loss of derivative.
    // Error: not all path propagates derivatives through.
    return val;
}

However, this compile error can be supressed with the no_diff keyword to clarify that the non-differentiable call is intentional.

[ForwardDifferentiable]
float f(float x)
{
    float val = 0;
    if (x > 5)
        val = x + 1;
    else
        val = no_diff nonDiff(x * 2.0f); // OK!
    return val;
}

Standard library functions

Most floating point math functions has builtin derivative implementation. Calls to these functions will be differentiated automatically. The following code is an example of calling stdlib function in a differentiable function.

[ForwardDifferentiable]
float f(float x)
{
    return exp(x);
}

The following stdlib functions have derivative implementations: log, sin, cos, exp, pow, abs, sqrt, max, min, dot.
You can add more stdlib support by adding definitions using the [ForwardDerivativeOf] attribute. For example, the following code defines the derivative implementation for sin:

__generic<T : __BuiltinFloatingPointType>
[ForwardDerivativeOf(sin)]
DifferentialPair<T> __d_sin(DifferentialPair<T> dpx)
{
    return DifferentialPair<T>(
        sin(dpx.p),
        T.dmul(cos(dpx.p), dpx.d));
}

All locations that sees the definition of __d_sin will treat sin as a differentiable function.

Dos and Don'ts in Differentiable Functions

Not all statements are supported inside a differentiable function. Here's an incomplete list of what is and isn't supported.

Okay to use inside differentiable function:

• Floating-point arithmetic operators (+, -, *, /)
• Integer arithmetic/logical operators (%, |, &) are allowed but will not be differentiated since they are discrete.
• Nested function calls: Call another function that is differentiable (i.e. has a [ForwardDifferentiable] or [ForwardDerivative] attribute)
• Member access (only for members defined with [DerivativeMember] attribute either explictly by the user or implicitly by the compiler through automatic IDifferentiable conformance synthesis)
• Array Indexing
• Control flow like: for loops as well as if-else conditional blocks.
• Calls to non-differentiable functions. Use no_diff keyword to clarify the intention when the result of a non-differentiable function contributes to a differentiable result value.

Not okay to use inside differentiable functions (error or undefined behaviour):

• Global variable access (textures, samplers or any other globally defined vars)
• Calling unsupported intrinsics such as TraceRay or TraceRayInline
• Switch-case statements