Can I return lists with different dimensions from a compiled function?

11

4

Sometimes I'll need several lists with different dimensions to be returned by a compiled function, but compiled function will fail if an irregular list contained:

(* The real case is of course much more complex. *)
a1 = ConstantArray[0., {10}];
a2 = ConstantArray[0., {9}];
Compile[{}, 
  (* for the real case, a lot of low level calculations will be done on a1 and a2 here *)
           {a1, a2}][]

CompiledFunction::cflist: Nontensor object generated; proceeding with uncompiled evaluation. >>

A possible (and maybe the simplest) work around is to add some extra elements to the lists and delete them(or just ignore them) later:

a1 = ConstantArray[0., {10}];
a2 = ConstantArray[0., {10}];
Delete[Compile[{}, {a1, a2}][], {2, -1}]

But it's a waste of memory (especially when lists are big and have higher dimensions).

So I wonder if other solution exists? Of course that solution should work on high dimensional lists like:

a1 = ConstantArray[0., {300, 299, 300}];
a2 = ConstantArray[0., {299, 300, 299}];

To avoid possible confusion I'd like to add a more specific example:

ie = 200;
ez = ConstantArray[0., {ie + 1}];
hy = ConstantArray[0., {ie}];

(* Notice the following function hasn't been fixed yet *)
fdtd1d = Compile[{{steps}}, 
   Module[{ie = ie, ez = ez, hy = hy}, 
    Do[
     ez[[2 ;; ie]] = ez[[2 ;; ie]] + (hy[[2 ;; ie]] - hy[[1 ;; ie - 1]]); 
     ez[[1]] = Sin[n/10]; 
     hy[[1 ;; ie]] = hy[[1 ;; ie]] + (ez[[2 ;; ie + 1]] - ez[[1 ;; ie]]), 
      {n, steps}]; 
    {ez, hy}]];

fdtd1d[1000]

Er… maybe I should add one more multidimensional sample?:

ie = 100; je = 100;
c = 1/Sqrt[2];
ez = ConstantArray[0., {ie + 1, je + 1}];
hx = ConstantArray[0., {ie + 1, je}];
hy = ConstantArray[0., {ie, je + 1}];

(* Notice the following function hasn't been fixed yet *)
fdtd2d = Compile[{{steps}}, 
   Module[{ez = ez, hx = hx, hy = hy, ie = ie, je = je, c = c},
    Do[
     ez[[2 ;; ie, 2 ;; je]] += 
      c (hx[[2 ;; ie, 1 ;; je - 1]] - hx[[2 ;; ie, 2 ;; je]] + 
         hy[[2 ;; ie, 2 ;; je]] - hy[[1 ;; ie - 1, 2 ;; je]]);
     ez[[Floor[ie/2], Floor[je/2]]] = Sin[n/10];
     hx += c (ez[[All, 1 ;; je]] - ez[[All, 2 ;; je + 1]]);
     hy += c (ez[[2 ;; ie + 1, All]] - ez[[1 ;; ie, All]]),
     {n, steps}]; 
   {ez, hx, hy}]];

AbsoluteTiming[dat = fdtd2d[1000];]
ArrayPlot /@ {ez2, hx2, hy2}

xzczd

Posted 2013-12-24T13:49:43.990

Reputation: 44 878

2Not the most elegant but could you flatten and join and then "unflatten" and separate afterwards? – s0rce – 2013-12-24T22:23:17.310

@s0rce But Flatten@{a1, a2} won't…… Oh, I see! This is indeed a possible approach! Why not give an answer? – xzczd – 2013-12-25T03:28:20.430

Answers

3

I suppose one could return values through an external call:

ClearAll["f`*", f];
f`compiled = Compile[{{dim1, _Integer}, {dim2, _Integer}},
   Module[{a1, a2},
    a1 = Range[dim1];
    a2 = Range[dim2];
    f`ret1 = a1;
    f`ret2 = a2;
    {Total@a1, Total@a2}
    ]];
f[dim1_, dim2_] := Block[{f`ret1, f`ret2, f`ans},
   f`ans = f`compiled[dim1, dim2];
   {f`ret1, f`ret2, f`ans}
   ];

The data is passed back through via calls to MainEvaluate. There is a cost to that, but perhaps in your actual use case, the cost is not too much.

Example

f[5, 3]

(* {{1, 2, 3, 4, 5}, {1, 2, 3}, {6, 15}} *)

Michael E2

Posted 2013-12-24T13:49:43.990

Reputation: 190 928

Can I summarize this as "using variables that haven't been compiled to store those lists"? And for my new added sample, the realization of this method is to change {ez, hy} into something like ez2 = ez; hy2 = hy? – xzczd – 2013-12-25T04:42:16.580

@xzczd Yes, that's about right. And if the return values are to be stored in e2z and hy2, then I think the answer to your second is yes, too. – Michael E2 – 2013-12-25T04:52:46.697

"perhaps in your actual use case, the cost is not too much." Your guess is right, I tested this method with my new added samples and it's surprising that the external call doesn't cause a significant speed-down. Maybe it's because it just happens at the end of the code in my specific case? – xzczd – 2013-12-25T05:46:40.993

@xzczd Yes, it's single call for each array. It's repeated external calls, or if somehow the arrays get unpacked (or maybe a few other things), that can really slow down a compiled function. – Michael E2 – 2013-12-25T05:50:02.467

4

As far as I know, the only case where this is possible is for compiled functions with RuntimeAttributes -> Listable. Example:

testComp  = 
   Compile[{{lst, _Integer, 1}}, lst^2, RuntimeAttributes -> Listable]

so that, for example:

testComp[Range /@ Range[4]]

(* {{1}, {1, 4}, {1, 4, 9}, {1, 4, 9, 16}} *)

I am not aware of other situations where CompiledFunction may return ragged lists.

Leonid Shifrin

Posted 2013-12-24T13:49:43.990

Reputation: 108 027

Er… my case is a little different, my a1 and a2 isn't irrelative (see my added sample above) so I can't take advantage of RuntimeAttributes -> Listable. – xzczd – 2013-12-25T04:00:46.030

3

2D update

Here is a way to make your example work with minimal overhead by joining the ragged output into a single larger list and storing the dimensions at the start and putting it back together with a post-processing function.

ie = 100; je = 100;
c = 1/Sqrt[2];
ez = ConstantArray[0., {ie + 1, je + 1}];
hx = ConstantArray[0., {ie + 1, je}];
hy = ConstantArray[0., {ie, je + 1}];

fdtd2d = Compile[{{steps, _Integer, 0}, {ie, _Integer, 
     0}, {je, _Integer, 0}, {c, _Real, 0}, {ez2, _Real, 
     2}, {hx2, _Real, 2}, {hy2, _Real, 2}}, 
   Module[{ez = ez2, hx = hx2, hy = hy2, ezdims, hxdims, hydims}, 
    Do[ez[[2 ;; ie, 2 ;; je]] += 
      c (hx[[2 ;; ie, 1 ;; je - 1]] - hx[[2 ;; ie, 2 ;; je]] + 
         hy[[2 ;; ie, 2 ;; je]] - hy[[1 ;; ie - 1, 2 ;; je]]);
     ez[[Floor[ie/2], Floor[je/2]]] = Sin[n/10];
     hx += c (ez[[All, 1 ;; je]] - ez[[All, 2 ;; je + 1]]);
     hy += c (ez[[2 ;; ie + 1, All]] - ez[[1 ;; ie, All]]), {n, 
      steps}];
    ezdims = Dimensions@ez;
    hxdims = Dimensions@hx;
    hydims = Dimensions@hy;
    Join @@ {{3, Length@ezdims, Length@hxdims, Length@hydims}, ezdims, hxdims, hydims, Flatten@ez, Flatten@hx, Flatten@hy}]];

positionindexer[data_] := 
 With[{accumulated = Accumulate@Join[{0}, data]}, 
  Table[{accumulated[[i]] + 1, accumulated[[i + 1]]}, {i, 
    Length@accumulated - 1}]]

postprocess[data_] := Module[{},
  output = data;
  numberofArrays = Round@First@output;
  output = Rest@output;
  numdims = Round@Take[output, numberofArrays];
  output = Drop[output, Length@numdims];
  dimpositions = positionindexer[numdims];
  dims = Round[Take[output, #] & /@ dimpositions];
  totallength = Times @@@ dims;
  output = Drop[output, Total@numdims];
  datapositions = positionindexer[totallength];
  dataflat = Take[output, #] & /@ datapositions;
  MapThread[ArrayReshape, {dataflat, dims}]]

postprocess@fdtd2d[1000, ie, je, c, ez, hx, hy]

s0rce

Posted 2013-12-24T13:49:43.990

Reputation: 9 292