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

Question

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

Answer 1

将this answer 中的信息与this answer 中的信息结合起来，可以找到您需要的大部分内容。

我们将使用第一个答案来说明如何获得设备计算能力以及流式多处理器的数量。我们将使用第二个答案（转换为 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
$