如何使用 Python 和 Numba 获取 GPU 中的 CUDA 内核数量？

Question

我想知道如何使用 Python、Numba 和 cudatoolkit 获取 GPU 中的 CUDA 核心总数。

Answer 1

大部分你所需要的可以通过在信息结合可以找到这个答案在信息沿着这个答案。

我们将使用第一个答案来指示如何获得设备计算能力以及流式多处理器的数量。我们将使用第二个答案（转换为 python）来使用计算能力来获取每个 SM 的“核心”计数，然后将其乘以 SM 的数量。

这是一个完整的例子：

$ cat t36.py
from numba import cuda


cc_cores_per_SM_dict = {
    (2,0) : 32,
    (2,1) : 48,
    (3,0) : 192,
    (3,5) : 192,
    (3,7) : 192,
    (5,0) : 128,
    (5,2) : 128,
    (6,0) : 64,
    (6,1) : 128,
    (7,0) : 64,
    (7,5) : 64,
    (8,0) : 64,
    (8,6) : 128
    }
# the above dictionary should result in a value of "None" if a cc match 
# is not found.  The dictionary needs to be extended as new devices become
# available, and currently does not account for all Jetson devices
device = cuda.get_current_device()
my_sms = getattr(device, 'MULTIPROCESSOR_COUNT')
my_cc = device.compute_capability
cores_per_sm = cc_cores_per_SM_dict.get(my_cc)
total_cores = cores_per_sm*my_sms
print("GPU compute capability: " , my_cc)
print("GPU total number of SMs: " , my_sms)
print("total cores: " , total_cores)

$ python t36.py
GPU compute capability:  (5, 2)
GPU total number of SMs:  8
total cores:  1024
$