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

Posted 2023-04-15

【中文标题】如何使用 Python 和 Numba 获取 GPU 中的 CUDA 内核数量？

【英文标题】：How can I get the number of CUDA cores in my GPU using Python and Numba?

【发布时间】：2020-12-28 13:55:22

【问题描述】：

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

【问题讨论】：

这能回答你的问题吗？ How can I get number of Cores in cuda device?

可能在***.com/questions/32530604/…重复

@MichaelJanz 不幸的是，我对使用 Python 和 Numba 有一个非常具体的要求。如果我使用 C，您建议的可能重复解决了问题。这是您提供的一个很好的信息来源，但不是我问题的答案。简而言之，我的问题不是重复的，它特别要求基于 Python 和 Numba 的解决方案。

【参考方案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
$

【讨论】：

这正是我想要的。此脚本的输出与以下输出相匹配：$ nvidia-settings -q CUDACores -t

'COMPUTE_CAPABILITY' 现在分为市长和次要财产：my_cc = (device.COMPUTE_CAPABILITY_MAJOR, device.COMPUTE_CAPABILITY_MINOR)

计算能力一直由主要部分和次要部分组成。在 numba 中，它们是元组的两个部分。作为元组检索是supported。