From 10f115fdb0c00b6420fa4eb63201583404394f2e Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 6 May 2025 17:09:36 +0300 Subject: [PATCH] usbgpu: USB_RESCAN_BUS envvar (#10177) --- tinygrad/runtime/support/amd.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tinygrad/runtime/support/amd.py b/tinygrad/runtime/support/amd.py index 9086d1c176..7863ad9d25 100644 --- a/tinygrad/runtime/support/amd.py +++ b/tinygrad/runtime/support/amd.py @@ -1,7 +1,7 @@ import functools, importlib, time from collections import defaultdict from dataclasses import dataclass -from tinygrad.helpers import getbits, round_up +from tinygrad.helpers import getbits, round_up, getenv from tinygrad.runtime.autogen import pci from tinygrad.runtime.support.usb import ASM24Controller @@ -33,8 +33,10 @@ def import_module(name:str, version:tuple[int, ...], version_prefix:str=""): raise ImportError(f"Failed to load autogen module for {name.upper()} {'.'.join(map(str, version))}") def setup_pci_bars(usb:ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, tuple[int, int]]: - try: usb.pcie_cfg_req(pci.PCI_VENDOR_ID, bus=gpu_bus, dev=0, fn=0, size=2) - except RuntimeError: + try: need_reset = (usb.pcie_cfg_req(pci.PCI_VENDOR_ID, bus=gpu_bus, dev=0, fn=0, size=2) != 0x1002) + except RuntimeError: need_reset = True + + if need_reset or getenv("USB_RESCAN_BUS", 0) == 1: for bus in range(gpu_bus): usb.pcie_cfg_req(pci.PCI_SUBORDINATE_BUS, bus=bus, dev=0, fn=0, value=gpu_bus, size=1) usb.pcie_cfg_req(pci.PCI_SECONDARY_BUS, bus=bus, dev=0, fn=0, value=bus+1, size=1)