Is there any way at all to report / share results from a Numba prange loop?

user_489088 0
Viewed 180

I have a large prange loop that efficiently does some process here. I'd like to track when one of them finds a better result. However, it seems writing to variables inside a prange loop will create a new private variable, so data can't be shared that way.

Here is a canonical example: this prange searches for the largest random value in a number of arrays. I'd like for them to "collaborate" so that only a new highest value is ever printed:

import numpy as np
from numba import jit, prange

@jit(nopython=True, parallel=True)
def test():
    arr = np.random.rand(1000, 15000000)

    largest_value = 0
    for i in prange(arr.shape[0]):
        for val in arr[i]:
            if val > largest_value:
                print('NEW LARGEST VALUE: ' + str(int(largest_value * 100000000)))
                largest_value = val

test()

Running this will print bogus values such as:

...
NEW LARGEST VALUE: 85809181
NEW LARGEST VALUE: 93504938
NEW LARGEST VALUE: 75796671
NEW LARGEST VALUE: 83944676
NEW LARGEST VALUE: 99068766

This is because each thread, when writing to largest_value creates a local private copy, so they can't share their findings that way.

Are there any tricks at all to make this possible? Shared memory? Anything fancy I'm missing? Something else?

numba python
1 Answers

The solution to share data between threads is to create an array before the parallel loop so the threads can then fill the shared array in the parallel loop. Note that you should be careful about false-sharing while writing in the shared array (for sake of performance).

In some cases, you need to write only per-thread informations. In that case, you can allocate an array of size get_num_threads() (possibly bigger to prevent false-sharing) and write at the location provided by get_thread_id(). These functions should be provided by Numba in parallel loops. You can find more information about this in the documentation.

Note that the old/manual way to do that is to use an additional encompassing loop for threadId in prange(get_num_threads()) and use threadId as the index for the shared array. As a result, the iteration space must be manually divided so that the work is evenly shared between thread (which is a bit cumbersome to do). Fortunately, this solution should not be needed anymore (at least in the most simple cases).

user_129395570
Related