From 299fa8f37b7ec2352f52f0dad60ee9a6945752ff Mon Sep 17 00:00:00 2001
From: nimlgen <138685161+nimlgen@users.noreply.github.com>
Date: Tue, 28 Jan 2025 01:15:56 +0300
Subject: [PATCH] am: unset high clocks for sleep (#8775)

---
 tinygrad/runtime/support/am/amdev.py |  2 ++
 tinygrad/runtime/support/am/ip.py    | 10 ++++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/tinygrad/runtime/support/am/amdev.py b/tinygrad/runtime/support/am/amdev.py
index 2f541d8766..9b14e024f8 100644
--- a/tinygrad/runtime/support/am/amdev.py
+++ b/tinygrad/runtime/support/am/amdev.py
@@ -290,11 +290,13 @@ class AMDev:
       ip.init()
       if DEBUG >= 2: print(f"am {self.devfmt}: {ip.__class__.__name__} initialized")
 
+    self.smu.set_clocks(perf=True)
     self.gfx.set_clockgating_state()
     self.reg("regSCRATCH_REG7").write(am_version)
     if DEBUG >= 2: print(f"am {self.devfmt}: boot done")
 
   def fini(self):
+    self.smu.set_clocks(perf=False)
     for ip in [self.sdma, self.gfx]: ip.fini()
 
   def paddr2cpu(self, paddr:int) -> int: return mv_address(self.vram) + paddr
diff --git a/tinygrad/runtime/support/am/ip.py b/tinygrad/runtime/support/am/ip.py
index b27b30579b..421b53bcfd 100644
--- a/tinygrad/runtime/support/am/ip.py
+++ b/tinygrad/runtime/support/am/ip.py
@@ -111,10 +111,6 @@ class AM_SMU(AM_IP):
     self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_SetDriverDramAddrLow, lo32(self.adev.paddr2mc(self.driver_table_paddr)), poll=True)
     self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_EnableAllSmuFeatures, 0, poll=True)
 
-    for clck in [0x00000C94, 0x000204E1, 0x000105DC, 0x00050B76, 0x00070B76, 0x00040898, 0x00060898, 0x000308FD]:
-      self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_SetSoftMinByFreq, clck, poll=True)
-      self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_SetSoftMaxByFreq, clck, poll=True)
-
   def is_smu_alive(self):
     with contextlib.suppress(RuntimeError): self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_GetSmuVersion, 0, timeout=100)
     return self.adev.mmMP1_SMN_C2PMSG_90.read() != 0
@@ -129,6 +125,12 @@ class AM_SMU(AM_IP):
     return table_t.from_buffer(to_mv(self.adev.paddr2cpu(self.driver_table_paddr), ctypes.sizeof(table_t)))
   def read_metrics(self): return self.read_table(smu_v13_0_0.SmuMetricsExternal_t, smu_v13_0_0.TABLE_SMU_METRICS)
 
+  def set_clocks(self, perf):
+    # TODO: Parse from bios.
+    for clck, (mn, mx) in {smu_v13_0_0.PPCLK_GFXCLK: (0, 3220), smu_v13_0_0.PPCLK_UCLK: (0, 1249), smu_v13_0_0.PPCLK_FCLK: (0, 2301)}.items():
+      self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_SetSoftMinByFreq, clck << 16 | (mx if perf else mn), poll=True)
+      self._smu_cmn_send_smc_msg_with_param(smu_v13_0_0.PPSMC_MSG_SetSoftMaxByFreq, clck << 16 | (mx if perf else mn), poll=True)
+
   def _smu_cmn_poll_stat(self, timeout=10000): self.adev.wait_reg(self.adev.mmMP1_SMN_C2PMSG_90, mask=0xFFFFFFFF, value=1, timeout=timeout)
   def _smu_cmn_send_msg(self, msg, param=0):
     self.adev.mmMP1_SMN_C2PMSG_90.write(0) # resp reg