元素矩阵乘法可能通过共享内存获得性能提升？

Question

最近几天我刚刚开始使用 Numba 进行 GPU 编程，并且我已经从博客周围的零散信息中学习了一些技术，其中一些在 C programming guide 中，在 Stack 社区中也有很多。

为了简化，我正在尝试提高我的模拟性能，而不是在使用常规Python 代码之前。使用Numba，我已经提高了我的代码的性能，现在在我的 Geforce GTX 1660TI 中运行速度提高了 45 倍，但现在我正试图进一步提高一点，正如 Here 提到的，我的内核没有良好的内存访问模式。

最近我试图了解在某些内核中使用共享内存来提高性能，就像在这个post 中一样，但我不知道这个例子是否对我有帮助，因为据我了解，它具有明显的优势共享内存，在我的常规内核中，我通常使用多个矩阵或向量进行元素乘法。

其实我不知道这是否应该在这里问，所以如果这里不是正确的地方，请原谅我。

我的代码的主要内核之一及其测试实现在下面的代码中

from timeit import default_timer as timer
import numba
from numba import jit, guvectorize, int32, int64, float64, prange
from numba import cuda
import numpy as np
from numpy import *
import math

stream = cuda.stream()

D = 9
nx = 20000
ny = 1000
ly = ny-1
uLB = 0.04

cx = np.array([0, 1,-1, 0, 0, 1,-1, 1,-1],dtype=np.float64);
cy = np.array([0, 0, 0, 1,-1, 1,-1,-1, 1],dtype=np.float64);
c = np.array([cx,cy]);
w = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36],dtype=np.float64);


def inivel(d, x, y):
    return (1-d) * uLB * (1 + 1e-4*sin(y/ly*2*pi))

@cuda.jit
def equilibrium_gpu(rho,u,c,w,feq):
    
    nx2 = rho.shape[0]
    ny2 = rho.shape[1]
    
    cuda.syncthreads()

    j, k = cuda.grid(2)
    if (j < nx2) & (k < ny2):
        for i in range(9):
            feq[i, j, k] = rho[j,k]*w[i] * (1 + (3 * (c[0,i]*u[0,j,k] + c[1,i]*u[1,j,k])) + 0.5*(3 * (c[0,i]*u[0,j,k] + c[1,i]*u[1,j,k]))**2 - (3/2 * (u[0,j,k]**2 + u[1,j,k]**2)))
            
    cuda.syncthreads()

vel = fromfunction(inivel, (2,nx,ny))
rho = np.ones([nx, ny], dtype='float64')
res = np.zeros([D, nx, ny], dtype='float64')
feq = np.zeros((9,nx,ny))

rho_device = cuda.to_device(rho, stream=stream)
u_device = cuda.to_device(vel, stream=stream)
c_device = cuda.to_device(c, stream=stream)
w_device = cuda.to_device(w, stream=stream)
feq_device = cuda.device_array(shape=(D,nx,ny,), dtype=np.float64, stream=stream)

threadsperblock = (16, 16)
blockspergrid_x = math.ceil(nx / threadsperblock[0])
blockspergrid_y = math.ceil(ny / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)

s = timer()
cuda.synchronize()
equilibrium_gpu[blockspergrid, threadsperblock,stream](rho_device,u_device,c_device,w_device,feq_device)
cuda.synchronize()
gpu_time = timer() - s

print(gpu_time)

我想知道如何通过共享内存或其他方式来提高这个内核的性能。

Answer 1

您可以将w[i] 和c[1,i] 放在共享内存中，但我怀疑这会明显更快，因为全局内存访问应该已经缓存在 L1 中。其他访问似乎没问题。

最大的性能问题来自于使用几乎无法有效计算双精度数的主流 GPU 计算。事实上，Geforce GTX 1660TI 可以实现高达 4.6 TFLOPS 的单精度数字，而只有 0.144 TFLOPS 的双精度数字。在这个 GPU 上，这种简单精度的计算速度提高了大约 32 倍！

请注意，有些 GPU 能够进行非常高效的双精度计算，例如 Nvidia Volta GPU，但价格也高得多（即使在这些 GPU 上，双精度仍然慢 2 倍）。

除此之外，最好自己展开循环，将一些变量放入寄存器中以免重新计算它们，并且在生成的代码方面更喜欢a*a 而不是a**2。此外，syncthreads 调用似乎毫无用处（而且代价高昂）。

---

编辑：我错过了非常低效的feq[i, j, k] 访问以及j 应该更好地用于连续维度的事实，因为grid 似乎返回值x, y 而不是y, x（其中x 应该是“连续”维度，尽管文档对此确实不清楚）。使用feq[i, k, j] 连续写入f（并读取其他数组，如u）然后转置数组（可能在GPU 上）可能要好得多。我认为优化的数组完全转置比使用共享内存在本地转置块更快。或者，您可以重新设计您的代码，以便根本不需要转置。

Answer 2

该示例具有other answer 中建议的大部分修改，除了切换到单精度（这很容易做到）。在我的特定机器（GTX 960）上，它的运行速度似乎比您的内核快 3 倍。您可能希望更改 nx 和 ny 以匹配您正在运行的任何测试用例。在我看来，我还调整了围绕内核启动的代码，以便更好地进行基准测试。据我所知，并非所有的变化似乎都有意义。例如转换 x**2 -> x*x 似乎并不重要。

from timeit import default_timer as timer
import numba
from numba import jit, guvectorize, int32, int64, float64, prange
from numba import cuda
import numpy as np
from numpy import *
import math

stream = cuda.stream()

D = 9
nx = 4000
ny = 1000
ly = ny-1
uLB = 0.04

cx = np.array([0, 1,-1, 0, 0, 1,-1, 1,-1],dtype=np.float64);
cy = np.array([0, 0, 0, 1,-1, 1,-1,-1, 1],dtype=np.float64);
c = np.array([cx,cy]);
w = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36],dtype=np.float64);


def inivel(d, x, y):
    return (1-d) * uLB * (1 + 1e-4*sin(y/ly*2*pi))

@numba.jit('void(float64[:,:], float64[:,:], float64[:,:], float64[:], float64[:], float64[:], float64[:,:,:])',parallel=True)
def equilibrium_cpu(rho, ux, uy, cx, cy, w, feq ):              # Equilibrium distribution function.
    nx2 = rho.shape[0]
    ny2 = rho.shape[1]

    for i in prange(9):
        for j in prange(0,nx2):
            for k in prange(0,ny2):
                feq[i,j,k] = rho[j,k]*w[i] * (1 + (3 * (cx[i]*ux[j,k] + cy[i]*uy[j,k])) + 0.5*(3 * (cx[i]*ux[j,k] + cy[i]*uy[j,k]))**2 - (3/2 * (ux[j,k]**2 + uy[j,k]**2)))



@cuda.jit
def equilibrium_gpu(rho,u,c,w,feq):

    nx2 = rho.shape[0]
    ny2 = rho.shape[1]

    k,j = cuda.grid(2)
    if (j < nx2) & (k < ny2):
        r = rho[j,k]
        ux = u[0,j,k]
        uy = u[1,j,k]
        for i in range(9):
             cv = (3 * (c[0,i]*ux + c[1,i]*uy))
             feq[i, j, k] = r*w[i] * (1 + cv + 0.5*(cv*cv) - (3/2 * (ux*ux + uy*uy)))

vel = fromfunction(inivel, (2,nx,ny))
rho = np.ones([nx, ny], dtype='float64')
res = np.zeros([D, nx, ny], dtype='float64')
feq = np.zeros((9,nx,ny))

rho_device = cuda.to_device(rho, stream=stream)
u_device = cuda.to_device(vel, stream=stream)
c_device = cuda.to_device(c, stream=stream)
w_device = cuda.to_device(w, stream=stream)
feq_device = cuda.device_array(shape=(D,nx,ny,), dtype=np.float64, stream=stream)

threadsperblock = (32, 32)
blockspergrid_x = math.ceil(ny / threadsperblock[0])
blockspergrid_y = math.ceil(nx / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)

s = timer()
equilibrium_cpu(rho, vel[0], vel[1], cx, cy, w, feq)

cpu_time = timer() - s

print(cpu_time)
equilibrium_gpu[blockspergrid, threadsperblock,stream](rho_device,u_device,c_device,w_device,feq_device)


cuda.synchronize()
s = timer()
equilibrium_gpu[blockspergrid, threadsperblock,stream](rho_device,u_device,c_device,w_device,feq_device)
cuda.synchronize()
gpu_time = timer() - s

print(gpu_time)

s = timer()
feq_host2 = feq_device.copy_to_host()
print(timer()-s)

gain = cpu_time/gpu_time
print(gain)
print(np.allclose(feq_host2, feq))