diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index 81480b55ed..29e41797ad 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -4,7 +4,7 @@ import itertools, math, functools from collections import defaultdict from enum import Enum, auto -from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, partition, prod, PtrDType +from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, partition, prod, PtrDType, all_same from tinygrad.ops import LazyOp, UnaryOps from tinygrad.ops import ReduceOps, BinaryOps, TernaryOps from tinygrad.runtime.lib import RawConst @@ -407,11 +407,13 @@ class Linearizer(OptimizedKernel): def uop(self, uop:UOps, dtype:Optional[DType], vin:Tuple[UOp, ...], arg:Any=None, cachable=False) -> UOp: key = (uop, dtype, vin, arg) if uop == UOps.STORE and len(vin) == 2 and vin[0] == vin[1]: return vin[0] # self store is noop + if uop == UOps.CAST and all(x.uop == UOps.GEP for x in vin) and all_same([x.vin[0] for x in vin]) and all(x.arg == i for i,x in enumerate(vin)): return vin[0].vin[0] + if uop == UOps.GEP and vin[0].uop == UOps.CONST: return self.const(vin[0].arg, dtype) if uop == UOps.ALU: # rewrites. NOTE: the rewritten NEG op is still around... if arg == BinaryOps.ADD and vin[1].uop == UOps.ALU and vin[1].arg == UnaryOps.NEG: return self.uop(UOps.ALU, dtype, (vin[0], vin[1].vin[0]), BinaryOps.SUB, cachable=cachable) # constant folding - if arg == UnaryOps.NEG and vin[0].uop == UOps.CONST: return self.uop(UOps.CONST, dtype, (), -vin[0].arg, cachable=True) + if arg == UnaryOps.NEG and vin[0].uop == UOps.CONST: return self.const(-vin[0].arg, dtype) # zero folding for x in [0,1]: if arg == BinaryOps.ADD and vin[x].uop == UOps.CONST and vin[x].arg == 0.0: return vin[1-x]