Removing calls to MainEvalute when using inlined compiled closures

Question

This question is tightly related to the answer Shaving the last 50 ms off NMinimize.

There @OleksandR shows how inlined closures can be used to eliminate calls to MainEvaluate. This is crucial for my application, since I need every last drop of performance mma can offer. Here is an extremely simplified working example (before you run this code, make sure to evaluate this package, containing the NelderMeadMinimize code):

Needs["CompiledFunctionTools`"]
With[{minimizer = 
    NelderMeadMinimize`Dump`CompiledNelderMead[
     Function[{a, b, c}, (a - d1)^2 + (b - d2)^2 + (c - d3)^2], {a, b, c},
        "ReturnValues" -> "OptimizedParameters"], 
   epsilon = $MachineEpsilon},
  orgFitter = 
   Compile[{{d1, _Real, 0}, {d2, _Real, 0}, {d3, _Real, 0}}, 
    minimizer[RandomReal[{0, 1}, {3 + 1, 3}], epsilon, -1], 
    CompilationOptions -> {"InlineCompiledFunctions" -> True}, 
    RuntimeOptions -> {"Speed", "EvaluateSymbolically" -> False}]];
StringMatchQ[CompilePrint[orgFitter], "*MainEvaluate*"]
(* -> False *)
orgFitter[12, 2, 3]

Now, if I change (a - d1)^2 + (b - d2)^2 + (c - d3)^2 to calling a compiled function the calls to MainEvalute are not eliminated! Please, instead of the simple myHi2 imagine a monster of a compiled procedure containing several hundred lines.

myHi2 = Compile[{{a, _Real, 0}, {b, _Real, 0}, {c, _Real, 
     0}, {d1, _Real, 0}, {d2, _Real, 0}, {d3, _Real, 
     0}}, (a - d1)^2 + (b - d2)^2 + (c - d3)^2];
With[{minimizer = 
    NelderMeadMinimize`Dump`CompiledNelderMead[
     Function[{a, b, c}, myHi2[a, b, c, d1, d2, d3]], {a, b, c}, 
     "ReturnValues" -> "OptimizedParameters"], 
   epsilon = $MachineEpsilon},
  orgFitter = 
   Compile[{{d1, _Real, 0}, {d2, _Real, 0}, {d3, _Real, 0}}, 
    minimizer[RandomReal[{0, 1}, {3 + 1, 3}], epsilon, -1], 
    CompilationOptions -> {"InlineCompiledFunctions" -> True, 
      "InlineExternalDefinitions" -> True}, 
    RuntimeOptions -> {"Speed", "EvaluateSymbolically" -> False}]];
StringMatchQ[CompilePrint[orgFitter], "*MainEvaluate*"]
(* -> True *)
orgFitter[12., 2., 3.]

How can I eliminate this calls to MainEvaluate?

Also, could someone help me understand, why this does not compile:

orgFitter = Compile[{{d1, _Real, 0}, {d2, _Real, 0}, {d3, _Real, 0}},
   With[{minimizer = 
      NelderMeadMinimize`Dump`CompiledNelderMead[
       Function[{a, b, c}, (a - d1)^2 + (b - d2)^2 + (c - d3)^2], {a, 
        b, c}, "ReturnValues" -> "OptimizedParameters"], 
     epsilon = $MachineEpsilon},
    minimizer[RandomReal[{0, 1}, {3 + 1, 3}], epsilon, -1]], 
   CompilationOptions -> {"InlineCompiledFunctions" -> True, 
     "InlineExternalDefinitions" -> True}, 
   RuntimeOptions -> {"Speed", "EvaluateSymbolically" -> False}];

CompilePrint[orgFitter]

Answer 1

This is actually not specific to the NelderMeadMinimize` package but rather a limitation of the behaviour of Compile with respect to CompilationOptions -> {"InlineCompiledFunctions" -> True}. Essentially, compiled functions will only be inlined (in most cases) if they appear lexically in the input to Compile. You can ensure that in three main ways:

Nested With

This is my preferred approach of the three shown here (although Leonid's LetL macro would be even more preferable). Unfortunately it's a little difficult to demonstrate it without copying your code out in full, but hopefully this abbreviated version is clear enough. Instead of writing

myHi2 = Compile[...];
With[{minimizer = ..., ...}, ...]

you can use

With[{myHi2 = Compile[...]},
 With[{minimizer = ..., ...}, ...]
]

which ensures that myHi2 is never encountered by Compile as an opaque function call, but rather as a compiled function. It is then inlined correctly.

"RuntimeOptions" -> {"EvaluateSymbolically" -> False} and Evaluate

Simply define

myHi2 = Compile[
 {{a, _Real, 0}, {b, _Real, 0}, {c, _Real, 0},
  {d1, _Real, 0}, {d2, _Real, 0}, {d3, _Real, 0}},
 (a - d1)^2 + (b - d2)^2 + (c - d3)^2,
 "RuntimeOptions" -> {"EvaluateSymbolically" -> False}
];

and you can then use

With[{minimizer = ... Evaluate@myHi2[...] ..., ...}, ...]

If you weren't to use "RuntimeOptions" -> {"EvaluateSymbolically" -> False}, you would see a message and the code may possibly be compiled incorrectly, especially if you have stray values in your session or the real myHi2 contains procedural code.

Manual inlining

Well, I don't mean completely manually. Here is a function that does what CompilationOptions -> {"InlineCompiledFunctions" -> True} does, without requiring that option (this was very useful prior to version 8, and indeed you can probably see that this code is of a certain vintage).

ClearAll[EvaluatePattern];
EvaluatePattern[expr_, patt_, opts___] :=
  Module[{holdTemporary},
   If[Hold /. {opts} /. Options[EvaluatePattern, Hold],
    SetAttributes[holdTemporary, HoldAll]
   ];
   Identity @@ ReplaceRepeated[
    holdTemporary[expr],
    p : patt :> With[{eval = p}, eval /; True],
    MaxIterations -> (
     MaxIterations /. {opts} /. Options[EvaluatePattern, MaxIterations]
    )
   ]
  ];
Attributes[EvaluatePattern] = {HoldFirst};
Options[EvaluatePattern] = {MaxIterations -> $IterationLimit, Hold -> False};

That was just a wrapper that gives more control over the Trott-Strzebonski device. It can be done more concisely (and, I think, clearly) using RuleCondition directly, but I didn't know about RuleCondition when I wrote the above, and as it's undocumented others may not recognise it either. Here comes the function we require in the present context:

ClearAll[InlineFunctionCalls];
InlineFunctionCalls[expr_, opts___] :=
 EvaluatePattern[
  expr, 
  sym_Symbol /; MemberQ[{Function, CompiledFunction, LibraryFunction}, Head[sym]],
  Sequence @@ Thread[
    {MaxIterations, Hold} -> (
      {MaxIterations, Hold} /. {opts} /.
       Options[InlineFunctionCalls, {MaxIterations, Hold}]
      )
    ]
  ];
Attributes[InlineFunctionCalls] = {HoldFirst};
Options[InlineFunctionCalls] = {MaxIterations -> $IterationLimit, Hold -> True};

Now we can just write

myHi2 = Compile[...];
With[{minimizer = ... // InlineFunctionCalls, ...}, ...]

which will catch and inline all external Function, CompiledFunction, and LibraryFunction objects (and more reliably so than Compile itself).

Edit

I forgot to mention why the final example doesn't compile. The reason is that Compile is HoldAll and it doesn't do a lot of analysis of its argument to find out whether it is compilable or not--if it's not on the list, it's not coming in. In this case Compile can't tell that NelderMeadMinimize`Dump`CompiledNelderMead produces a compiled function, so here it's your responsibility to generate the closure first:

orgFitter = With[{
    minimizer = NelderMeadMinimize`Dump`CompiledNelderMead[
      Function[{a, b, c}, (a - d1)^2 + (b - d2)^2 + (c - d3)^2], {a, b, c},
      "ReturnValues" -> "OptimizedParameters"
     ],
    epsilon = $MachineEpsilon
   },
   Compile[{{d1, _Real, 0}, {d2, _Real, 0}, {d3, _Real, 0}},
    minimizer[RandomReal[{0, 1}, {3 + 1, 3}], epsilon, -1],
    CompilationOptions -> {
     "InlineCompiledFunctions" -> True, "InlineExternalDefinitions" -> True
    },
    RuntimeOptions -> {"Speed", "EvaluateSymbolically" -> False}
   ]
  ];