Troubleshooting Compile making things slower

I'm trying to speed up the following piece of code which runs in about 0.5 seconds on my machine.

numSteps = 5;
numSamples = 10000;
b = 2;
d = 2;

step = Function[{Typed[w, "Real32"]},
Module[{X}, X = RandomVariate[NormalDistribution[], {b, d}];
w - .1 X\[Transpose] . X . w]
];
batchStep = Function[{Typed[wb, "Real32"]}, step /@ wb];

wb0 = RandomVariate[NormalDistribution[], {numSamples, d}];
NestList[batchStep, wb0, numSteps]; // Timing


Using FunctionCompile fails with Cannot find a definition for RandomVariate. Meanwhile, using Compile below makes it run in about 1 second. What are the tips to troubleshoot this?

numSteps = 20;
numSamples = 10000;
b = 2;
d = 2;

step = Compile[{{w, _Real, 1}},
Module[{X}, X = RandomVariate[NormalDistribution[], {b, d}];
w - .1 X\[Transpose] . X . w]
];
batchStep = Compile[{{wb, _Real, 2}}, step /@ wb];

wb0 = RandomVariate[NormalDistribution[], {numSamples, d}];
NestList[batchStep, wb0, numSteps]; // Timing


Summarizing lessons on slow Compile

• Huge penalty for modifying global variable from inside compiled function, use Module to localize variables

• Get better performance for pre-generating random data and threading over corresponding lists, instead of generating on-demand from inside the compiled function

• matrix multiplication for small matrices inefficient, can hand-code solution like in Henrik Schumacher solution below

There is no point in compiling RandomVariate as its backend is already compiled. Putting it into Compile just leads to some call from the CompiledFunction to the Mathematica kernel -- which just adds overhead without providing any advantages. What you could do instead is compile the loop generated NestList:

cApplyNested = Compile[{{X, _Real, 3}, {w, _Real, 1}},

Block[{m, n, w1, w2, result, X11, X12, X21, X22},
m = Dimensions[X][[1]];
result = Table[0., {m + 1}, {2}];
result[[1, 1]] = w1 = CompileGetElement[w, 1];
result[[1, 2]] = w2 = CompileGetElement[w, 2];
Do[
X11 = CompileGetElement[X, j, 1, 1];
X12 = CompileGetElement[X, j, 1, 2];
X21 = CompileGetElement[X, j, 2, 1];
X22 = CompileGetElement[X, j, 2, 2];
result[[j + 1, 1]] = w1 = w1 - 0.1 (w1 (X11 X11 + X21 X21) + w2 (X11 X12 + X21 X22));
result[[j + 1, 2]] = w2 = w2 - 0.1 (w1 (X11 X12 + X21 X22) + w2 (X12 X12 + X22 X22));
, {j, 1, m}];
result
],

CompilationTarget -> "C",
RuntimeAttributes -> {Listable},
Parallelization -> True,
RuntimeOptions -> "Speed"
];


Now you can generate all samples first and then thread cf over this data in parallel like this:

Xlist = RandomVariate[NormalDistribution[], {numSamples, numSteps, b, d}];
wb0 = RandomVariate[NormalDistribution[], {numSamples, d}];
result = cApplyNested[Xlist, wb0];


On my machine this takes about 0.016409 seconds for numSamples = 10000 and numSteps = 20, while OP's original code takes 1.13243 seconds.

Edit

If you want to be a bit more flexible, you can employ Mathematica to symbolically generate the code. Then you can employ Compile as a JIT-compiler to create a function for each d you need (and only if and when you need it). This could be done as follows:

ClearAll[cApplyNested];
cApplyNested[d_] := cApplyNested[d] = Block[{XX, X, j, w, ww, code},
XX = Table[CompileGetElement[X, j, k, l], {k, 1, d}, {l, 1, d}];
ww = Table[CompileGetElement[w, k], {k, 1, d}];
With[{code = ww - 0.1 XX\[Transpose] . XX . ww, dim = d},
Compile[{{X, _Real, 3}, {w0, _Real, 1}},

Block[{m, w, result},
m = Dimensions[X][[1]];
result = Table[0., {m + 1}, {dim}];
w = Table[0., {dim}];

result[[1]] = w = w0;

Do[result[[j + 1]] = w = code;, {j, 1, m}];

result
],

CompilationTarget -> "C",
RuntimeAttributes -> {Listable},
Parallelization -> True,
RuntimeOptions -> "Speed"
]
]
];


Now you can do

numSteps = 20;
numSamples = 10000;
d = 2;

Xlist = RandomVariate[NormalDistribution[], {numSamples, numSteps, d, d}];
wb0 = RandomVariate[NormalDistribution[], {numSamples, d}];
result = cApplyNested[d][Xlist, wb0];


without any severe performance degression. The only thing that you will observe is that call cApplyNested[d] for the first time will require some extra time for compilation process. But since we use memoization here, the time of the next call cApplyNested[d] the CompiledFunction will be already known.

• ok, that's a big speed-up .... I'm surprised that you are hardcoding matrix multiplication for d=2, wondering if that's really necessary...this hardcoding reduces utility of this function Feb 7 at 7:27
• The Dot operations for matrices and vectors use BLAS routines and those routines are notoriously slow for small matrix sizes (because they were designed for large matrix sizes). So it does make quite a lot of sense to use hand-coded routines, in particular because one can completely unroll the loops if the sizes are known at compile time. Feb 7 at 7:42
• I suspect a bug in that code, cf is missing. If I use cApplyNested instead of cf I get the wrong shape out Feb 7 at 17:38
• Oh I see. m = Dimensions[X][[2]]; should have been m = Dimensions[X][[1]];. Now it should work as expected. Feb 7 at 18:18
• thanks, works as expected now, summarizing high-level lessons in question for posterity Feb 7 at 18:23