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

Question

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

Answer 1

Most of what you need can be found by combining the information in this answer along with the information in this answer .

We'll use the first answer to indicate how to get the device compute capability and also the number of streaming multiprocessors. We'll use the second answer (converted to python) to use the compute capability to get the "core" count per SM, then multiply that by the number of SMs.

Here is a full example:

$ 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
$