From 52a8ecb15ebac955623133e8a20b5afb314d0da5 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Thu, 3 Apr 2025 20:50:09 +0800
Subject: [PATCH] record unittest location in process replay [pr] (#9727)

---
 tinygrad/codegen/kernel.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py
index 719f2ea863..5b648dfd7c 100644
--- a/tinygrad/codegen/kernel.py
+++ b/tinygrad/codegen/kernel.py
@@ -689,12 +689,11 @@ class Kernel:
     src = self.opts.render(self.uops)
 
     if CAPTURE_PROCESS_REPLAY:
-      # NOTE: calling traceback.extract_stack() is very slow, recording backtraces isn't included by default yet
-      if getenv("RECORD_TRACEBACKS"):
-        import traceback
-        stack = "\n".join(traceback.format_list(traceback.extract_stack()[:-1]))
-      else: stack = None
-      diskcache_put("kernel_process_replay", str(id(self)), (self.ast, self.opts, self.applied_opts, self.uops[0].arg, stack, ContextVar._cache, src))
+      import sys
+      frm = sys._getframe(1)
+      while (f_back:=frm.f_back) is not None and "unittest" not in f_back.f_code.co_filename: frm = f_back
+      loc = (frm.f_code.co_filename, frm.f_lineno)
+      diskcache_put("kernel_process_replay", str(id(self)), (self.ast, self.opts, self.applied_opts, self.uops[0].arg, loc, ContextVar._cache, src))
 
     # group non-local bufs by the op type (LOAD or STORE) and the buffer arg. take the max access of that buffer in bytes
     # TODO: these max and min don't work on symbolic, and results are very wrong.