From efbf4fca055876b3cb40486504a86fa1bbeaac01 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cyril=20Roum=C3=A9gous?=
 <63861105+croumegous@users.noreply.github.com>
Date: Sat, 15 Jun 2024 01:37:27 +0200
Subject: [PATCH] perf: graph_rewrite line reduction and make it a little bit
 faster [run_process_replay] (#4958)

---
 tinygrad/codegen/uops.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py
index 6ceda98164..430e1a7e70 100644
--- a/tinygrad/codegen/uops.py
+++ b/tinygrad/codegen/uops.py
@@ -296,14 +296,10 @@ class UOpGraph:
         changed += recurse_cnt
         # NOTE: this changes UOp, so we have to delete caches
         up.vin = tuple(rewrite(x) for x in up.vin)
-        try: del up.parents
-        except AttributeError: pass
-        try: del up.cmp_tuple
-        except AttributeError: pass
+        if 'parents' in up.__dict__: delattr(up, 'parents')
+        if 'cmp_tuple' in up.__dict__: delattr(up, 'cmp_tuple')
         # replace with cached nodes
-        if found:=self.nodes.get(key:=up.tuple()): return found
-        self.nodes[key] = up
-        return up
+        return self.nodes.setdefault(up.tuple(), up)
       sink = rewrite(sink)
       run_cnt += 1
       assert run_cnt < 100, "exceeded 100 rewrite loops!"