diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py index 98e12e21bc..6f82bbd70a 100644 --- a/tinygrad/codegen/uopgraph.py +++ b/tinygrad/codegen/uopgraph.py @@ -286,7 +286,7 @@ constant_folder = PatternMatcher([ (NOp.var('x') * 0, lambda x: x.const(float('nan') if isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)), # x-x -> 0 (NOp.var('x') - NOp.var('x'), lambda x: x.const(0)), - # min==max -> CONST + # min==max -> CONST (slow!) (UPat({UOps.ALU, UOps.DEFINE_VAR}, name='x'), lambda x: x.const(x.vmin.arg) if x.vmin.arg == x.vmax.arg else None), # ** load/store folding ** (NOp.store(NOp.var("buf"), NOp.var("idx"), NOp.load(NOp.var("buf"), NOp.var("idx"))), lambda buf,idx:UOp(UOps.NOOP)), diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 685a1cdc91..64931618fc 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -37,11 +37,14 @@ class PtrDType(DType): class dtypes: @staticmethod - def is_float(x: DType) -> bool: return x.scalar() in (dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64) + @functools.lru_cache(None) + def is_float(x: DType) -> bool: return x.scalar() in {dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64} @staticmethod # static methds on top, or bool in the type info will refer to dtypes.bool - def is_int(x: DType) -> bool: return x.scalar() in (dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.pyint) or dtypes.is_unsigned(x) + @functools.lru_cache(None) + def is_int(x: DType) -> bool: return x.scalar() in {dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.pyint} or dtypes.is_unsigned(x) @staticmethod - def is_unsigned(x: DType) -> bool: return x.scalar() in (dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64) + @functools.lru_cache(None) + def is_unsigned(x: DType) -> bool: return x.scalar() in {dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64} @staticmethod def from_py(x) -> DType: if x.__class__ is float: return dtypes.default_float diff --git a/tinygrad/ops.py b/tinygrad/ops.py index f6ee8786fd..ba8da80804 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -193,9 +193,9 @@ class UOp: if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1] if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1 return None # generic None if we aren't sure - @functools.cached_property + @property def vmin(self) -> UOp: return x if (x:=self._min_max[0]) is not None and not math.isnan(x.arg) else self.sconst(dtypes.min(cast(DType, self.dtype))) - @functools.cached_property + @property def vmax(self) -> UOp: return x if (x:=self._min_max[1]) is not None and not math.isnan(x.arg) else self.sconst(dtypes.max(cast(DType, self.dtype))) @functools.cached_property def _min_max(self) -> Tuple[Optional[UOp], Optional[UOp]]: