Troubleshooting Compile making things slower

Question

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

Answer 1

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 = Compile`GetElement[w, 1];
    result[[1, 2]] = w2 = Compile`GetElement[w, 2];
    Do[
     X11 = Compile`GetElement[X, j, 1, 1];
     X12 = Compile`GetElement[X, j, 1, 2];
     X21 = Compile`GetElement[X, j, 2, 1];
     X22 = Compile`GetElement[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[Compile`GetElement[X, j, k, l], {k, 1, d}, {l, 1, d}];
    ww = Table[Compile`GetElement[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.