From e645a0e7661e751ef23b0877df831a207707e7e7 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Wed, 11 Sep 2024 14:32:46 +0800
Subject: [PATCH] allow selecting UPat files in TRACK_MATCH_STATS
 [run_process_replay] (#6470)

---
 tinygrad/ops.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index 94672625cf..d4b36ef7f5 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -735,6 +735,7 @@ if TRACK_MATCH_STATS:
     ret = [0,0,0.0,0.0]
     for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]):
       loc_str = f"{k.location[0].split('/')[-1]}:{k.location[1]}"
+      if getenv("UPAT_FILE", loc_str) not in loc_str: continue
       print(f"{v[0]:6d} / {v[1]:7d} -- {v[3]*1000.:9.2f} / {v[2]*1000.:9.2f} ms -- {loc_str:15s}", k.printable())
       ret = [x+y for x,y in zip(ret, v)]
     print(f"{ret[0]:6d} / {ret[1]:7d} -- {ret[3]*1000.:9.2f} / {ret[2]*1000.:9.2f} ms -- TOTAL")