From 2c80583e145b32f675b46832f28c8bb42e6d51aa Mon Sep 17 00:00:00 2001
From: Roelof van Dijk <3604013+roelofvandijk@users.noreply.github.com>
Date: Wed, 26 Jun 2024 20:13:14 +0200
Subject: [PATCH] perf: cache const UOp creation [run_process_replay] (#5156)

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

diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py
index d9971f1c8e..9fb4a0ed1f 100644
--- a/tinygrad/codegen/uops.py
+++ b/tinygrad/codegen/uops.py
@@ -60,6 +60,7 @@ class UOp:
   def max(self, x): return UOp.alu(BinaryOps.MAX, self, x)
   def min(self, x): return -UOp.alu(BinaryOps.MAX, -self, -x)
   @staticmethod
+  @functools.lru_cache(maxsize=None)
   def const(dtype:Optional[DType], b:ConstType|Variable):
     if isinstance(b, Variable): return UOp(UOps.DEFINE_VAR, dtype, (), b)
     return UOp(UOps.CONST, dtype, arg=dtypes.as_const(b, dtype) if dtype is not None else b)