CUPY：矩阵向量乘法比向量向量乘法和小尺寸的 l2norm 更快

Question

我正在将我的 CPU 代码传输到 GPU。 在优化它时，我发现了一个有争议的性能行为：

考虑一个计算向量 L2 范数的简单任务。 对于具有大量元素的向量，我的性能按预期扩展，但是对于少量（256）它不是：

import cupy as cp
a=cp.random.rand(256)

%timeit cp.linalg.norm(a)
32.3 µs ± 159 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

现在，让我们将其与矩阵向量点积进行比较：

b=cp.random.rand(256,256)
%timeit cp.dot(a,b)
8.36 µs ± 80.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

你可以看到出乎意料的矩阵向量乘积快了 4 倍。 为什么会这样？

我开始研究这个玩具问题。 首先，我创建了自定义缩减 kernel：

l2norm = cp.ReductionKernel('T x','T y',  'x * x','a + b', 'y = sqrt(a)', '0', 'l2norm')

使用这个 kernel，我的执行时间约为 17 微秒，比 linalg.norm 好两倍，但仍然比矩阵向量点积差两倍。 我确信这个 kernel 优化得非常好，因此 C++ Thurst 实现不会做得更好。

我也试过用cp.sqrt(a.dot(a))计算范数。 我发现这是非常低效的，因为向量向量点积a.dot(a)比矩阵向量积a.dot(b)花费更长的时间！！！

我确实明白，对于这个小问题，性能受到带宽限制，因此大部分时间可以用于创建 arrays、复制/获取数据，而不是算术。 但即使在这种情况下，我也希望 L2 范数比矩阵向量乘积快一点，因为它只需要 O(N) 次操作和获取，结果是一个数字。 在矩阵向量积的情况下，我什至不预先分配结果，我执行 N^2 次操作并从 memory 中获取 O(N^2) 个数字。

对于大量元素（>1000 个元素），性能按预期扩展。

Ubuntu 18.05，anaconda 分布，python 3.8.3，cupy 8.2.，nvcc 11.0

Answer 1

首先，您只是在测量 CPU 时间，内核是异步执行的，并且您的测量仅包括准备 kernel 启动的部分时间，但您不是在等待实际的 kernel 执行。

如果我们通过使用cupyx.time.repeat测量来更改代码以考虑到这一点，我们会得到

import cupy as cp
import cupyx

a = cp.random.rand(256)
cp.linalg.norm(a)
print(cupyx.time.repeat(cp.linalg.norm, (a,)))
b = cp.random.rand(256, 256)
print(cupyx.time.repeat(cp.dot, (a, b)))
c = cp.zeros(())
l2norm = cp.ReductionKernel(
    "T x", "T y", "x * x", "a + b", "y = sqrt(a)", "0", "l2norm"
)
print(cupyx.time.repeat(l2norm, (a, c)))

结果是

norm                :    CPU:   32.077 us   +/- 2.206 (min:   30.961 / max:   64.160) us     GPU-0:   36.275 us   +/- 2.223 (min:   34.880 / max:   68.512) us
dot                 :    CPU:    9.572 us   +/- 0.261 (min:    9.235 / max:   15.934) us     GPU-0:   13.640 us   +/- 0.347 (min:   12.896 / max:   21.440) us
l2norm              :    CPU:   10.216 us   +/- 0.578 (min:    9.847 / max:   23.790) us     GPU-0:   14.396 us   +/- 0.591 (min:   13.504 / max:   27.936) us

cupy.linalg.norm 正在启动几个内核来计算范数，因此 CPU 时间高达 32 us，而 GPU 时间累积为 36 us。 这里的数组大小是如此之小，以至于这主要是添加了几个内核的恒定开销。

dot 只是调用 cublas function，因此它的 cpu 时间大大减少，GPU 时间非常快，但是随着尺寸的减小，这纯粹是开销。

最后，由于生成实际 kernel 所需的步骤，您的缩减 kernel 有更多的 CPU 时间，但 gpu 的执行与点积大致相同。

如果我们将 arrays 大小增加到 4096，结果如下：

norm                :    CPU:   31.637 us   +/- 2.200 (min:   30.487 / max:   62.955) us     GPU-0:   35.741 us   +/- 2.215 (min:   34.336 / max:   67.008) us
dot                 :    CPU:    9.547 us   +/- 3.753 (min:    9.051 / max:  370.309) us     GPU-0:  244.535 us   +/- 3.791 (min:  241.952 / max:  598.624) us
l2norm              :    CPU:   10.170 us   +/- 0.542 (min:    9.845 / max:   17.006) us     GPU-0:   16.106 us   +/- 0.725 (min:   15.168 / max:   29.600) us

请注意，GPU 执行时间仅针对点积发生变化，这与您的观察结果一致：)。 对于其他内核，与初始开销相比，实际 kernel 执行时间的大小仍然太小。