#!/usr/bin/env python3

import os
from tinygrad.helpers import Context
from tinygrad.runtime.support.system import System, PCIDevice
from tinygrad.runtime.support.hcq import FileIOInterface
from tinygrad.runtime.support.am.amdev import AMDev

if __name__ == "__main__":
  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
  for gpu in gpus:
    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
  pcidevs = [PCIDevice("AM", gpu) for gpu in gpus]
  amdevs = []
  with Context(DEBUG=2):
    for pcidev in pcidevs:
      amdevs.append(AMDev(pcidev, reset_mode=True))
    for amdev in amdevs: amdev.smu.mode1_reset()