diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index 4428a71e0a..42ea349a58 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -63,9 +63,7 @@ class Kernel: print(self.ast) raise e - @functools.lru_cache(None) - def ordered_parents(op:UOp) -> List[UOp]: return dedup([item for x in op.src for item in ordered_parents(x)] + [op]) - self.reduceops = dedup([x for x in ordered_parents(self.ast) if x.op is Ops.REDUCE_AXIS]) + self.reduceops = [x for x in self.ast.toposort if x.op is Ops.REDUCE_AXIS] self.vars: List[Variable] = self.ast.variables() # NOTE: this requires a specific order with the [::-1], this is likely a bug