运行时环境层
编写了一个程序,列出系统中可用的 CUDA 设备,获取设备的名称、计算能力和全局内存大小等信息。
- 检查 CUDA 设备: 定义
check_cuda
函数以尝试获取第一个 CUDA 设备并收集其信息。 - 获取设备信息: 包括计算能力、最大线程数、共享内存等。
- 异常处理: 捕获并打印获取设备信息时的错误。
- 输出设备信息: 打印设备的详细信息和名称。
代码:
import tvm
from tvm.contrib import nvcc
# 检查 TVM 是否支持 CUDA 并返回设备详细信息
def check_cuda():
try:
# 尝试获取 CUDA 设备
device = tvm.cuda(0)
# 获取设备的详细信息
device_info = {
"compute_capability": device.compute_version, # 计算能力
"max_threads_per_block": device.max_threads_per_block, # 每个块的最大线程数
"max_shared_memory_per_block": device.max_shared_memory_per_block, # 每个块的最大共享内存
"multi_processor_count": device.multi_processor_count, # 多处理器数量
"warp_size": device.warp_size, # warp 大小
"total_global_memory": device.total_global_memory, # 总全局内存
}
print("CUDA check success")
print("Device Info:")
for key, value in device_info.items():
print(f" {key}: {value}")
return device_info
except Exception as e:
print(f"CUDA check failed: {e}")
return None
check_cuda()
# 获取当前可用的设备
dev = tvm.cuda(0) # 获取第一个 GPU 设备
# 输出设备名称
device_name = dev.device_name
print("Device Name:", device_name)
# 获取设备的详细信息
device_info = {
"Device Type": dev.device_type,
"Device ID": dev.device_id, # 使用 device_id 替代 device_index
}
# 输出设备详细信息
for key, value in device_info.items():
print(f"{key}: {value}")
结果:
CUDA check success
Device Info:
compute_capability: 8.9
max_threads_per_block: 1024
max_shared_memory_per_block: 49152
multi_processor_count: 80
warp_size: 32
total_global_memory: 16852844544
Device Name: NVIDIA GeForce RTX 4080 SUPER
Device Type: 2
Device ID: 0