Effective use of the Mathematica Compiler and Code Generation

Rob Knapp

Wolfram Research, Inc.

Even if you have never used the Mathematica function Compile directly, it has probably been used “under the hood” for a variety of things you have done.

It is used in plotting functions

It is used for numerical computation functions like NDSolve, FindRoot, etc.

It is used extensively in the part of Mathematica system code written in the Mathematica language.

In all of these cases, the compiler is used to speed up just a part of the computation rather than being used for the entirety.

The most common difficulty I see with many people’s use of the compiler is treating it like a C or Java compiler where an entire program has to be compiled.  In Mathematica the compiler is well integrated with the rest of the system via basic number types and PackedArrays so it is practical and very often more efficient to compile small numerically intensive parts of a computation and leave the rest to ordinary Mathematica evaluation.  An additional advantage of this approach is that you can concentrate on making small pieces of a computation fast, something that is much easier and more likely to produce results than trying to work with the whole computation at once.

Some things to keep in mind for effective use of the Mathematica Compiler

I’ll keep harping on these throughout the talk, so it seemed worthwhile to do a quick summary now.

You don’t have to compile your whole program (I’m repeating myself since I think this is so important)

The internal data structure for lists in Compile is exactly the same as the internal data structure for PackedArray

This means that passing PackedArray in and out of CompiledFunction is quite efficient

The template for calling functions in the virtual machine is the same as the template for LibraryLink functions

This means that setting up a LibraryLink function to do a critical part of your computation is essentially as fast as if it were functionality that the Compiler has built in

An easy way to make a LibraryLink function is to use Compile with CompilationTarget->”C”

Compiling in listability gives you several advantages

Flexibility: a single function can work for lists or scalars

Flexibility: in some cases this can be used to work around the restriction that nested lists be rectangular

Speed: parallelization can give some nice speedups

Compiling to C code can make things much faster, but the C compilation can be expensive.

I suggest On[“CompilerWarnings”]

On["CompilerWarnings"]

To see what has fully compiled use CompilePrint

Needs["CompiledFunctionTools`"]

Example: Chebyshev Approximations

The goal in this example is to define a function Tfun[f, {x0,x1}] that represents a function f of x on an interval [(Subscript[x, 0]), (Subscript[x, 1])] as a sum of Chebyshev polynomials on [-1, 1]

f(x) = (∑_(i=0)^n)(Subscript[a, i])(Subscript[T, i])(x)

A really simple implementation of this is to use

Clear[Tfun];
SetAttributes[Tfun, HoldAll];
Tfun[fx_, {x_, x0_,x1_}, a_List][t_] := Module[{s = (2t - (x1+x0))/(x1-x0)},tsum[a,s]]

tsum[a_,x_] := a.Table[ChebyshevT[i,x],{i,0,Length[a]-1}];

Format[Tfun[fx_, {x_,x0_,x1_}, a_List]] := Subscript[Tfun[fx, {x, x0,x1}],"<"<>ToString[Length[a]-1]<>">"]

You can use Fourier cosine transforms to generate sets of coefficients (I won’t go into the math behind this, but it is fun).
ToTfun implements this

Clear[ToTfun];
ToTfun[fx_, {x_, x0_, x1_}, n_] := Module[{cnodes=N[Cos[Pi Range[0,n]/n]],xnodes, fnodes,coeffs, dx},
xnodes = (x1 + x0 + cnodes(x1-x0))/2;
fnodes = Table[fx,{x,xnodes}];
coeffs = FourierDCT[fnodes,1]*Sqrt[2/n];
coeffs[[{1,-1}]]/=2;
Tfun[fx,{x, x0, x1}, Evaluate[coeffs]]
]

Another possible way to generate them would be to do a polynomial fit with a Chebyshev basis to make an approximate function for data.

As a test consider Runge’s function that causes problems for regularly spaced interpolation points.

runge[x_] := 1/(1+(5x)^2);

trunge = ToTfun[runge[x-3],{x,0,2 Pi},47]

Compare the function and its approximation

Plot[{trunge[x],runge[x-3]},{x,0,2 Pi}, PlotRange->All]

As currently implemented, Tfun is not very fast:

test = RandomReal[{0,2 Pi}, 10000];

AbsoluteTiming[Map[trunge,test];]

AbsoluteTiming[trunge[test];]

To be useful, the approximation needs to be as fast or faster than the original function.

AbsoluteTiming[Map[runge[#-3]&,test];]

It looks even worse against:

AbsoluteTiming[runge[test-3];]

The reason that it is so slow is that it has to evaluate (Subscript[T, 0]), (Subscript[T, 1]), ... (Subscript[T, n]) separately.
There is a faster and more numerically stable way, via Clenshaw’s algorithm that uses the recurrence relation of the Chebyshev functions.

Here is an implementation for Clenshaw’s algorithm

clenshawVM = Compile[{{coeffs, _Real, 1}, {x,_Real}}, Module[{n = Length[coeffs]}, 
Which[
n==1, coeffs[[1]],
n==2, coeffs[[1]] + x coeffs[[2]],
True, Module[{a1,a2,twox = 2. x, t=x},
a2 = coeffs[[n]];
a1 = twox*a2 + coeffs[[n - 1]];
Do[t = a1; a1 = twox*a1 - a2 + coeffs[[i]]; a2 = t,
{i,n-2,2,-1}];
x*a1-a2+coeffs[[1]]
]
]
], RuntimeAttributes->Listable]

The RuntimeAttribute Listable is given because you want Tfun to be Listable -- it also gives some advantages like automatic parallelization

Replace the definition of tsum with this and time ...

Clear[tsum]; tsum = clenshawVM;
AbsoluteTiming[Map[trunge,test];]

AbsoluteTiming[trunge[test];]

So its nearly a factor of 10 faster.  We can do better yet by compiling to C:

clenshawC = Compile[{{coeffs, _Real, 1}, {x,_Real}}, Module[{n = Length[coeffs]}, 
Which[
n==1, coeffs[[1]],
n==2, coeffs[[1]] + x coeffs[[2]],
True, Module[{a1,a2,twox = 2. x, t=x},
a2 = coeffs[[n]];
a1 = twox*a2 + coeffs[[n - 1]];
Do[t = a1; a1 = twox*a1 - a2 + coeffs[[i]]; a2 = t,
{i,n-2,2,-1}];
x*a1-a2+coeffs[[1]]
]
]
],  RuntimeAttributes->Listable, CompilationTarget->"C"];

Replace the definition of tsum with this and time ...

Clear[tsum]; tsum = clenshawC;
AbsoluteTiming[Map[trunge,test];]

Clear[tsum]; tsum = clenshawC;
AbsoluteTiming[trunge[test];]

Pretty good!  However we can go even faster yet.  Since the idea here is to make an approximation that is as fast as possible, strict accuracy and error checking is not necessary.  When this is the case, use RuntimeOptions->”Speed” to get a little bit extra:

clenshawC = Compile[{{coeffs, _Real, 1}, {x,_Real}}, Module[{n = Length[coeffs]}, 
Which[
n==1, coeffs[[1]],
n==2, coeffs[[1]] + x coeffs[[2]],
True, Module[{a1,a2,twox = 2. x, t=x},
a2 = coeffs[[n]];
a1 = twox*a2 + coeffs[[n - 1]];
Do[t = a1; a1 = twox*a1 - a2 + coeffs[[i]]; a2 = t,
{i,n-2,2,-1}];
x*a1-a2+coeffs[[1]]
]
]
],  RuntimeAttributes->Listable, CompilationTarget->"C", RuntimeOptions->"Speed"];

Clear[tsum]; tsum = clenshawC;
AbsoluteTiming[Map[trunge,test];]

Clear[tsum]; tsum = clenshawC;
AbsoluteTiming[trunge[test];]

Some of the time may be in the scaling, so it makes sense to define a version that does that too:

sclenshawC = Compile[{{coeffs, _Real, 1}, {s,_Real}, {x0, _Real},{x1, _Real}}, Module[{n = Length[coeffs], x = (2s - (x1+x0))/(x1-x0)}, 
Which[
n==1, coeffs[[1]],
n==2, coeffs[[1]] + x coeffs[[2]],
True, Module[{a1,a2,twox = 2. x, t=x},
a2 = coeffs[[n]];
a1 = twox*a2 + coeffs[[n - 1]];
Do[t = a1; a1 = twox*a1 - a2 + coeffs[[i]]; a2 = t,
{i,n-2,2,-1}];
x*a1-a2+coeffs[[1]]
]
]
] ,RuntimeAttributes->Listable, CompilationTarget->"C", RuntimeOptions->"Speed", Parallelization->False];

Clear[Tfun];
SetAttributes[Tfun, HoldAll];
Tfun[fx_, {x_, x0_,x1_}, a_List][t_] :=sclenshawC[a,t,x0,x1]

AbsoluteTiming[Map[trunge,test];]

AbsoluteTiming[trunge[test];]

If you want to make this into a more general function, we want it to work for different types and precision.  
The best way to do this is to set it up to compile as needed.

Using some tricks, we can do this using a minimal amount of code

First set up a general definition for the algorithm

clenshaw = Function[{coeffs, s, x0,x1}, 
Module[{n = Length[coeffs], x = (2s - (x1+x0))/(x1-x0)}, 
Which[
n==1, coeffs[[1]],
n==2, coeffs[[1]] + x coeffs[[2]],
True, Module[{a1,a2,twox = 2 x, t},
a2 = coeffs[[n]];
a1 = twox*a2 + coeffs[[n - 1]];
Do[t = a1; a1 = twox*a1 - a2 + coeffs[[i]]; a2 = t,
{i,n-2,2,-1}];
x*a1-a2+coeffs[[1]]
]
]
]];

This sets up a function that takes the function and compiles it

HoldPattern[clenshawCompile[ctype_, Function[args_, body_]]] := Block[args,Module[{cargs = Table[{arg, _Real},{arg,args}]},cargs[[1]] = {args[[1]],Blank[ctype], 1}; Compile[Evaluate[cargs], body, RuntimeAttributes->Listable, CompilationTarget->"C", RuntimeOptions->"Speed"]]];

Delayed definitions for real and complex versions -- these will only actually be compiled if used.

clenshawR := (clenshawR = clenshawCompile[Real, clenshaw]);
clenshawC := (clenshawC = clenshawCompile[Complex, clenshaw]);

Clear[Tfun];

SetAttributes[Tfun, HoldAll];

Tfun[fx_, {x_, x0_,x1_}, a_List][t_] := 
Which[
Developer`PackedArrayQ[a,Real],clenshawR[a,t,x0,x1],
Developer`PackedArrayQ[a,Complex],clenshawC[a,t,x0,x1],
True, Module[{prec = Precision[a]}, Block[{$MinPrecision = prec, $MaxPrecision = prec},clenshaw[a,t,x0,x1]]]
]

Format[Tfun[fx_, {x_,x0_,x1_}, a_List]] := Subscript[Tfun[fx, {x, x0,x1}],"<"<>ToString[Length[a]-1]<>">"]

Now redefine the constructor so that it handles complex valued functions and takes a WorkingPrecision option

Clear[ToTfun];

Options[ToTfun] = {WorkingPrecision->MachinePrecision};

ToTfun[fx_, {x_, x0_, x1_}, n_, OptionsPattern[]] := Module[{cnodes,xnodes, fnodes,coeffs, dx, wprec},
wprec = OptionValue[WorkingPrecision];
cnodes=N[Cos[Pi Range[0,n]/n], wprec];
xnodes = (x1 + x0 + cnodes(x1-x0))/2;
fnodes = Table[fx,{x,xnodes}];
Block[{$MinPrecision = wprec, $MaxPrecision = wprec},
coeffs = FourierDCT[fnodes,1]*Sqrt[2/n];
coeffs[[{1,-1}]]/=2
];
Tfun[fx,{x, x0, x1}, Evaluate[coeffs]]
]

Now try it out ...

thyper = ToTfun[HypergeometricPFQ[{1,1},{3,3,3},x],{x,-30,30},100]

AbsoluteTiming[thyper[2.]]

Plot[thyper[x] - HypergeometricPFQ[{1,1},{3,3,3},x],{x,-30,30}, PlotRange->All]

test = RandomReal[{-30,30},1000];

AbsoluteTiming[HypergeometricPFQ[{1,1},{3,3,3},test];]

AbsoluteTiming[thyper[{1,1},{3,3,3},test];]

texp = ToTfun[Exp[I x],{x,0,2 Pi},20]

AbsoluteTiming[texp[1]]

ParametricPlot[{Re[texp[t]],Im[texp[t]]},{t,0, 2 Pi}]

It also seamless works for higher precision ...

runge[x_] := 1/(1+(5x)^2);

trunge = ToTfun[runge[x-3],{x,0,10},1234, WorkingPrecision->32]

... albeit a bit slower

AbsoluteTiming[Max[Table[RealExponent[runge[x- 3] - trunge[x]],{x,RandomReal[{0,10},1000, WorkingPrecision->32]}]]]

Example: Making a simple fixed step ODE solver for autonomous equations

As a bigger example of what you can do with the compiler and how to use its modularity, lets construct a simple ODE solver based on the Heun method (a second order RungeKutta method that I’m not particularly worried that it will compete with NDSolve.)

Lets suppose that we have a function f[v] (v a vector) and the ODEs are given by v’[t]=f[v[t]]

A step is done by

heunstep = Compile[{{h, _Real},{v,_Real,1}},
Module[{k0,k1},
k0 = h f[v];
k1=h f[v+k0];
v + (k0+k1)/2]]

Now lets try it on the Lorenz equations

f = Function[{v},Module[{x,y,z},{x,y,z}=v;{-3 (x-y),-x z+26.5 x-y,x y-z}]]

And try a step ...

x0 = {0.,1.,0.};
h = 0.01;

heunstep[h,x0]

Seems to work, so maybe that warning was just so much noise ...

This iterates a solution for n steps

n = 20000;
AbsoluteTiming[solution = Module[{x = x0},Table[x = heunstep[h,x],{n}]];]

Its not terribly fast, but it looks like you’d expect

Graphics3D[Line[solution]]

It is more than a little slower than NDSolve (with less numerical accuracy)

AbsoluteTiming[s=First[NDSolve[{x'[t]==-3 (x[t]-y[t]),y'[t]==-x[t] z[t]+26.5 x[t]-y[t],z'[t]==x[t] y[t]-z[t],x[0]==z[0]==0,y[0]==1},{x,y,z},{t,0,200},MaxSteps->∞]]]

Lets see if we can get this to work a bit faster

One of the first things to do when a CompiledFunction is running slower than expected is to see what it does with CompilePrint

CompilePrint[heunstep]

The speed killers here are the MainEvaluates to evaluate the function

MainEvaluate

MainEvaluate indicates that the compiler resorts to Mathematica to evaluate an expression that it could not compile.  
There is nothing really wrong with this, but a Mathematica evaluation is much slower than interpreting an instruction in the virtual machine.

Sometimes it is unavoidable -- in those cases it is occasionally possible to restructure code so the evaluation is done once instead of many times in a loop.

Many times, however, it is avoidable and we’ll investigate that.

If the function called is a LibraryFunction or CompiledFunction, the virtual machine can avoid the evaluation adaptively

cf = Compile[{{v, _Real, 1}},Module[{x,y,z},{x,y,z}=v;{-3 (x-y),-x z+26.5 x-y,x y-z}]]

AbsoluteTiming[Block[{f = cf},solution = Module[{x = x0},Table[x = heunstep[h,x],{n}]];]]

So we get a huge speedup here

Note that it has to be explicitly assigned to a CompiledFunction -- this does the evaluation

AbsoluteTiming[Block[{f},f[v_] := cf[v];solution = Module[{x = x0},Table[x = heunstep[h,x],{n}]];]]

Generated code can not be as adaptive, so if we want to use code generation, we need to compile the step function after the rhs function is compiled and using the definition.  
You can have Mathematica inline the definition using

cheunstep = Compile[{{h, _Real},{v,_Real,1}},
Module[{k0,k1,k2,k3},
k0 = h cf[v];
k1=h cf[v+k0];
v + (k0+k1)/2], CompilationOptions->{"InlineExternalDefinitions"->True}]

CompilePrint[cheunstep]

This is comparable in timing to adaptively using the CompiledFunction, but you have lost the flexibility

AbsoluteTiming[solution = Module[{x = x0},Table[x = cheunstep[h,x],{n}]];]

However, what you have gained is that you can effectively compile to C

AbsoluteTiming[cheunstep = Compile[{{h, _Real},{v,_Real,1}},
Module[{k0,k1,k2,k3},
k0 = h cf[v];
k1=h cf[v+k0];
v + (k0+k1)/2], CompilationOptions->{"InlineExternalDefinitions"->True}, CompilationTarget->"C", RuntimeOptions->"Speed"];]

AbsoluteTiming[solution = Module[{x = x0},Table[x = cheunstep[h,x],{n}]];]

If the function being called is not terribly big, it may be advantageous to inline the code and not just the CompiledFunction

cheunstep = Compile[{{h, _Real},{v,_Real,1}},
Module[{k0,k1,k2,k3},
k0 = h cf[v];
k1=h cf[v+k0];
v + (k0+k1)/2], CompilationOptions->{"InlineExternalDefinitions"->True, "InlineCompiledFunctions"->True}, CompilationTarget->"C", RuntimeOptions->"Speed"];

AbsoluteTiming[solution = Module[{x = x0},Table[x = cheunstep[h,x],{n}]];]

Is effectively equivalent to

heunstep = Compile[{{h, _Real},{v,_Real,1}},
Module[{k0,k1,k2,k3},
k0 = h f[v];
k1=h f[v+k0];
v + (k0+k1)/2], CompilationOptions->{"InlineExternalDefinitions"->True}, CompilationTarget->"C", RuntimeOptions->"Speed"];

AbsoluteTiming[solution = Module[{x = x0},Table[x = heunstep[h,x],{n}]];]

Now the stepcode is fast enough that it is worth compiling the iteration too.  
I’ve done this for  a generic step function for now

AbsoluteTiming[citerate = Compile[{{h, _Real}, {n, _Integer},{x0,_Real,1}},Module[{x = x0}, Table[x = step[h,x],{n}]], {{_step, _Real, 1}}];]

Timing[solution = Block[{step = cheunstep}, citerate[h, n,x0]];]

Now we could get a tiny bit of extra speed by inlining everything, but at the cost of extra compilation time.

AbsoluteTiming[citerateinline = Compile[{{h, _Real}, {n, _Integer},{x0,_Real,1}},Module[{x = x0}, Prepend[Table[x = cheunstep[h,x],{t,0,200,h}],x0]], CompilationOptions->{"InlineExternalDefinitions"->True, "InlineCompiledFunctions"->True}, CompilationTarget->"C"];]

Timing[solution = citerateinline[h, n,x0];]

So you gain .002 second at the extra cost of .18 second!

So the thing to do now is to put together a function that solves an autonomous ODE specified by a function

asolve[f_Function, v0_, h_, n_, opts___] := 
Module[{citerate, cstep},
cstep = Compile[{{hh, _Real},{v,_Real,1}},
Module[{k0,k1,k2,k3},
k0 = hh f[v];
k1=hh f[v+k0];
v + (k0+k1)/2], opts, CompilationOptions->"InlineExternalDefinitions"->True];
citerate = Compile[{{hh, _Real}, {nn, _Integer},{x0,_Real,1}},Module[{x = x0}, Prepend[Table[x = cstep[hh,x],{nn}],x0]], opts, CompilationOptions->{"InlineExternalDefinitions"->True}];
citerate[h,n,v0]
]

AbsoluteTiming[solution = asolve[f,x0,h,n, RuntimeOptions->"Speed"];]

Has it got it correct?

Graphics3D[Line[solution]]

Compare to the C compilation

AbsoluteTiming[solution = asolve[f,x0,h,20000, CompilationTarget->"C", RuntimeOptions->"Speed"];]

Doesn’t make sense to use C!!

Rossler Equation

r = With[{a = .1, b = .1, c = 14},Function[{v}, Module[{x,y,z},{x,y,z} = v; {-y-z,x + a y, b + z(x - c)}]]];
r0 = {0.,0.,0.};

AbsoluteTiming[solution = asolve[r,x0,h,n, RuntimeOptions->"Speed"];]

Graphics3D[Line[solution]]

Example: Symplectic Euler Method

This uses n steps of size h of the symplectic Euler method starting with init to approximate a solution of the pendulum equation.

psolve = Function[{h, n, init}, Module[{x, y}, {x,y} = init;Table[x = Mod[x + h y, 2 π];y = y + h Sin[x];{x,y},{n}]]]

Run the function

AbsoluteTiming[s1 = psolve[.00025,100000,{.001,0}];]

ListPlot[s1, PlotRange->All]

Just with Mathematica this is not too terribly slow, though it is not really fast either -- you can do better

All you have to do to compile it is to replace Function with Compile and give types for the arguments ...

vmpsolve = Compile[{{h,_Real},{ n,_Integer}, {init, _Real, 1}}, Module[{x, y}, {x,y} = init;Table[x = Mod[x + h y, 2 π];y = y + h Sin[x];{x,y},{n}]]]

AbsoluteTiming[s2 = vmpsolve[.00025,100000,{0.001,0}];]

Not much work for a big jump in speed.

Timing[cpsolve = Compile[{{h,_Real},{ n,_Integer}, {init, _Real, 1}}, Module[{x, y}, {x,y} = init;Table[x = Mod[x + h y, 2π];y = y + h Sin[x];{x,y},{n}]], CompilationTarget->"C"]]

AbsoluteTiming[s3 = cpsolve[.00025,100000,{0.001,0}];]

Now its fast enough to be somewhat interactive

Manipulate[Show[Graphics[{PointSize[0],Point[cpsolve[10.^lh, 100000, init]]}, PlotRange->{{0,2π},{-3,3}}, Axes->True]],{{init,{0.001,0}},Locator},{{lh,0},-5,0}]

It’ll be nice to make an image with a bunch of initial conditions

init = Table[{x,0.},{x, 0.01,1,0.01}];
colors = Map[ColorData["Rainbow"],Range[Length[init]]/Length[init]];
np =10000Ceiling[Table[1/(Pi-x),{x,0.01,1,0.01}]];

AbsoluteTiming[data = MapThread[cpsolve[1,#1,#2]&, {np, init}];]

It is possible to make the computation faster at the expense of some overhead:

Timing[lsolve = Compile[{{h,_Real},{ n,_Integer}, {init, _Real, 1}}, Module[{x, y}, {x,y} = init;Table[x = Mod[x + h y, 2 Pi];y = y + h Sin[x];{x,y},{n}]], CompilationTarget->"C", RuntimeAttributes->Listable]]

AbsoluteTiming[data=lsolve[1, np, init];]

Ans now a nice image:

Show[Graphics[{
PointSize[0],MapThread[{#1,Point[#2]}&, {colors, data}]}], PlotRange->{{0,2π},{-3,3}}, Axes->True]

Summary

Effective use of the Mathematica compiler and code generation can speed up calculations tremendously.

Some things to consider are:

Be aware of types

Use CompiledFunctionTools and PrintCompile to see what a CompiledFunction is doing

Speed things up.  Here are things to try:

Try to avoid Mathematica evaluations

Check that types are as consistent as possible

If you want to run many instances think about using RuntimeAttributes->Listable possibly with Parallelization->True

If you expect to use the function many times, try CompilationTarget->”C”

Here is what I think is the most important thing of all:  For most computations, you’ll have far better success in using the compiler to speed up parts of the computation rather than trying to use Compile to generate one function that does it all.

For lots more information, go to the The Mathematica Compiler User Guide