Speeding Up the Wolfram Language

An Introduction to the Wolfram Function Compiler

Arnoud Buzing
Towards Data Science

--

At the core of the Wolfram Language sits an expression evaluator that transforms input expressions into computed results using standard evaluation rules. This evaluator is what makes the Wolfram Language (abbreviated here as WL) incredibly flexible and powerful, but there is a performance cost to running individual expressions through this evaluator. This story shows how you can use the new Wolfram function compiler to get blazing fast performance. The techniques described here formed the computational basis for the creation of the following video:

The Basics

Let’s take a look at a very basic example. Suppose we want to numerically square lots of complex numbers, not an unusual task in mathematics, physics, and other sciences. The naive way to do this is to define a function that takes one complex number and squares it:

f[z_] := z^2

Then we can call that function and run it repeatedly in a Do loop:

Do[ f[0.0+0.0*I], 1000000 ]

Writing WL code like this is very inefficient though. On my Linux box (Intel Xeon CPU E3–1245) this evaluation takes about 700 milliseconds. To quote Data from Star Trek: “For an android, that’s nearly an eternity”.

A typical improvement to this code is to use the fact that many WL functions are listable, which means that they can operate on lists sequentially the same way they can operate on individual expressions.

In this case, we can use this listability to get a small performance improvement:

Attributes[f] = {Listable}c = ConstantArray[0.0 + 0.0*I, 1000000]f[c]

This performs the same operation in about 500 milliseconds, so 200 ms less than before. In simple cases, you can also take advantage of functions with built-in performance optimizations. In this case, squaring a number (a special case of the Power function) has a built-in optimization so you can simply write:

c^2

This takes the evaluation time down a lot, to 2 milliseconds.

But these special cases are not available for when your code becomes more complex. To get the next level of performance improvements you need to compile your code to optimized machine instructions, and this is what the WL function compiler does.

To compile WL into machine-code, you have to add type information but otherwise, you can use standard WL code. Here is an example of how you can write the original example with the WL function compiler:

f = Function[{Typed[z, "ComplexReal64"]}, Do[z^2, 1000000]]cf = FunctionCompile[f]cf[0.0+0.0*I]

This code evaluates in 0.3 milliseconds, so it is faster than any of the uncompiled methods. There is a cost of a one-time compilation step which can take anywhere from a few hundred milliseconds to several seconds or minutes in the case of very complicated code. But the end result is always very fast and optimized code.

There is one more optimization to be made. The code above still includes extra machine code to allow it to be aborted safely in the middle of a computation. You can choose to remove this safeguard and get even higher performance:

cf = FunctionCompile[ f, 
CompilerOptions -> {"AbortHandling" -> False}
]
cf[0.0+0.0*I]

This code evaluates in about 0.009 milliseconds (or 9 microseconds).

Example

Now let’s put this new knowledge into practice. For example, we can create a visualization for a complex-valued function by computing many points in the complex plane. To create an image of the result, we take the absolute value after squaring each complex number. Note that the Table limits (x and y) are given such that they create an image with the proper orientation:

f = Function[{},
Table[ Abs[(x+I*y)^2],
{y, 1.0, -1.0, -0.01},
{x, -1.0, 1.0, 0.01}
]
cf = FunctionCompile[ f,
CompilerOptions -> {"AbortHandling" -> False}
]
Image[cf[]]

The result is a simple image with gray levels indicating the absolute value of z², where 0.0 represents black and 1.0 represents white. Absolute values that are greater than 1.0 are clipped and displayed as white here.

(All images by author)

We can make this look a little bit more exciting by changing the color space from gray levels to the HSB color space. Here, HSB stands for Hue, Saturation, and Brightness. We will fix saturation and brightness for the moment and only vary the hue. Each function evaluation now generates an HSB triplet:

f = Function[{}, Table[
{Abs[(x + I*y)^2], 1.0, 1.0},
{y, 1.0, -1.0, -0.01},
{x, -1.0, 1.0, 0.01}
]]
cf = FunctionCompile[f,
CompilerOptions -> {"AbortHandling" -> False}]
Image[cf[], ColorSpace -> "HSB"]

In the WL the HSB color space is cyclic, which means that the colors in the 0–1 range repeat in the 1–2 range and beyond. This means that the red color represents 0, or 1, or 2, and so on. In the image below the red color in the center represents 0 and the red circle around it represents 1, and the red in the very corners represent 2:

Next, instead of computing the absolute values of z², we can look at its argument (or Arg in the WL).

f = Function[{}, Table[
{0.5*(1+Arg[(x + I*y)^2]/Pi), 1.0, 1.0},
{y, 1.0, -1.0, -0.01},
{x, -1.0, 1.0, 0.01}
]]
cf = FunctionCompile[f,
CompilerOptions -> {"AbortHandling" -> False}]
Image[cf[], ColorSpace -> "HSB"]

We can also use both the absolute value and the argument of z², for example: Use the absolute value for hue and use arg for the brightness.

f = Function[{}, Table[
{Abs[(x + I*y)^2], 1.0, 1.0 + Arg[(x + I*y)^2]/Pi},
{y, 1.0, -1.0, -0.01},
{x, -1.0, 1.0, 0.01}
]]

Because the range of Arg is -Pi to Pi we need to apply some rescaling to avoid negative brightness values (which are completely black). Now we have a visualization of the complex z² function which both shows the absolute value (in hues) and the argument (in brightness, where small values near zero are black):

With this preliminary exploration completed, we can now move to more interesting visualizations, for example, f(z)=z³+1. The code below is slightly improved for readability. First, the complex number z is created from its real and imaginary parts. Next, its function value is computed. And finally, its color value is assigned:

f = Function[{}, Table[
Module[{z, fz},
z = x + I*y;
fz = z^3 + 1;
{Abs[fz], 1.0, 1.0 + Arg[fz]/Pi}]
, {y, 2.0, -2.0, -0.005}, {x, -2.0, 2.0, 0.005}]]

To better show the zeros of the function, the domain is also increased from -2 to 2 for the real and imaginary parts:

At this point, it becomes easy to experiment with any function. Below is an example with higher powers of z and trigonometric functions:

f = Function[{}, Table[
Native`UncheckedBlock@Module[{z = x + I*y, fz},
fz = Sin[z^5] + Cos[z^3];
{Log[1 + Abs[fz]], 1.0, 1.0 + Arg[fz]/Pi}]
, {y, 2.0, -2.0, -0.005}, {x, -2.0, 2.0, 0.005}]]

And the resulting image:

Despite its complexity, this final image takes only 180 milliseconds to compute. For more information on function compilation, visit the following web page:

--

--

I create awesome software at Wolfram Research, makers of Mathematica, Wolfram|Alpha, Wolfram Cloud, and many other products and services.