Efficient compilation of a list of functions

Question

Let's say that we have at our disposal

nbfunc = 4;

functions, whose exact expressions are gathered in the list

tabfunc = Table[LegendreP[n,x], {n, 1, nbfunc}];

For the purpose of the example, I gave myself the list of Legendre polynomials, whose evaluation can be compiled. For a given run, the number of functions nbfunc and their expressions tabfunc are fixed forever. But, they may change from one run to another.

Now, I would like to evaluate efficiently the function f[neval_,xeval_] defined as

f[neval_,xeval_] := tabfunc[[neval]] /. {x -> xeval};

where the integer neval is always assumed to satisfy 1 <= neval <= nbfunc. I can compile this evaluation by defining the function fC as

fC = Compile[{{n, _Integer}, {x, _Real}},
tabfunc[[n]],
CompilationTarget -> "C",
CompilationOptions -> {"ExpressionOptimization" -> True, "InlineCompiledFunctions" -> True, "InlineExternalDefinitions" -> True},
RuntimeOptions -> {"CatchMachineOverflow" -> False , "CatchMachineUnderflow" -> False, "CatchMachineIntegerOverflow" -> False, "CompareWithTolerance" -> False, "EvaluateSymbolically" -> False}];

The compilation works correctly (in particular no calls to MainEvaluate), and the timings are improved

f[4, 0.2] == fC[4, 0.2]
Table[f[4, 0.2], {i, 1, 1000}]; // AbsoluteTiming // First
Table[fC[4, 0.2], {i, 1, 1000}]; // AbsoluteTiming // First
(*True
0.005839
0.000169*)

Yet, when inspecting the compiled code of fC,

Needs["CompiledFunctionTools`"];
CompilePrint[fC]

we note that the code contains lines of the form

22  T(R1)0 = {R0, R6, R8, R5}
23  R6 = Part[ T(R1)0, I0]
24  Return

This implies that to compute fC[4,0.2], the compiled function first computes the list {f[1,0.2],f[2,0.2],f[3,0.2],f[4,0.2]}, then returns the fourth element of this list, i.e. the value f[4,0.2]. Unfortunately, this is unsatisfactory, as to compute fC[4,0.2], there should be no need to compute all the others f[i,0.2] for 1<=i<=3.

My question is therefore as follows:

How should one proceed to compile the function fC[n_,x_], so that only the needed expression for the n that is asked is effectively evaluated?

In particular, I face the difficulty that the number nbfunc may change from one run to another (and can be quite large), so that this cannot be done by hand.

Answer 1

Since Which is compilable, you can try this:

fC = Block[{n},
   With[{

    code = Which @@Join[Riffle[Thread[n == Range[nbfunc]], N[tabfunc]], {True,0.}]

    },
    Compile[{{n, _Integer}, {x, _Real}},
     code,
     CompilationTarget -> "C",
     Parallelization -> True,
     RuntimeAttributes -> {Listable},
     RuntimeOptions -> "Speed"
     ]
    ]
   ];

This way, the function is both listable and parallelized (although parallelization does not help very much):

k = 1000000;
fC[RandomInteger[{1, 4}, k], RandomReal[{-1, 1}, k]]; // 
  AbsoluteTiming // First

Don't forget to give a default case (with True); it might not work otherwise. Notice also how I inlined the actual code with With in order to "precompute" what can be computed in advance. The respective "CompilationOptions" for inlining are infamous for their lack of robustness.