From 2f06eccf1da766306626e9d5ee36042e26c135e2 Mon Sep 17 00:00:00 2001
From: nimlgen <138685161+nimlgen@users.noreply.github.com>
Date: Sat, 25 Jan 2025 00:33:00 +0300
Subject: [PATCH] am: script and vfio msg (#8742)

* am: script and vfio msg

* use sysfs bars always for now

* tiny chnages
---
 extra/amdpci/am_smi.py           | 4 ++--
 extra/amdpci/setup_python_cap.sh | 3 +++
 extra/amdpci/setup_vfio.sh       | 2 ++
 tinygrad/runtime/ops_amd.py      | 7 ++-----
 4 files changed, 9 insertions(+), 7 deletions(-)
 create mode 100755 extra/amdpci/setup_python_cap.sh
 create mode 100755 extra/amdpci/setup_vfio.sh

diff --git a/extra/amdpci/am_smi.py b/extra/amdpci/am_smi.py
index 9c82db3ca9..f5ffbc6da3 100644
--- a/extra/amdpci/am_smi.py
+++ b/extra/amdpci/am_smi.py
@@ -12,14 +12,14 @@ def bold(s): return f"\033[1m{s}\033[0m"
 def color_temp(temp):
   if temp >= 87: return colored(f"{temp:>4}", "red")
   elif temp >= 80: return colored(f"{temp:>4}", "yellow")
-  return colored(f"{temp:>4}", "white")
+  return f"{temp:>4}"
 
 def color_voltage(voltage): return colored(f"{voltage/1000:>5.3f}V", "cyan")
 
 def draw_bar(percentage, width=40, fill='█', empty='░'):
   filled_width = int(width * percentage)
   bar = fill * filled_width + empty * (width - filled_width)
-  return f'[{bar}] {percentage*100:.1f}%'
+  return f'[{bar}] {percentage*100:5.1f}%'
 
 def same_line(strs:list[list[str]], split=8) -> list[str]:
   ret = []
diff --git a/extra/amdpci/setup_python_cap.sh b/extra/amdpci/setup_python_cap.sh
new file mode 100755
index 0000000000..2ef3c5f6b7
--- /dev/null
+++ b/extra/amdpci/setup_python_cap.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+PYTHON_PATH=$(readlink -f $(which python3))
+sudo setcap 'cap_dac_override,cap_sys_rawio,cap_sys_admin=ep' $PYTHON_PATH
diff --git a/extra/amdpci/setup_vfio.sh b/extra/amdpci/setup_vfio.sh
new file mode 100755
index 0000000000..6ca786edaa
--- /dev/null
+++ b/extra/amdpci/setup_vfio.sh
@@ -0,0 +1,2 @@
+#!/bin/bash
+sudo modprobe vfio-pci disable_idle_d3=1
diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py
index e1a821d624..360e408cf3 100644
--- a/tinygrad/runtime/ops_amd.py
+++ b/tinygrad/runtime/ops_amd.py
@@ -464,7 +464,7 @@ class PCIIface:
 
         iommu_group = HWInterface.readlink(f"/sys/bus/pci/devices/{self.pcibus}/iommu_group").split('/')[-1]
       except OSError:
-        if DEBUG >= 1: print(f"am {self.pcibus}: failed to init vfio-pci module (not inserted or no-iommu mode is not supported).")
+        if DEBUG >= 1: print(f"am {self.pcibus}: failed to init vfio-pci module (run `sudo modprobe vfio-pci`).")
         PCIIface.vfio = False
 
     # Init vfio for the device
@@ -498,10 +498,7 @@ class PCIIface:
                   'array_count': 12, 'simd_arrays_per_engine': 2, 'lds_size_in_kb': 64}
 
   def _map_pci_range(self, bar, off=0, addr=0, size=None):
-    if PCIIface.vfio:
-      vfio.VFIO_DEVICE_GET_REGION_INFO(self.vfio_dev, reg:=vfio.struct_vfio_region_info(argsz=ctypes.sizeof(vfio.struct_vfio_region_info), index=bar))
-      fd, sz, off = self.vfio_dev, size or reg.size, reg.offset + off
-    else: fd, sz = self.bar_fds[bar], size or self.pcidev.regions[bar].size
+    fd, sz = self.bar_fds[bar], size or self.pcidev.regions[bar].size
     return to_mv(fd.mmap(addr, sz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if addr else 0), off), sz)
 
   def alloc(self, size:int, host=False, uncached=False, cpu_access=False):