diff --git a/test/test_fusion_op.py b/test/test_fusion_op.py index 372fc46056..8c720f079c 100644 --- a/test/test_fusion_op.py +++ b/test/test_fusion_op.py @@ -27,7 +27,7 @@ class TestFusionOp(unittest.TestCase): def test_recursive_add(self): st = time.perf_counter() a = Tensor([1,2,3,4]) - for _ in range(12): a = a + a + for _ in range(20): a = a + a sched = create_schedule([a.lazydata], None) ji = lower_schedule_item(sched[-1]) et = time.perf_counter() diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index f30a7d45f2..03d1bcaa1e 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -516,7 +516,9 @@ class Linearizer(Kernel): if cachable: self.saved_exprs[key] = ret return ret - def ast_parse(self, x:LazyOp, acc: List[UOp], offs:Optional[List[int]], loaded_buffers:Dict[Union[MemBuffer, ConstBuffer, LocalBuffer], List[UOp]], do_reduce=False, loop_ctx=tuple()) -> List[UOp]: # noqa: E501 + def ast_parse(self, x:LazyOp, acc: List[UOp], offs:Optional[List[int]], loaded_buffers:Dict[Union[MemBuffer, ConstBuffer, LocalBuffer], List[UOp]], do_reduce=False, loop_ctx=tuple(), cache=None) -> List[UOp]: # noqa: E501 + if cache is None: cache = {} + if x in cache: return cache[x] if x.op in BufferOps: return loaded_buffers[x.arg] if x.op == UnaryOps.CAST: return [self.uop(UOps.CAST, x.arg[0], (u,), x.arg) if not isinstance(x.arg[0], ImageDType) else u for u in self.ast_parse(x.src[0], acc, offs, loaded_buffers)] # noqa: E501 if x.op in ReduceOps and not do_reduce: @@ -527,7 +529,7 @@ class Linearizer(Kernel): x = LazyOp(TernaryOps.MULACC, x.src[0].src, x.arg) if x.op == ReduceOps.SUM and x.src[0].__class__ is LazyOp and x.src[0].op == UnaryOps.CAST and x.src[0].src[0].__class__ is LazyOp and x.src[0].src[0].op == BinaryOps.MUL: # noqa: E501 x = LazyOp(TernaryOps.MULACC, x.src[0].src[0].src, x.arg) - values = [self.ast_parse(v, acc, offs, loaded_buffers, loop_ctx=loop_ctx) for v in x.src] + values = [self.ast_parse(v, acc, offs, loaded_buffers, loop_ctx=loop_ctx, cache=cache) for v in x.src] ops = {ReduceOps.SUM:BinaryOps.ADD, ReduceOps.MAX:BinaryOps.MAX, TernaryOps.MULACC:TernaryOps.MULACC} if x.op in ops: ret: List[UOp] = [] @@ -540,4 +542,5 @@ class Linearizer(Kernel): acc[off] = self.uop(UOps.PHI, input_acc[off].dtype, (input_acc[off], acc[off]) + tuple(loop_ctx)) else: ret = [self.uop(UOps.ALU, dtype=dtypes.bool if x.op == BinaryOps.CMPLT else None, vin=val, arg=x.op) for val in zip(*values)] + cache[x] = ret return ret diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 345ef2ce08..c23da40664 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -36,20 +36,20 @@ class LazyBuffer: self.device, self.st, self.dtype, self.shape = device, st, dtype, st.shape if base is None: # properties on base - self.children: WeakSet[LazyBuffer] = WeakSet() - for x in srcs: x.base.children.add(self.base) self.op, self.arg, self.srcs = op, arg, srcs # this is a LazyOp, except the src is LazyBuffers and not LazyOps self.realized: Optional[Buffer] = None self.output_buffer: Optional[Buffer] = None self.forced_realize = False self.contiguous_child: Optional[Tuple[ReferenceType[LazyBuffer], ShapeTracker]] = None + self.children: WeakSet[LazyBuffer] = WeakSet() + for x in srcs: x.base.children.add(self.base) else: # properties on view assert base.base == base, "base must be a base itself" self._base = base def __repr__(self) -> str: - return f"" + return f"" @property def base(self) -> LazyBuffer: return self._base if hasattr(self, '_base') else self @@ -141,7 +141,10 @@ class LazyBuffer: # *** schedule creation *** # recursively create a lazyop -def _recursive_lazyop(buf:LazyBuffer, inputs:List[LazyBuffer], var_vals:Dict[Variable, int], st:ShapeTracker, realizes:Set[LazyBuffer], first=True): +def _recursive_lazyop(buf:LazyBuffer, inputs:List[LazyBuffer], var_vals:Dict[Variable, int], st:ShapeTracker, + realizes:Set[LazyBuffer], first=True, cache=None) -> LazyOp: + if cache is None: cache = {} + if (buf, st) in cache: return cache[(buf, st)] if buf != buf.base: var_vals.update(merge_dicts([var_vals, buf.st.var_vals])) st = buf.st.unbind()+st @@ -161,7 +164,7 @@ def _recursive_lazyop(buf:LazyBuffer, inputs:List[LazyBuffer], var_vals:Dict[Var # if a CONTIGUOUS made it all the way here, just skip it if buf.op == LoadOps.CONTIGUOUS: assert first - return _recursive_lazyop(buf.srcs[0], inputs, var_vals, st, realizes, False) + return _recursive_lazyop(buf.srcs[0], inputs, var_vals, st, realizes, False, cache) # if it's a reduce, we have to change the shapetracker if buf.op in ReduceOps: @@ -169,7 +172,8 @@ def _recursive_lazyop(buf:LazyBuffer, inputs:List[LazyBuffer], var_vals:Dict[Var st = ShapeTracker.from_shape(buf.srcs[0].shape).unbind() # otherwise we fuse it like normal - return LazyOp(buf.op, tuple(_recursive_lazyop(x, inputs, var_vals, st, realizes, False) for x in buf.srcs), buf.arg) + cache[(buf, st)] = ret = LazyOp(buf.op, tuple(_recursive_lazyop(x, inputs, var_vals, st, realizes, False, cache) for x in buf.srcs), buf.arg) + return ret # recursively walk back in the graph to create the schedule def _recursive_schedule(out:LazyBuffer, seen:Set[LazyBuffer], realizes:Set[LazyBuffer],