nv: fix fault info (#8587)

* nv: fix fault info

* and emu for amd

* skip if not mock

test/mockgpu/amd/amddriver.py

@@ -1,4 +1,4 @@
-import pathlib, re, ctypes, mmap, collections, functools, copy
+import pathlib, re, ctypes, mmap, collections, functools, copy, os
 import tinygrad.runtime.autogen.kfd as kfd
 from tinygrad.helpers import from_mv
 from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
@@ -49,6 +49,7 @@ class AMDDriver(VirtDriver):
     self.object_by_handle = {}
     self.doorbells = {}
     self.next_doorbell = collections.defaultdict(int)
+    self.mmu_event_ids = []
 
     for i in range(gpus): self._prepare_gpu(i)
 
@@ -114,6 +115,8 @@ class AMDDriver(VirtDriver):
     elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_EVENT:
       struct.event_slot_index = self._alloc_next_event_slot()
       struct.event_id = struct.event_slot_index
+
+      if struct.event_type == kfd.KFD_IOC_EVENT_MEMORY: self.mmu_event_ids.append(struct.event_id)
     elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_QUEUE:
       gpu = self.gpus[struct.gpu_id]
       if struct.queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
@@ -126,7 +129,12 @@ class AMDDriver(VirtDriver):
       struct.doorbell_offset = self._alloc_doorbell(struct.gpu_id)
       self.track_address(struct.doorbell_offset, struct.doorbell_offset + 8, lambda mv,off: None, lambda mv, off: self._emulate_execute())
     elif nr == kfd_ioctls.AMDKFD_IOC_WAIT_EVENTS:
-      pass
+      evs = (kfd.struct_kfd_event_data * struct.num_events).from_address(struct.events_ptr)
+      for ev in evs:
+        if ev.event_id in self.mmu_event_ids and "MOCKGPU_EMU_FAULTADDR" in os.environ:
+          ev.memory_exception_data.gpu_id = 1
+          ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
+          ev.memory_exception_data.failure.NotPresent = 1
     else:
       name = "unknown"
       for k,v in kfd_ioctls.__dict__.items():

test/mockgpu/nv/nvdriver.py

@@ -1,4 +1,4 @@
-import ctypes, mmap, collections, functools
+import ctypes, mmap, collections, functools, os
 import tinygrad.runtime.autogen.nv_gpu as nv_gpu
 from typing import Any
 from tinygrad.helpers import to_mv
@@ -183,6 +183,15 @@ class NVDriver(VirtDriver):
     elif struct.cmd == nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE: pass
     elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_PERF_BOOST: pass
     elif struct.cmd == nv_gpu.NV2080_CTRL_CMD_FB_FLUSH_GPU_CACHE: pass
+    elif struct.cmd == nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_ALL_SM_ERROR_STATES:
+      params = nv_gpu.NV83DE_CTRL_DEBUG_READ_ALL_SM_ERROR_STATES_PARAMS.from_address(params_ptr)
+      params.mmuFault.valid = bool("MOCKGPU_EMU_FAULTADDR" in os.environ)
+    elif struct.cmd == nv_gpu.NV83DE_CTRL_CMD_DEBUG_READ_MMU_FAULT_INFO:
+      params = nv_gpu.struct_NV83DE_CTRL_DEBUG_READ_MMU_FAULT_INFO_PARAMS.from_address(params_ptr)
+      params.count = 1
+      params.mmuFaultInfoList[0].faultAddress = int(os.environ['MOCKGPU_EMU_FAULTADDR'], base=16)
+      params.mmuFaultInfoList[0].faultType = 1
+      params.mmuFaultInfoList[0].accessType = 1
     else: raise RuntimeError(f"Unknown {struct.cmd} to rm_control")
     return 0
 

test/test_hcq.py

@@ -1,4 +1,4 @@
-import unittest, ctypes, struct
+import unittest, ctypes, struct, os
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import getenv
 from tinygrad.device import Buffer, BufferSpec
@@ -486,5 +486,18 @@ class TestHCQ(unittest.TestCase):
 
       assert buf2.as_buffer()[0] == i
 
+  @unittest.skipUnless(MOCKGPU, "Emulate this on MOCKGPU to check the path in CI")
+  def test_on_device_hang(self):
+    if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
+
+    os.environ["MOCKGPU_EMU_FAULTADDR"] = "0xDEADBEE1"
+
+    # Check api calls
+    with self.assertRaises(RuntimeError) as ctx:
+      self.d0.on_device_hang()
+
+    assert "0xDEADBEE1" in str(ctx.exception)
+    os.environ.pop("MOCKGPU_EMU_FAULTADDR")
+
 if __name__ == "__main__":
   unittest.main()

tinygrad/runtime/ops_nv.py

@@ -540,7 +540,7 @@ class NVDevice(HCQCompiled[NVSignal]):
     if sm_errors.mmuFault.valid:
       mmu_info = rmctrl.debug_read_mmu_fault_info(self.fd_ctl, self.root, self.debugger)
       for i in range(mmu_info.count):
-        pfinfo = mmu_info.mmuFaultInfolist[i]
+        pfinfo = mmu_info.mmuFaultInfoList[i]
         report += [f"MMU fault: 0x{pfinfo.faultAddress:X} | {NV_PFAULT_FAULT_TYPE[pfinfo.faultType]} | {NV_PFAULT_ACCESS_TYPE[pfinfo.accessType]}"]
         if DEBUG >= 5:
           report += ["GPU mappings:\n"+"\n".join(f"\t0x{x:X} - 0x{x+y-1:X} | {self._debug_mappings[(x,y)]}" for x,y in sorted(self._debug_mappings))]