Automatic differentiation is an exciting emerging technology which enables deep learning applications and is of particular value to PassiveLogic’s smart building platform.

The Swift language has first class support for autodiff. Using this, we can write a simple function and make it differentiable by decorating it with the differentiable annotation:

// This is our differentiable function, y = a² + b³. So complicated, // thank goodness we have autodiff to determine the derivatives! @differentiable(reverse) public func awesome(a: Float, b: Float) -> Float { a * a + b * b * b }

Now we can obtain the value and derivatives of awesome() simply by calling the magic valueWithGradient() function provided by the Swift compiler:

let (value, derivatives) = valueWithGradient(at: 3, 5, of: awesome) print(“\(value), dA = \(derivatives.0), dB = \(derivatives.1)”)

This is wonderfully simple and effective (shout-out to the Google Swift for Tensorflow team who drove the effort to get this integrated into the mainline Swift compiler).

But… what if it could be faster?

Enzyme is a project that integrates tightly with the LLVM optimizer to produce autodiff functions. Because it operates at the low level of LLVM IR, it can potentially generate faster autodiff code than implementations that work at an earlier point in the compilation pipeline, such as the current incarnation of differentiable Swift.

Since Swift uses the LLVM backend, it ought to in principle be possible to use Enzyme to produce the derivatives of awesome(). Let’s make it happen!

We start from the same Swift function (minus the differentiable annotation):

// Our enzyme differentiable function. Behold the awesomeness! public func awesome(a: Float, b: Float) -> Float { a * a + b * b * b }

Enzyme imposes a couple of requirements on the differentiable function. It must have a C calling convention; and any parameters for which we will get a derivative must be passed by reference (in order to use the Enzyme duplicated argument convention). So we write a small wrapper function:

@_cdecl(“awesomeWrapper”) func awesomeWrapper(a: UnsafeMutablePointer?, b: UnsafeMutablePointer?) -> Float { awesome(a: a!.pointee, b: b!.pointee) }

Now that we have a wrapper that Enzyme is willing to differentiate, we can ask Enzyme to generate a function to take the wrapper and a set of input parameters, and return the derivatives (we will discuss the mechanics of the Enzyme generation step in a subsequent post). It is required that the generated function name must start with the prefix __enzyme_autodiff (the Enzyme code generation pass looks for calls to any function starting with this prefix and takes this as an indication that it needs to generate code). There might of course be multiple generated functions in a project with different signatures, so we can use the part of the function name following the required prefix to disambiguate these functions.

In this case, awesome() is a function which takes two float parameters and returns a float, so using an invented naming convention we will call the generated function __enzyme_autodiff_Float_FloatFloat(). Because we are using the Enzyme duplicated argument convention, each of the input parameters will be followed by a “shadow” output parameter which will return the derivative of the corresponding input parameter.

The Enzyme generated function will use the C calling convention. So this leads us to a C prototype for __enzyme_autodiff_Float_FloatFloat():

float __enzyme_autodiff_Float_FloatFloat( float(*)(float*, float*), // function to differentiate: it // accepts two floats as input and // returns a float. const float*, // The first input parameter. float*, // (output): the derivative of the // first parameter. const float*, // The second input parameter. float*); // (output): the derivative of the // second parameter.

Since Swift has the ability to call C directly, we will place this prototype in a .h file and tell Swift it is an external C library (even though the actual implementation of the function exists in no library but will be generated by the Enzyme step during compilation).

Now we can finally write a Swift function to call __enzyme_autodiff_Float_FloatFloat() and return the derivatives of awesome()’s inputs:

You might think that the return value of the Enzyme generated function was the result of awesome(), but not so. So to implement the equivalent of Swift’s valueWithGradient() we will have to call awesome() directly to get the value and awesomeDerivative() to get the gradient:

Whew! That was a lot of work to reach feature parity with a single line of Swift code! Why on earth would we go through this misery? Well, here’s a simple benchmark that computes the derivatives of awesome() 100,000,000 times in a loop:

Oh. 477 times faster?! That’s worth a fair amount of misery.

This is of course a trivial example of a differentiable function, but we have observed similar speedups in significantly complex functions with hundreds of parameters. Enzyme is pretty impressive!

Despite the performance benefits observed from Enzyme-generated derivatives, you may want to hold off on replacing all of your differentiable Swift code for two reasons: first, that work is underway to significantly improve native differentiable Swift performance in ways that may narrow or remove this gap (some of which will be described later by fellow PassiveLogic employee Brad Larson); and second, that there are currently significant limitations on the Swift / Enzyme integration that we will discuss in a subsequent post.

The code presented in this post can be found here.