Optimizing my Cuda kernel to sum varying index ranges inside a torch tensor

user_9261771 0
Viewed 391

I'd like to write a Cuda kernel to sum given (contiguous) index ranges in an array. For example, the input array is arr=[1]*10 and I want 3 sums - sum(arr[0:2]), sum(arr[2:3]), sum(arr[3:10]), so the output should be [2, 1, 7].

My arrays are large 2-dimensional arrays (so I want to do this summation for each row, with the same indices), dimensions are typically around 1,000 by 100,000 with the index sub-ranges to be summed varying a lot (between 1 and >1,000). The arrays are already on the GPU as Pytorch tensors so moving them back and forth to/from the CPU for this purpose is costly.

I wrote the following Numba Kernel (here with a minimal working example). Basically, each thread is responsible for a single source column. It finds the relevant target column (w.r.t. index ranges) and adds the column to the target.

from numba import cuda
import numpy as np

@cuda.jit
def sum_idxs(arr, idxs, sum_arr):
    pos = cuda.grid(1)
    if pos>=arr.shape[1]: return
    for i in range(len(idxs)):
        if idxs[i]<=pos<idxs[i+1]:
            thread_idx = i
            break
    for i in range(arr.shape[0]):
        cuda.atomic.add(sum_arr, (i, thread_idx), arr[i, pos])

arr = np.ones(shape=(3, 10))
idxs = np.array([0, 2, 3, 10])
sum_arr = np.zeros(shape=(arr.shape[0], len(idxs)-1))
threads_per_block = 32
blocks_per_grid = ceil(arr.shape[1] / threads_per_block)
sum_idxs[threads_per_block, blocks_per_grid](arr, idxs, sum_arr)
print(sum_arr)

which gives the correct result

[[2. 1. 7.]
 [2. 1. 7.]
 [2. 1. 7.]]

and allows me to do keep my tensors on the GPU as desired.

(I've used numpy arrays here for simplicity. In my code I use cuda.as_cuda_array(tensor) for my pytorch tensor)

However, this is still a major performance bottleneck of my code, is there any way to further optimize it?

numba pytorch cuda
1 Answers

Here is one possible approach. Segmented reductions can often be implemented fairly efficiently by using one block per segment (or in this case, we will use one block per row). If the number of segments/rows is large enough, this will tend to saturate the GPU.

The code design I will suggest will use one block per row, and each block will process the 3 segments of that row in order. To process a segment, the block will use a canonical CUDA reduction implemented using a block-stride loop to do the initial data gather.

Here is an example, fixing some things in your code that have been mentioned in the comments (correct grid dimensioning, conversion to float32):

$ cat t73.py
from numba import cuda,float32,int32
import numpy as np
import math
#TPB = threads per block, max of 1024
#TPB must be the power-of-2 expressed in TPBP2, i.e. TPB = 2**TPBP2    
TPB   = 1024
H     = TPB//2
TPBP2 = 10
@cuda.jit
def sum_idxs(arr, idxs, sum_arr):
    pos = cuda.grid(1)
    if pos>=arr.shape[1]: return
    for i in range(len(idxs)):
        if idxs[i]<=pos<idxs[i+1]:
            thread_idx = i
            break
    for i in range(arr.shape[0]):
        cuda.atomic.add(sum_arr, (i, thread_idx), arr[i, pos])

@cuda.jit
def sum_idxs_i(arr, idxs, sum_arr):
    s = cuda.shared.array(shape=(TPB), dtype=float32)
    tx = cuda.threadIdx.x
    row = cuda.blockIdx.x
#process each of the 3 segments in a row
    for j in range(3):
        lower = idxs[j]
        upper = idxs[j+1]
        val = float32(0)
#block-stride loop to gather data from the segment
        for i in range(tx+lower, upper, TPB):
            val += arr[row, i]
#canonical shared-memory parallel reduction
        s[tx] = val
        mid = H
        for i in range(TPBP2):
            cuda.syncthreads()
            if tx < mid:
                s[tx] += s[tx+mid]
            mid >>= 1
        if tx == 0:
            sum_arr[row, j] = s[0]

rows = 1000
cols = 100000
arr = np.ones(shape=(rows, cols),dtype=np.float32)
idxs = np.array([0, 2, 3, cols],dtype=np.int32)
sum_arr = np.zeros(shape=(arr.shape[0], len(idxs)-1),dtype=np.float32)
blocks_per_grid = math.ceil(arr.shape[1] / TPB)
sum_idxs[blocks_per_grid, TPB](arr, idxs, sum_arr)
print(sum_arr)
sum_arr = np.zeros(shape=(arr.shape[0], len(idxs)-1),dtype=np.float32)
blocks_per_grid = (arr.shape[0])
sum_idxs_i[blocks_per_grid, TPB](arr, idxs, sum_arr)
print(sum_arr)
$ nvprof python t73.py
==4383== NVPROF is profiling process 4383, command: python t73.py
[[2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]
 ...
 [2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]]
[[2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]
 ...
 [2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]
 [2.0000e+00 1.0000e+00 9.9997e+04]]
==4383== Profiling application: python t73.py
==4383== Profiling result:
            Type  Time(%)      Time     Calls       Avg       Min       Max  Name
 GPU activities:   45.92%  287.93ms         6  47.988ms  1.1520us  144.09ms  [CUDA memcpy HtoD]
                   44.88%  281.42ms         6  46.903ms  1.4720us  140.74ms  [CUDA memcpy DtoH]
                    8.46%  53.052ms         1  53.052ms  53.052ms  53.052ms  cudapy::__main__::sum_idxs$241(Array<float, int=2, C, mutable, aligned>, Array<int, int=1, C, mutable, aligned>, Array<float, int=2, C, mutable, aligned>)
                    0.75%  4.6729ms         1  4.6729ms  4.6729ms  4.6729ms  cudapy::__main__::sum_idxs_i$242(Array<float, int=2, C, mutable, aligned>, Array<int, int=1, C, mutable, aligned>, Array<double, int=2, C, mutable, aligned>)
      API calls:   43.26%  339.61ms         6  56.602ms  20.831us  193.89ms  cuMemcpyDtoH
                   36.75%  288.52ms         6  48.087ms  15.434us  144.35ms  cuMemcpyHtoD
                   18.66%  146.51ms         1  146.51ms  146.51ms  146.51ms  cuDevicePrimaryCtxRetain
                    0.93%  7.3083ms         5  1.4617ms  4.8120us  6.7314ms  cuMemFree
                    0.23%  1.8049ms         6  300.81us  9.4520us  778.85us  cuMemAlloc
                    0.04%  327.52us         2  163.76us  156.34us  171.19us  cuLinkAddData
                    0.04%  299.72us         2  149.86us  148.92us  150.80us  cuModuleLoadDataEx
                    0.04%  276.32us         2  138.16us  131.16us  145.16us  cuLinkComplete
                    0.02%  123.96us         2  61.978us  61.252us  62.704us  cuLinkCreate
                    0.01%  64.406us         2  32.203us  29.439us  34.967us  cuLaunchKernel
                    0.01%  63.184us         2  31.592us  30.251us  32.933us  cuDeviceGetName
                    0.00%  29.454us         1  29.454us  29.454us  29.454us  cuMemGetInfo
                    0.00%  20.732us        26     797ns     477ns  2.0320us  cuCtxGetCurrent
                    0.00%  12.852us        25     514ns     363ns  1.0920us  cuCtxGetDevice
                    0.00%  12.429us         2  6.2140us  1.7830us  10.646us  cuDeviceGetPCIBusId
                    0.00%  5.0950us        10     509ns     302ns  1.0770us  cuFuncGetAttribute
                    0.00%  3.9600us         2  1.9800us  1.8000us  2.1600us  cuModuleGetFunction
                    0.00%  3.5630us         2  1.7810us  1.7510us  1.8120us  cuLinkDestroy
                    0.00%  1.8970us         1  1.8970us  1.8970us  1.8970us  cuCtxPushCurrent
                    0.00%  1.8370us         4     459ns     226ns     697ns  cuDeviceGet
                    0.00%  1.6080us         6     268ns     181ns     481ns  cuDeviceGetAttribute
                    0.00%  1.5060us         3     502ns     230ns     795ns  cuDeviceGetCount
                    0.00%  1.2390us         2     619ns     428ns     811ns  cuDeviceComputeCapability
$

This code was run on a GTX960, which happens to report ~84GB/s of device memory bandwidth via bandwidthTest CUDA sample code. In the above example, we see that the improved kernel runs in ~4.7ms (about 10x faster than the original atomic kernel) and this translates to (1000*100000*4)bytes/4.7ms ~= 85GB/s, so we can conclude for this specific test case, this kernel is approximately "optimal".

user_16959600
Related