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

问题描述

我想知道如何使用Python，Numba和cudatoolkit在我的GPU中获得CUDA内核的总数。

解决方法

您可以通过将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,5) : 192,7) : 192,(5,0) : 128,2) : 128,(6,0) : 64,1) : 128,(7,5) : 64,(8,0) : 64
    }
# 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 = getattr(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
$