From aee4ebe52fb813d3f4ef65eb1766f9820946949c Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 31 Jul 2025 19:37:38 -0700 Subject: [PATCH] linearize works --- test/test_tiny.py | 15 ++++++++++----- tinygrad/codegen/linearize.py | 4 ++-- tinygrad/renderer/cstyle.py | 4 +++- tinygrad/schedule/kernelize.py | 33 ++++++++++++++++++++++++++++++--- 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/test/test_tiny.py b/test/test_tiny.py index 864be7614c..ba161ae959 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -38,7 +38,10 @@ class TestTiny(unittest.TestCase): b = Tensor.eye(N).contiguous().realize() c = Tensor.eye(N).contiguous().realize() d = Tensor.eye(N).contiguous().realize() - out = (((a@b).relu()@c).relu()@d).contiguous().realize() + e = Tensor.eye(N).contiguous().realize() + f = Tensor.eye(N).contiguous().realize() + g = Tensor.eye(N).contiguous().realize() + out = (a@b@c@d@e@f@g).contiguous().realize() self.assertListEqual(out.flatten().tolist(), [1.0]*(N*N)) def test_conv2d(self): @@ -50,10 +53,12 @@ class TestTiny(unittest.TestCase): def test_double_conv2d(self): N = 64 a = Tensor.ones(1,4,N,N).contiguous().realize() - w1 = Tensor.ones(16,4,3,3).contiguous().realize() - w2 = Tensor.ones(24,16,3,3).contiguous().realize() - w3 = Tensor.ones(32,24,3,3).contiguous().realize() - out = a.conv2d(w1).conv2d(w2).conv2d(w3).contiguous().realize() + w1 = Tensor.ones(4,4,3,3).contiguous().realize() + w2 = Tensor.ones(4,4,3,3).contiguous().realize() + w3 = Tensor.ones(4,4,3,3).contiguous().realize() + w4 = Tensor.ones(4,4,3,3).contiguous().realize() + w5 = Tensor.ones(4,4,3,3).contiguous().realize() + out = a.conv2d(w1).conv2d(w2).conv2d(w3).conv2d(w4).conv2d(w5).contiguous().realize() # *** randomness *** diff --git a/tinygrad/codegen/linearize.py b/tinygrad/codegen/linearize.py index 5483a5a91b..53fedcb1cc 100644 --- a/tinygrad/codegen/linearize.py +++ b/tinygrad/codegen/linearize.py @@ -90,9 +90,9 @@ class BlockContext: ctx.block_ctxs[u] = _sort_ctx(this_block_ctx) if u.op is not Ops.SINK else () # RANGE/IF add to the next ctx - # STORE/ASSIGN subtract from the next ctx + # STORE/REDUCE_AXIS subtract from the next ctx if u.op in {Ops.RANGE, Ops.IF}: ctx.child_ctxs[u] = _sort_ctx(ctx.block_ctxs[u] + (u,)) - elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([y for y in ctx.block_ctxs[u] if y not in u.src]) + elif u.op in {Ops.STORE, Ops.REDUCE_AXIS, Ops.CONTIGUOUS}: ctx.child_ctxs[u] = tuple([y for y in ctx.block_ctxs[u] if y not in u.src]) return ctx # ***** make blocks ***** diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index ff45e9d609..6f8fa7094c 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -141,7 +141,9 @@ class CStyleLanguage(Renderer): c: defaultdict[str, int] = defaultdict(int) name = "test" for u in uops: - if u.op is Ops.NOOP: continue + if u.op is Ops.NOOP: + if len(u.src): r[u] = r[u.src[0]] + continue if u.op is Ops.SINK: if u.arg is not None: name = u.arg.function_name continue diff --git a/tinygrad/schedule/kernelize.py b/tinygrad/schedule/kernelize.py index b0148489bb..f522844149 100644 --- a/tinygrad/schedule/kernelize.py +++ b/tinygrad/schedule/kernelize.py @@ -422,11 +422,14 @@ def contiguous_create_ranges(ctx:list[int], x:UOp): ranges = [UOp.range(dtypes.int, s, ctx[0]+i) if resolve(s!=1) else UOp.const(dtypes.int, 0) for i,s in enumerate(x.shape)] ctx[0] += len(ranges) mm = UOp(Ops.MAP, dtype=x.src[0].dtype, src=(x.src[0],)+tuple(ranges)) - return x.replace(src=(mm,)+tuple(ranges)) + buf = UOp.new_buffer(x.device, prod(x.shape), x.dtype) + mm2 = UOp(Ops.MAP, dtype=x.src[0].dtype, src=(buf,)+tuple(ranges)) + return UOp(Ops.STORE, src=(mm2, mm)+tuple(ranges)) + #return x.replace(src=(mm,)+tuple(ranges)) def map_reshape(x:UOp): # don't push on the final buffer reshape for readable graph - if x.src[0].src[0].op is Ops.BUFFER: return None + #if x.src[0].src[0].op is Ops.BUFFER: return None acc = 1 to_sum = [] for s,src in list(zip(x.shape, x.src[1:]))[::-1]: @@ -441,6 +444,7 @@ def map_reshape(x:UOp): else: ret.append(UOp.const(dtypes.int, 0)) ret = UOp.sink(*ret).simplify().src[::-1] if len(ret) else () + #return UOp(Ops.MAP, dtype=x.src[0].src[0].dtype, src=(x.src[0].src[0],)+ret) # generic return x.src[0].replace(src=tuple([UOp(Ops.MAP, dtype=s.dtype, src=(s,)+ret) for s in x.src[0].src])) @@ -464,6 +468,7 @@ def map_permute(x:UOp): print(perm, x.src[0].arg) ret = tuple([ret[p] for p in perm]) print(x.src[0].src[0].shape, ret) + #return UOp(Ops.MAP, dtype=x.src[0].src[0].dtype, src=(x.src[0].src[0],)+ret) return x.src[0].replace(src=tuple([UOp(Ops.MAP, dtype=s.dtype, src=(s,)+ret) for s in x.src[0].src])) def map_expand(x:UOp): @@ -493,6 +498,7 @@ def map_shrink(ctx:list[int], x:UOp): ctx[0] += 1 ret[i] = UOp(Ops.CATRANGE, src=tuple(new_ret_i)) mm = UOp(Ops.MAP, r.dtype, src=(r.src[0],)+tuple(ret)) + #return mm # TODO: put the ranges on the shrink? return UOp(Ops.SHRINK, r.dtype, src=(mm,), arg=r.arg) @@ -508,6 +514,14 @@ index_pushing = PatternMatcher([ lambda x: x.src[0].replace(src=tuple([UOp(Ops.MAP, dtype=s.dtype, src=(s,)+x.src[1:]) for s in x.src[0].src]))), ]) +fix_buffers = PatternMatcher([ + (UPat(Ops.BUFFER, name="x"), lambda x: UOp(Ops.DEFINE_GLOBAL, dtype=x.dtype.ptr(x.arg), arg=x.src[0].arg)), + (UPat(Ops.MAP, name="x"), lambda x: x.replace(op=Ops.INDEX, dtype=x.src[0].dtype)), + (UPat((Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE), name="x"), lambda x: x.src[0]), + (UPat(Ops.EXPAND, name="x"), lambda x: x.replace(arg=None, op=Ops.NOOP)), + (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: x.replace(op=Ops.REDUCE, arg=x.arg[0])), +]) + @track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}") def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]: """ @@ -519,7 +533,20 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]: Returns: Map transforming each UOp in the sink to the Ops.KERNEL graph. """ - graph_rewrite(sink, index_pushing, ctx=[0], name="index pushing", bottom_up=True) + pushed = graph_rewrite(sink, index_pushing, ctx=[0], name="index pushing", bottom_up=True) + pushed = graph_rewrite(pushed, fix_buffers, name="fix buffers") + from tinygrad.codegen.devectorizer import pm_reduce, ReduceContext + pushed = graph_rewrite(pushed, pm_reduce, ctx=ReduceContext(), name="remove reduce") + from tinygrad.codegen.linearize import block_create, BlockContext, pm_blockend_merge, block_merge, pm_finalize + pushed = graph_rewrite(pushed, block_create, ctx=BlockContext.from_sink(pushed), name="block create", bottom_up=True) + pushed = graph_rewrite(pushed, pm_blockend_merge, name="blockend merge") + pushed = graph_rewrite(pushed, block_merge, name="block merge") + pushed = graph_rewrite(pushed, pm_finalize, name="finalize") + from tinygrad.device import Device + try: + print(Device['CPU'].renderer.render(pushed.arg.lst)) + except Exception as e: + print("render fail", e) # multi + merge_views + simplify tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+sym+replace_contiguous, ctx={}, name="merge_views")