From 83ae83d871a3dcf5c19c6d0c76864b4c17f0f2c9 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Tue, 15 Apr 2025 13:32:18 +0300 Subject: [PATCH] compare amd and am to cpu as well (#9896) --- test/external/speed_compare_amd_am.py | 52 ++++++++++++++++++++++----- 1 file changed, 43 insertions(+), 9 deletions(-) diff --git a/test/external/speed_compare_amd_am.py b/test/external/speed_compare_amd_am.py index 0808253721..1080f295e6 100644 --- a/test/external/speed_compare_amd_am.py +++ b/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)