Shared Memory

Shared Memory

我们知道，CPU和GPU组成异构计算架构，如果想从内存上优化程序，我们必须尽量减少主机与GPU设备间的数据拷贝，并将更多计算从主机端转移到GPU设备端，我们要尽量在设备端初始化数据，并计算中间数据，并尽量不做无意义的数据回写。

GPU的内存结构如图所示：GPU的计算核心都在Streaming Multiprocessor（SM）上，SM里有计算核心可直接访问的寄存器（Register）和共享内存（Shared Memory）；多个SM可以读取显卡上的显存，包括全局内存（Global Memory）。每个SM上的Shared Memory相当于该SM上的一个缓存，一般都很小，Telsa V100的Shared Memory也只有96KB。注意，Shared Memory和Global Memory的字面上都有共享的意思，但是不要将两者的概念混淆，Shared Memory离计算核心更近，延迟很低；Global Memory是整个显卡上的全局内存，延迟高。

从软件角度来看，CUDA的线程可以访问不同级别的存储，每个Thread有独立的私有内存；每个Block中多个Thread都可以在该Block的Shared Memory中读写数据；整个Grid中所有Thread都可以读写Global Memory。Shared Memory的读写访问速度会远高于Global Memory。内存优化一般主要利用Shared Memory技术。下文将以矩阵乘法为例，展示如何使用Shared Memory来优化程序。

普通矩阵乘法

一个C = AB的矩阵乘法运算，需要我们把A的某一行与B的某一列的所有元素一一相乘，求和后，将结果存储到结果矩阵C的(row, col)上。在这种实现中，每个线程都要读取A的一整行和B的一整列，共计算M行*P列。以计算第row行为例，计算C[row, 0]、C[row, 1]…C[row, p-1]这些点时都需要从显存的Global Memory中把整个第row行读取一遍。可以算到，A矩阵中的每个点需要被读 B.width 次，B矩阵中的每个点需要被读 A.height 次。这样比较浪费时间。因此，可以将多次访问的数据放到Shared Memory中，减少重复读取的次数，并充分利用Shared Memory的延迟低的优势。

from numba import cuda
import numpy as np
import math
from time import time

@cuda.jit
def matmul(A, B, C):
    """  矩阵乘法 C = A * B
    """
    # Numba库提供了更简易的计算方法
    # x, y = cuda.grid(2)