How to force Compile to return multiple results?

Question

I would like to return two separate values from a compiled function, but Mathematica refuses to use the compiled version. Here is a narrowed-down example:

    cFunc = Compile[{{a, _Integer, 1}}, {a, a.a}];

(* 
  ==> CompiledFunction::cfex: Could not complete external evaluation at instruction 3; proceeding with uncompiled evaluation. >>

  ==> {{1, 1, 1}, 3}
*)

Here a and a.a are hypothetic intermediate results of some long computation, that must be externalized in order for computation to go on outside the compiled function. I assume it is not possible to return lists that are not uniform (data type and dimensions). Is there any way to return more than one results from a compiled computation?

Answer 1

What about making the result uniform inside Compile and constructing things back afterwards? Like

   cFunc = Compile[
       {{a, _Integer, 1}}, 
       Join @@ {a, {a . a}}]; 
cFunc[{1, 3}]
Function[z, {Most[z], 
       Last[z]}][%]

which has the nice feature of not calling MainEvaluate:

    Needs["CompiledFunctionTools`"]; cFunc // CompilePrint

(* 
==>
    1 argument
        2 Integer registers
        3 Tensor registers
        Underflow checking off
        Overflow checking off
        Integer overflow checking on
        RuntimeAttributes -> {}

        T(I1)0 = A1
        I0 = 4
        Result = T(I1)2

1   I1 = DotVV[ T(I1)0, T(I1)0, I0]]
2   T(I1)1 ={ I1 }
3   T(I1)2 = Join[ T(I1)0, T(I1)1]]
4   Return

*)

Answer 2

You can use a technique like this:

cf = Compile[{{n, _Integer}, {m, _Integer}}, 
   Module[{cpos = RandomSample[Range[n], m]},
    Set[pos, cpos];
    RandomReal[1, Length[cpos]]]];

{valres, posres} = Block[{pos}, {cf[10, 5], pos}]

The set works on a variable which is scoped with Block. Since this is done only once, the call to MainEvaluate is not a problem.

Answer 3

A couple of ideas:

cfun = Compile[{{a, _Integer, 1}}, x[1] = a; x[2] = a.a;]
fun[a_?VectorQ] :=
 Block[{x},
  cfun[a];
  {x[1], x[2]}
 ]

fun[{1,2,3,4,5}]

{{1, 2, 3, 4, 5}, 55}

cf2 = Compile[{{a, _Integer, 1}}, Sow[a]; Sow[a.a];];

Reap[cf2@{1, 2, 3, 4, 5}][[2, 1]]

{{1, 2, 3, 4, 5}, 55}

Answer 4

I have been encountering this problem a lot recently. Since many related questions are linked to this post and I will share my solution here.

Advantages

• Can return multiple results (not necessarily have the same shape)

• Can return ragged lists

• Does not call MainEvaluate

Basic idea

Flatten all the tensors into one single list, and include enough information to reconstruct them.

The first element of my list is the number of tensors / variables to return.

The $2$nd to $2 + var - 1$ th element corresponds to the rank of each tensor

The $2 + var$ to the $2 + var + rank_i -1$ th elements corresponds to the dimension of each tensor

Construction inside Compile

1. Multiple return with different-dimension tensors (of different types)

Note:In my example there is no Complex or True|False, but since Re, Im and Boole are all compilable, they can be transformed to a real tensor and a integer tensor respectively.

This example illustrates returning 3 tensors with different dimensions.

cf1=Compile[{},
  Module[{
   m={{0,8,1,7},{1,9,2,6}},
   n={0.301,0.98},
   p={{{1,0},{2,7}},{{2,0},{0,0}}}},
  Join[{3},
   {TensorRank[m]},{TensorRank[n]},{TensorRank[p]},
   Dimensions[m],Dimensions[n],Dimensions[p],
   Flatten@m,Flatten@n,Flatten@p]
 ]
]

2. Return a ragged list (of arbitrary length)

This is a common case when a collection of Positions should be returned. This example illustrates adding 1D list of arbitrary length to the result programmatically.

cf2=Compile[{},Module[{var=0,rank={},dim={},res={},temp},
  Do[temp=RandomReal[{0,1},RandomInteger[{1,10}]];
   var++;
   AppendTo[rank,TensorRank[temp]];
   dim=Join[dim,Flatten@Dimensions[temp]];
   res=Join[res,Flatten@temp];
  ,{i,1,3}];
 Join[{var},rank,dim,res]]]

Neither of the examples have MainEvaluate when examining with CompilePrint.

Extracting the lists

extractLists[list_?VectorQ] := 
 Module[{vars = Round@First@list, rank, dim}, 
  rank = Round@list[[2 ;; 1 + vars]]; 
  dim = Round@
    Internal`PartitionRagged[
     list[[2 + vars ;; 1 + vars + Total@rank]], rank]; 
  MapThread[
   ArrayReshape, {Internal`PartitionRagged[
     list[[2 + vars + Total@rank ;;]], Times @@@ dim], dim}]]

The results (the result of cf2 is random):

extractLists[cf1[]]
(*{{{0., 8., 1., 7.}, {1., 9., 2., 6.}}, {0.301, 
  0.98}, {{{1., 0.}, {2., 7.}}, {{2., 0.}, {0., 0.}}}}*)
extractLists[cf2[]]
(*{{0.895086, 0.716247, 0.626751, 0.457065, 0.709812, 0.118539, 
  0.504491, 0.40369}, {0.2376}, {0.159539, 0.398285, 0.0233042, 
  0.246191, 0.351316, 0.580408}}*)

Notes

The type of the result is not conserved (Integer is converted to Real). This can be implemented by adding extra parameter before the rank info (I did not include it because it's not useful for my cases). Also I am not sure whether the performance of the code inside Compile is optimal. Feel free to edit if there are improvements.