compare amd and am to cpu as well (#9896)

test/external/speed_compare_amd_am.py

@@ -9,7 +9,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_amd import AMDDevice
 from contextlib import contextmanager
 import numpy as np
-import os
+import os, random, statistics
 
 am_signal_pages, am_signal_pool, am_devices = [], [], []
 amd_signal_pages, amd_signal_pool, amd_devices = [], [], []
@@ -27,7 +27,6 @@ def rebind_vfio(pcibus="0000:44:00.0"):
 
   os.system("sudo modprobe amdgpu")
   os.system("rocm-smi --setprofile compute")
-  os.system("rocm-smi --setmclk 3")
   os.system("rocm-smi --setperflevel high")
 
 @contextmanager
@@ -49,7 +48,14 @@ def run_am():
   AMDDevice.signal_pages, AMDDevice.signal_pool, AMDDevice.devices = [], [], []
 
 if __name__ == "__main__":
+  CHECK_CPU = getenv("CHECK_CPU", 0)
+  SEED = getenv("SEED", 42)
+  CNT = getenv("CNT", 7)
+  random.seed(SEED)
+  np.random.seed(SEED)
+
   # TODO: NUM=780 is super slow
+  # NUM=1907 is broken on AMD and AM have some mismatches (0 vs 1)
   # kfd feels so bad when taking gpu out while it's running... Need hacks to rebind it before running.
   rebind_vfio(pcibus="0000:44:00.0")
 
@@ -61,6 +67,8 @@ if __name__ == "__main__":
   with run_amd():
     amddev = Device["AMD"]
 
+  if CHECK_CPU: cpudev = Device["CPU"]
+
   single = getenv("NUM", -1)
   if single != -1: ast_strs = ast_strs[single:single+1]
 
@@ -87,12 +95,23 @@ if __name__ == "__main__":
       if not has_bf16:
         for i,rawbuf in enumerate(test_ambufs): rawbuf.copyin(contents[i])
 
+    if CHECK_CPU:
+      cpu_rdr = cpudev.renderer
+      cpu_rdr.device = "CPU"
+      cpulin = ast_str_to_lin(ast, opts=cpu_rdr)
+      cpulin = hand_coded_optimizations(cpulin)
+      cpu_prg = CompiledRunner(cpulin.to_program())
+      cpubufs = bufs_from_lin(cpulin)
+      test_cpubufs = get_fuzz_rawbufs(cpulin) if not has_bf16 else ambufs
+      if not has_bf16:
+        for i,rawbuf in enumerate(test_cpubufs): rawbuf.copyin(contents[i])
+
     # warmup
     tm_amd, tm_am, failed = [], [], False
     with run_amd():
       try:
         amd_prg(test_amdbufs, {}, wait=True)
-        for i in range(7): tm_amd.append(amd_prg(amdbufs, {}, wait=True))
+        for i in range(CNT): tm_amd.append(amd_prg(amdbufs, {}, wait=True))
       except RuntimeError:
         print("AMD FAILED")
         tm_amd = [1e9]
@@ -101,12 +120,16 @@ if __name__ == "__main__":
     with run_am():
       try:
         am_prg(test_ambufs, {}, wait=True)
-        for i in range(7): tm_am.append(am_prg(ambufs, {}, wait=True))
+        for i in range(CNT): tm_am.append(am_prg(ambufs, {}, wait=True))
       except RuntimeError:
         print("AM FAILED")
         tm_am = [1e9]
         failed = True
 
+    if CHECK_CPU:
+      cpu_prg(test_cpubufs, {}, wait=True)
+      for i in range(1): cpu_prg(cpubufs, {}, wait=True)
+
     if not failed and not has_bf16:
       with run_amd():
         curesult = np.frombuffer(test_amdbufs[0].as_buffer(), _to_np_dtype(test_amdbufs[0].dtype))
@@ -114,10 +137,21 @@ if __name__ == "__main__":
       with run_am():
         amresult = np.frombuffer(test_ambufs[0].as_buffer(), _to_np_dtype(test_ambufs[0].dtype))
 
-      np.testing.assert_allclose(curesult, amresult, rtol=1e-2, atol=1e-2)
+      if CHECK_CPU:
+        cpuresult = np.frombuffer(test_cpubufs[0].as_buffer(), _to_np_dtype(test_cpubufs[0].dtype))
+        np.testing.assert_allclose(amresult, cpuresult, rtol=1e-2, atol=1e-2)
+        np.testing.assert_allclose(curesult, cpuresult, rtol=1e-2, atol=1e-2)
 
-    average_tm_amd += min(tm_amd)
-    average_tm_am += min(tm_am)
-    ratio = min(tm_am)/min(tm_amd)
-    print(f"{average_tm_am/average_tm_amd:5.2f}x -- {num:4d} {colorize_float(ratio)}  {min(tm_am)*1e6:7.2f} vs {min(tm_amd)*1e6:7.2f} us", amlin.name)
+      try:
+        np.testing.assert_allclose(curesult, amresult, rtol=1e-2, atol=1e-2)
+      except AssertionError as e:
+        print("AM and AMD results do not match")
+        print(e)
+
+    bam = statistics.median(tm_am)
+    bamd = statistics.median(tm_amd)
+    average_tm_amd += bamd
+    average_tm_am += bam
+    ratio = bam/bamd
+    print(f"{average_tm_am/average_tm_amd:5.2f}x -- {num:4d} {colorize_float(ratio)} {bam*1e6:7.2f} vs {bamd*1e6:7.2f} us", amlin.name)
     if DEBUG > 3 and ratio > 1.04: print(f"AM slower {ratio}", amlin.ast, amlin.applied_opts)