From 81410befc27ee849bdde4cb7623ef0f69d439255 Mon Sep 17 00:00:00 2001
From: nimlgen <138685161+nimlgen@users.noreply.github.com>
Date: Fri, 2 May 2025 15:46:29 +0300
Subject: [PATCH] am: remove sleep from wait_reg (#10139)

* am: remove sleep from wait_reg

* fst

* ooops
---
 tinygrad/runtime/support/am/amdev.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tinygrad/runtime/support/am/amdev.py b/tinygrad/runtime/support/am/amdev.py
index fcea8f7d35..3431cd8799 100644
--- a/tinygrad/runtime/support/am/amdev.py
+++ b/tinygrad/runtime/support/am/amdev.py
@@ -347,9 +347,9 @@ class AMDev:
     self.reg("regBIF_BX_PF0_RSMU_DATA").write(val)
 
   def wait_reg(self, reg:AMRegister, value:int, mask=0xffffffff, timeout=10000) -> int:
-    for _ in range(timeout):
+    start_time = int(time.perf_counter() * 1000)
+    while int(time.perf_counter() * 1000) - start_time < timeout:
       if ((rval:=reg.read()) & mask) == value: return rval
-      time.sleep(0.001)
     raise RuntimeError(f'wait_reg timeout reg=0x{reg.addr:X} mask=0x{mask:X} value=0x{value:X} last_val=0x{rval}')
 
   def _read_vram(self, addr, size) -> bytes: