Convert function to exploit parallelization of the GPU

user_1678467 0
Viewed 44

I have a function that uses values stored in one array to operate on another array. This behaves similar to the numpy.hist function. For example:

import numpy as np
from numba import jit

@jit(nopython=True)
def array_func(x, y, output_counts, output_weights):
    for row in range(x.size):
        col = int(x[row] * 10)
        output_counts[col] += 1
        output_weights[col] += y[row]

    return (output_counts, output_weights)

# in the current code these arrays exists ad pytorch tensors
# on the GPU and get converted to numpy arrays on the CPU before 
# being passed to "array_func"
x = np.random.randint(0, 11, (1000)) / 10
y = np.random.randint(0, 100, (10000))

output_counts, output_weights = array_func(x, y, np.zeros(y.size), np.zeros(y.size))

While this works for arrays it does not work for torch tensors that are on the GPU. This is close to what histogram functions do, but I also need the summation of binned values (i.e., the output_weights array/tensor). The current function requires me to continually pass the data from GPU to CPU, followed by the CPU function being run in series.

Can this function be converted to run in parallel on the GPU?

##EDIT##

The challenge is caused by the following line:

output_weights[col] += y[row]

If it weren't for that line I could just use the torch.histc function.

Here's my thought: GPUs are "fast" because they have hundreds/thousands of threads available and can run parts of a big job (or many smaller jobs) on these threads. However, if I convert the function above to work on torch tensors then there is no benefit to running on the GPU (it actually kills the performance). I wonder if there is a way I can break of x so each value gets sent to different threads (similar to how apply_async does within multiprocessing)?

I'm open to other options.

In it's current form the function is fast, but the GPU-to-CPU data transfer is killing me.

numba pytorch
1 Answers

Your computation is indeed a general histogram operation. There are multiple ways to compute this on a GPU regarding the number of items to scan, the size of the histogram and the distribution of the values.

For example, one solution consist in building local histograms in each separate kernel blocks and then perform a reduction. However, this solution is not well suited in your case since len(x) / len(y) is relatively small.

An alternative solution is to perform atomic updates of the histogram in parallel. This solutions only scale well if there is no atomic conflicts which is dependent of the actual input data. Indeed, if all value of x are equal, then all updates will be serialized which is slower than doing the accumulation sequentially on a CPU (due to the overhead of the atomic operations). Such a case is frequent on small histograms but assuming the distribution is close to uniform, this can be fine.

This operation can be done with Numba using CUDA (targetting Nvidia GPUs). Here is an example of kernel solving your problem:

@cuda.jit
def array_func(x, y, output_counts, output_weights):
    tx = cuda.threadIdx.x  # Thread id in a 1D block
    ty = cuda.blockIdx.x   # Block id in a 1D grid
    bw = cuda.blockDim.x   # Block width, i.e. number of threads per block
    pos = tx + ty * bw     # Compute flattened index inside the array
    if pos < x.size:
        col = int(x[pos] * 10)
        cuda.atomic.add(output_counts, col, 1)
        cuda.atomic.add(output_weights, col, y[pos])

For more information about how to run this kernel, please read the documentation. Note that the arrays output_counts and output_weights can possibly be directly created on the GPU so to avoid transfers. x and y should be on the GPU for better performance (otherwise a CPU reduction will be certainly faster). Also note that the kernel should be pretty fast so the overhead to run/wait it and allocate/free temporary array may be significant and even possibly slower than the kernel itself (but certainly faster than doing a double transfer from/to the CPU so to compute things on the CPU assuming data was on the GPU). Note also that such atomic accesses are only fast on quite recent Nvidia GPU that benefit from specific computing units for atomic operations.

user_129395570
Related