Is there a clean way to parallelize multi-index tables?

Question

The documentation of ParallelTable indicates that if you want to parallelize the evaluation of a table of more than one dimension, such as

Table[f[i,j], {i, 5}, {j, 5}]

it is often not enough to simply use ParallelTable, as this will in general only parallelize over the outermost index.

Parallelization happens along the outermost (first) index:

ParallelTable[Labeled[Framed[f[i, j]], $KernelID], {i, 4}, {j, i}]

---

Using multiple iteration specifications is equivalent to nesting Table functions:

ParallelTable[i + j, {i, 3}, {j, i}]

ParallelTable[Table[i + j, {j, i}], {i, 3}]

In some use cases (such as a relatively long computation to be run for a limited number of instances, say this five by five array, with more cores than each dimension, say eight kernels) this can be a significant under-use of system resources.

This can be overcome by joining the two indices into one multi-index and parallelizing over that, in the form

ParallelTable[
  With[{i=index[[1]], j=index[[2]]},
    f[i,j]
  ],
{index, Flatten[Array[List, {5, 5}], 1] }]

but this feels like a bad hack.

Is there a cleaner way to make ParallelTable throw those bounds to the wind and simply hand out computation to any kernel that's free for it?

Answer 1

I came across this problem today and was unpleasantly surprised about this functionality of ParallelTable. The Method->"FinestGrained" does not wokr as one would naively expect due to the parallelization over the outermost index only. I came up the following solution.

ClearAll[NonNestedParallelTable];
NonNestedParallelTable[a_,b__][balance_:Identity,opts___]:=Module[{args,res,idxs,i,rule},
    args={b}[[All,1]];
    res=Table[Unique[]->args,b];
    idxs=Flatten[res]//balance;
    res=res/.Rule[x_,y_]:>x;
    res/.ParallelTable[i[[1]]->Inactive[With][MapThread[Inactive[Set][#1,#2]&,{args,i[[2]]}],Hold[a]] // Activate // ReleaseHold,{i,idxs}]
]
NonNestedParallelTable//Attributes={HoldFirst};

The main idea is to flatten out the table before computation and reconstruct the non-flat table by means of unique identifier. The opts-argument can be used to pass options to the flat ParallelTable inside the module and the balance-argument can be used to manipulate evaluation order manually. This implementation should work on non-regular arrays of arbitrary dimension. Here an example using three parallel kernels

NonNestedParallelTable[{$KernelID, f[i, j]}, {i, 1, 2}, {j, 1, 3}][]//MatrixForm

ParallelTable[{$KernelID, f[i, j]}, {i, 1, 2}, {j, 1, 3}] // MatrixForm

where we clearly see that ParallelTable does not utilize all kernels while NonNestedParallelTable does. Using the balance-argument of NonNestedParallelTable one can affect the evaluation order manually:

NonNestedParallelTable[{$KernelID, f[i, j]}, {i, 1, 2}, {j, 1, 3}][]//MatrixForm

NonNestedParallelTable[{$KernelID, f[i, j]}, {i, 1, 2}, {j, 1, 3}][Reverse]//MatrixForm

NonNestedParallelTable[{$KernelID, f[i, j]}, {i, 1, 2}, {j, 1, 3}][(RandomSample[#, Length[#]] &)] // MatrixForm

Using (RandomSample[#, Length[#]] &) could be used as a naive load balancing method in case of significantly differing computation times for f[i,j].

One disclaimer: I have not tested the implementation of NonNestedParallelTable to my satisfaction and I can not guarantee that it works for all functions and ranges. I will start using the method for my parallel computations/tables and will report (and hopefully fix) errors should I encounter some.