优化我的 Cuda 内核以对火炬张量内的不同索引范围求和答案

【问题标题】：Optimizing my Cuda kernel to sum varying index ranges inside a torch tensor优化我的 Cuda 内核以对火炬张量内的不同索引范围求和
【发布时间】：2021-07-12 18:38:08
【问题描述】：

我想编写一个 Cuda 内核来对数组中给定的（连续的）索引范围求和。例如，输入数组是 arr=[1]*10，我想要 3 个和 - sum(arr[0:2]), sum(arr[2:3]), sum(arr[3:10])，所以输出应该是 [2, 1, 7]。

我的数组是大型二维数组（所以我想对每一行进行求和，具有相同的索引），维度通常在 1,000 x 100,000 左右，要求和的索引子范围不同 a手数（介于 1 和 >1,000 之间）。这些数组已经作为 Pytorch 张量在 GPU 上，因此为此目的在 CPU 之间来回移动它们的成本很高。

我编写了以下 Numba 内核（这里有一个最小的工作示例）。基本上，每个线程负责一个源列。它找到相关的目标列（w.r.t. 索引范围）并将该列添加到目标。

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)

给出正确的结果

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

并允许我根据需要将张量保留在 GPU 上。

（为了简单起见，我在这里使用了 numpy 数组。在我的代码中，我使用 cuda.as_cuda_array(tensor) 作为我的 pytorch 张量）

但是，这仍然是我的代码的主要性能瓶颈，有没有办法进一步优化它？

【问题讨论】：

您展示的这段代码使用 64 位浮点数据类型。我假设您正在使用的火炬张量不是 64 位的？
@RobertCrovella 它们是 32 位浮点数
该代码中有一些非常明显的错误，这意味着它对于任何非平凡的尺寸都不能像您描述的那样工作。当然不是您在问题中描述的问题大小。就像您在内核调用中反转了块和线程尺寸的事实一样。如果你能提供工作代码，以及你在实际问题规模下获得的性能等，那就太好了。
@talonmies “显然”在旁观者的眼中（显然）。当我编写最小的示例时，块和线程的尺寸以某种方式混合在一起。您能详细说明其他错误吗？我是 Cuda 和这类代码的新手，如果有任何有用的建议，我将不胜感激。

标签： cuda pytorch numba

【解决方案1】：

这是一种可能的方法。分段缩减通常可以通过每段使用一个块来相当有效地实现（或者在这种情况下，我们将每行使用一个块）。如果段/行的数量足够大，这将使 GPU 趋于饱和。

我建议的代码设计将每行使用一个块，每个块将按顺序处理该行的 3 个段。为了处理一个段，该块将使用一个规范的 CUDA 缩减，使用一个块步长循环来执行初始数据收集。

这是一个示例，修复了 cmets 中提到的代码中的一些内容（正确的网格尺寸，转换为 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
$

此代码在 GTX960 上运行，它恰好通过 bandwidthTest CUDA 示例代码报告了大约 84GB/s 的设备内存带宽。在上面的示例中，我们看到改进后的内核运行时间约为 4.7 毫秒（大约比原始原子内核快 10 倍），这转换为(1000*100000*4)bytes/4.7ms ~= 85GB/s，因此我们可以得出结论，对于这个特定的测试用例，这个内核大约是“最佳的” ”。

【讨论】：

一个值得一提的决定性区别是，warp 中的（相邻）线程协作，以便通过一个内存请求访问更大的 128 字节大小的连续内存块，而不是每个 4 字节的 32 个请求。跨度>
我不确定你在说什么。也许你应该研究一下 SASS 代码。原子操作的第一步是加载相关值（arr[i, pos] 会很好地合并）。这将通过一条普通的 LDG 指令在整个 Warp 范围内发生，一次 32 个线程/128 个字节。这部分代码在这两种情况下不会有任何不同。所以我不确定你所说的“决定差异”是什么意思。这里的区别在于是否使用原子，以及该吞吐量与规范共享内存并行减少的比较。