运行时环境层

通过 TVM 来检查是否支持 ROCm 设备（Radeon Open Compute），并输出设备的相关信息。

• 函数 check_rocm 用于检查当前是否有可用的 ROCm 设备。 tvm.runtime.device("rocm", 0) 用于获取第 0 个 ROCm 设备（设备 ID 为 0），并返回一个设备对象。
• device.exist 用于检查获取的设备是否存在。如果设备不可用（device.exist 为 False），则输出 "ROCm 设备不可用" 并返回 None，表示没有找到可用的 ROCm 设备。
• 调用 check_rocm 函数并将返回值赋给变量 device。如果 device 为 None，说明没有可用的 ROCm 设备。
• 如果设备存在（device 不为 None），则创建一个 device_info 字典，包含设备的类型 (device.device_type) 和设备 ID (device.device_id)。 device.device_type: 返回设备的类型，这里应该是 "rocm"。 device.device_id: 返回设备的 ID，这里应该是 0，因为设备是通过 tvm.runtime.device("rocm", 0) 获取的。

代码：

import tvm

# 检查 TVM 是否支持 ROCm 并返回设备详细信息
def check_rocm():
    try:
        # 获取 ROCm 设备
        device = tvm.runtime.device("rocm", 0)
        
        # 检查设备是否可用
        if not device.exist:
            print("ROCm 设备不可用")
            return None

        print("ROCm check success")
        return device
    except Exception as e:
        print(f"ROCm check failed: {e}")
        return None

device = check_rocm()

# 获取当前可用的设备并输出设备信息
if device:
    device_info = {
        "device_type": device.device_type,
        "device_id": device.device_id
    }
    print("Device Info:")
    for key, value in device_info.items():
        print(f"  {key}: {value}")

结果：

ROCm check success
Device Info:
  device_type: rocm
  device_id: 0