Efficiently collecting results inside a compiled function

Question

When we don't know the number of results that will be generated, the usual way to collect results is Reap/Sow. Another alternative is linked lists. Neither of these are available in compiled functions.

AppendTo does work, but it has an $O(n^2)$ complexity, so it will be unacceptably slow for long result lists.

There was a very clever suggestion to use Internal`Bag in these situations:

• Internal`Bag inside Compile

Unfortunately a compiled Bag will only hold scalars. Let's try to use vectors:

cf = Compile[{}, 
  Module[{bag = Internal`Bag[{{0, 0, 0}}]}, 
   Do[Internal`StuffBag[bag, {i, i, i}], {i, {1, 2, 3}}]; 
   Internal`BagPart[bag, All]
  ]]

CompilePrint shows that this calls MainEvaluate, so this does not work.

What is the best way to collect a large number of results in a compiled function when the number of results is not known before the computation and the result type is not a (fixed size) vector or matrix?

---

Benchmarking the answers

• Andy's answer

  cf = Compile[{len},
         Module[{bag = Internal`Bag[Most[{0}]]}, 
           Do[Internal`StuffBag[bag, {i, i, i}, 1], {i, len}];
           Partition[Internal`BagPart[bag, All], 3]
         ]
       ];
  
  Do[cf[500000], {100}]; // Timing
  
  (* ==> {2.87, Null} *)
  

  I needed to initialize the Bag using bag = Internal`Bag[Most[{0}]] to let Compile know that it is holding integers, not reals (see here).

• Leonid's answer

  cf2 = 
    Compile[{len},
     Module[{arr, lim, ctr},
      arr = ConstantArray[{0, 0, 0}, 10];
      lim = Length[arr];
      ctr = 1;
      Do[
       If[ctr == lim,
        arr = Join[arr, Table[{0, 0, 0}, {lim}]];
        lim = Length[arr]];
       arr[[ctr++]] = {i, i, i},
  
       {i, len}
      ];
      Take[arr, ctr - 1]
     ]
    ];
  
  Do[cf2[500000], {100}]; // Timing
  
  (* ==> {16.474, Null} *)
  

Comparing the computational complexity of the two solutions by direct measurement:

data = Table[
   {Round[2^k], First@AbsoluteTiming@Do[cf[Round[2^k]], {100}]}, 
   {k, 13, 19, 1/2}];
data2 = Table[
   {Round[2^k], First@AbsoluteTiming@Do[cf2[Round[2^k]], {100}]}, 
   {k, 13, 19, 1/2}];

ListLogLogPlot[{data, data2}]

(They're the same.)

Answer 1

Assuming you know the dimensions of the pieces that you want to come out you can always add a second argument to Internal`StuffBag that indicates the rank of the elements going in. The result is still flat so you have to partition after the fact.

cf = Compile[{}, Module[{bag = Internal`Bag[]},
    Do[Internal`StuffBag[bag, {i, i, i}, 1], {i, {0, 1, 2, 3}}];
    Partition[Internal`BagPart[bag, All], 3]]];

Here I'm indicating that vectors will be going in to the bag. I would specify 2 for a matrix. Notice that it no longer calls MainEvaluate.

In[25]:= StringFreeQ[CompilePrint[cf], "MainEvaluate"]

Out[25]= True

In[20]:= cf[]

Out[20]= {{0., 0., 0.}, {1., 1., 1.}, {2., 2., 2.}, {3., 3., 3.}}

Answer 2

As I mentioned in my recent answer, another alternative is to implement a version of a dynamic array inside Compile, as say arr = Table[{0,0},{10}] (collecting vectors of length 2, in this example) . Set up a variable (say lim) which gives the current size limit, initialize it to the initial size of the allocated array, and another one (say ctr) which counts the current maximal used position. Then, you can do something like

If[ctr==lim,
  arr = Join[arr,Table[{0,0},{lim}]];
  lim = Length[arr]
]

or, instead of doubling, use some other array expansion policy (other exponent, or additive, may depend on the problem). This adds one extra instruction (check) in your inner loop, but potentially saves you memory, and, unlike Internal`Bag, can be returned from Compile. This won't work when your results are lists of different lengths (dimensions), however.