运行时环境层

编写了一个程序,列出系统中可用的 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