mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

perf: cache const UOp creation [run_process_replay] (#5156)

Browse source
This commit is contained in:
Roelof van Dijk 2024-06-26 20:13:14 +02:00 • committed by GitHub
commit 2c80583e14
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
tinygrad/codegen/uops.py
View file

@ -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)