From 39924387b1f2ecec2d5d132aca7560ba6becda25 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 21 Jul 2026 18:41:31 -0700 Subject: [PATCH] LOOP is srcless RANGE (kimi) (#17129) * LOOP is srcless RANGE (kimi) * upd * cleanups * cleanups * we don't float anymore --- test/backend/test_wait_loop.py | 104 ++++++++++++++++++++++++++++ test/unit/test_wait_loop.py | 56 --------------- tinygrad/codegen/late/linearizer.py | 8 +-- tinygrad/codegen/opt/postrange.py | 5 +- tinygrad/codegen/simplify.py | 2 +- tinygrad/renderer/__init__.py | 8 +-- tinygrad/renderer/cstyle.py | 6 +- tinygrad/renderer/llvmir.py | 10 +-- tinygrad/renderer/ptx.py | 7 +- tinygrad/runtime/ops_python.py | 9 ++- tinygrad/uop/__init__.py | 2 +- tinygrad/uop/ops.py | 8 +-- tinygrad/uop/render.py | 2 +- tinygrad/uop/spec.py | 8 +-- tinygrad/uop/symbolic.py | 2 +- tinygrad/viz/serve.py | 2 +- 16 files changed, 143 insertions(+), 96 deletions(-) create mode 100644 test/backend/test_wait_loop.py delete mode 100644 test/unit/test_wait_loop.py diff --git a/test/backend/test_wait_loop.py b/test/backend/test_wait_loop.py new file mode 100644 index 0000000000..65ee164740 --- /dev/null +++ b/test/backend/test_wait_loop.py @@ -0,0 +1,104 @@ +import unittest +from tinygrad import Tensor, UOp +from tinygrad.device import Device +from tinygrad.dtype import AddrSpace, dtypes +from tinygrad.renderer.nir import NIRRenderer +from tinygrad.renderer.isa.x86 import X86Renderer +from tinygrad.uop.ops import KernelInfo + +def wait_loop_kernel(C:UOp) -> UOp: + N = 10 + + # a RANGE with no src is a bound-less loop header: a jump target with no induction variable. + # the compare and conditional backedge are expanded by the renderers from the loop RANGE/END + l = UOp.loop(0) + + i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG) + + # i = 0 + i = i.after(i[0].store(0)) + + # i + 1, read loop-carried through after(l) + inc = i.after(l)[0].load() + 1 + + # i = inc; END(store, l, cond): conditional backedge, loop again while inc < N (do-while) + # NOTE: the cond uses the computed value, not a reload of the register + st = i[0].store(inc) + i = i.after(st.end(l, inc < N)) + + return C[0].store(i[0].load()).sink(arg=KernelInfo(name="wait_loop")) + +def nested_loop_kernel(C:UOp) -> UOp: + r = UOp.range(4, 0) + l = UOp.loop(1) + + i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG) + i = i.after(i[0].store(0)) + + inc = i.after(l, r)[0].load() + 1 + st = i[0].store(inc) + + lend = st.end(l, inc < (r.cast(dtypes.int)+1)*3) + i = i.after(lend.end(r)) + + return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=())) + +def two_loops_kernel(C:UOp) -> UOp: + # two sequential loops on the same counter: ++ until 10, then ++ until 25 + l1, l2 = UOp.loop(0), UOp.loop(1) + + i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG) + i = i.after(i[0].store(0)) + + inc1 = i.after(l1)[0].load() + 1 + i = i.after(i[0].store(inc1).end(l1, inc1 < 10)) + + inc2 = i.after(l2)[0].load() + 1 + i = i.after(i[0].store(inc2).end(l2, inc2 < 25)) + + return C[0].store(i[0].load()).sink(arg=KernelInfo(name="two_loops", opts_to_apply=())) + +def loop_in_loop_kernel(C:UOp) -> UOp: + # outer loop while i < 12, inner loop increments until i % 4 == 0 -> 12 + l1, l2 = UOp.loop(0), UOp.loop(1) + + i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG) + i = i.after(i[0].store(0)) + + inc = i.after(l1, l2)[0].load() + 1 + st = i[0].store(inc) + + # the outer END closes the inner END, and its cond reloads the register after the inner loop (in scope at the outer level) + e2 = st.end(l2, inc % 4 != 0) + oc = i.after(e2)[0].load() + i = i.after(e2.end(l1, oc < 12)) + + return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=())) + +@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, X86Renderer)), "loops are not supported in LVP and X86") +class TestWaitLoop(unittest.TestCase): + def test_wait_loop(self): + c = Tensor.empty(1, dtype=dtypes.int) + c = Tensor.custom_kernel(c, fxn=wait_loop_kernel)[0] + c.realize() + self.assertEqual(c.item(), 10) + + def test_nested_loop_in_range(self): + c = Tensor.empty(1, dtype=dtypes.int) + c = Tensor.custom_kernel(c, fxn=nested_loop_kernel)[0] + c.realize() + self.assertEqual(c.item(), 12) + + def test_two_sequential_loops(self): + c = Tensor.empty(1, dtype=dtypes.int) + c = Tensor.custom_kernel(c, fxn=two_loops_kernel)[0] + c.realize() + self.assertEqual(c.item(), 25) + + def test_loop_in_loop(self): + c = Tensor.empty(1, dtype=dtypes.int) + c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0] + c.realize() + self.assertEqual(c.item(), 12) + +if __name__ == "__main__": unittest.main() diff --git a/test/unit/test_wait_loop.py b/test/unit/test_wait_loop.py deleted file mode 100644 index 0ba7f38088..0000000000 --- a/test/unit/test_wait_loop.py +++ /dev/null @@ -1,56 +0,0 @@ -import unittest -from tinygrad import Tensor, UOp -from tinygrad.dtype import AddrSpace, dtypes -from tinygrad.uop.ops import KernelInfo - -def wait_loop_kernel(C:UOp) -> UOp: - N = 10 - - # LOOP is a bound-less loop header: a jump target with no induction variable. - # the compare and conditional backedge are expanded by the renderers from LOOP/END - l = UOp.loop(0) - - i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG) - - # i = 0 - i = i.after(i[0].store(0)) - - # i + 1, read loop-carried through after(l) - inc = i.after(l)[0].load() + 1 - - # i = inc; END(store, l, cond): conditional backedge, loop again while inc < N (do-while) - # NOTE: the cond uses the computed value, not a reload of the register - st = i[0].store(inc) - i = i.after(st.end(l, inc < N)) - - return C[0].store(i[0].load()).sink(arg=KernelInfo(name="wait_loop")) - -def nested_loop_kernel(C:UOp) -> UOp: - r = UOp.range(4, 0) - l = UOp.loop(1) - - i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG) - i = i.after(i[0].store(0)) - - inc = i.after(l, r)[0].load() + 1 - st = i[0].store(inc) - - lend = st.end(l, inc < (r.cast(dtypes.int)+1)*3) - i = i.after(lend.end(r)) - - return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=())) - -class TestWaitLoop(unittest.TestCase): - def test_wait_loop(self): - c = Tensor.empty(1, dtype=dtypes.int) - c = Tensor.custom_kernel(c, fxn=wait_loop_kernel)[0] - c.realize() - self.assertEqual(c.item(), 10) - - def test_nested_loop_in_range(self): - c = Tensor.empty(1, dtype=dtypes.int) - c = Tensor.custom_kernel(c, fxn=nested_loop_kernel)[0] - c.realize() - self.assertEqual(c.item(), 12) - -if __name__ == "__main__": unittest.main() diff --git a/tinygrad/codegen/late/linearizer.py b/tinygrad/codegen/late/linearizer.py index 1f11cbc733..e16778a367 100644 --- a/tinygrad/codegen/late/linearizer.py +++ b/tinygrad/codegen/late/linearizer.py @@ -27,7 +27,7 @@ def linearize(sink:UOp) -> list[UOp]: case Ops.BUFFER: priority = -17 if u.addrspace == AddrSpace.LOCAL else -18 case Ops.LOAD: priority = -1 # place loads early case Ops.STORE: priority = 1 # place stores late - case Ops.RANGE | Ops.LOOP: priority = 5 # placing RANGE/LOOP is good + case Ops.RANGE: priority = 5 # placing RANGE is good case Ops.END: priority = -5 # placing END is bad case _: priority = 0 # everything else has priority 0 priorities[u] = (run_count, priority, extra) @@ -66,7 +66,7 @@ class CFGContext: if u.op in (Ops.END, Ops.SINK): nesting |= {x:u for x in deps[u] if x.op is Ops.END and (u.op is Ops.SINK or u.src[1] in deps[x]) and x not in nesting} - if u.op in (Ops.RANGE, Ops.LOOP, Ops.END): deps[u][u] = None + if u.op in (Ops.RANGE, Ops.END): deps[u][u] = None self.edges: dict[UOp, UOp] = {} siblings: dict[UOp, list[UOp]] = {} @@ -81,11 +81,11 @@ class CFGContext: self.edges[y.src[1]] = x pm_add_control_flow = PatternMatcher([ - (UPat((Ops.RANGE, Ops.LOOP), name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None), + (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None), ]) def do_split_ends(e:UOp): - ret, backedge = e.src[0], tuple(x for x in e.src[1:] if x.op is Ops.LOOP or x.dtype == dtypes.bool) + ret, backedge = e.src[0], tuple(x for x in e.src[1:] if x.dtype in (dtypes.void, dtypes.bool)) for r in sorted(UOp.sink(*[x for x in e.src[1:] if x not in backedge]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r) return ret.end(*backedge) if len(backedge) else ret diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 952877eb93..79f77595f3 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -21,8 +21,9 @@ class Scheduler: @property def rngs(self): - # always in order by axistype - return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1]) + # always in order by axistype. void RANGEs are loops, not opt axes + return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.dtype is not dtypes.void and u.vmax > 0], + key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1]) @property def shape_len(self) -> int: return len(self.rngs) @property diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 594fb7b735..72f04255be 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -10,7 +10,7 @@ def flatten_range(r:UOp) -> UOp|None: rngs = r.src[off:] if not len(rngs): return None # ranges in the cond should not be ended - backedge = tuple(x for x in rngs if x.op is Ops.LOOP or x.dtype == dtypes.bool) + backedge = tuple(x for x in rngs if x.dtype in (dtypes.void, dtypes.bool)) return r.replace(src=r.src[:off]+tuple(UOp.sink(*[x for x in rngs if x not in backedge]).ranges)+backedge) pm_flatten_range = PatternMatcher([ diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 7edaae369b..2946315da7 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -39,11 +39,11 @@ class Estimates: mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize) if u.op is Ops.RANGE: mult_stack.append(mults) - mults *= cast(sint, u.src[0].ssimplify()) - # SPECIAL are already counted in mults - mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults + if u.dtype is not dtypes.void: # unbounded loop, unknown trip count + mults *= cast(sint, u.src[0].ssimplify()) + # SPECIAL are already counted in mults + mults = mults.substitute({x:x.const_like(0) for x in mults.toposort() if x.op is Ops.SPECIAL}) if isinstance(mults, UOp) else mults elif u.op is Ops.END: mults = mult_stack.pop(-1) - elif u.op is Ops.LOOP: mult_stack.append(mults) # unbounded loop, unknown trip count elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1 elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG: diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 0e5aa76581..367645776b 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -13,10 +13,10 @@ base_rewrite = PatternMatcher([ (UPat(Ops.BUFFER, name="x"), lambda ctx,x: ctx.render_buffer(x)), # range/loop/if/endif + (UPat(Ops.RANGE, dtypes.void, name="x"), lambda ctx,x: "for (;;) {"), (UPat(Ops.RANGE, name="x"), lambda ctx,x: f"for ({ctx.render_dtype(x.dtype)} {ctx[x]} = 0; {ctx[x]} < {ctx[x.src[0]]}; {ctx[x]}++) {{"), - (UPat(Ops.LOOP, name="x"), lambda ctx,x: "for (;;) {"), - (UPat(Ops.END, src=(UPat(), UPat(Ops.LOOP), UPat(name="c", dtype=dtypes.bool))), lambda ctx,c: f" if (!({ctx[c]})) {{ break; }}\n}}"), + (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE), UPat(name="c", dtype=dtypes.bool))), lambda ctx,c: f" if (!({ctx[c]})) {{ break; }}\n}}"), (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"), (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"), @@ -238,7 +238,7 @@ class CStyleLanguage(Renderer): l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "") kernel.append("\n".join(" "*depth + line for line in l.split("\n"))) if prefix: c[prefix] += 1 # if it was used, increment - if u.op in {Ops.IF, Ops.RANGE, Ops.LOOP}: depth += 1 + if u.op in {Ops.IF, Ops.RANGE}: depth += 1 del self.r # NOTE: this relies on bufs dict preserving order diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 686f24d813..9188d3045f 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -96,6 +96,11 @@ base_rewrite = PatternMatcher([ (UPat(Ops.WHERE, name="x"), lambda ctx,x: f" {ctx[x]} = select {ldt(x.src[0].dtype)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}, {ldt(x.src[2].dtype)} {ctx[x.src[2]]}"), + # loop (a RANGE with no src is an unbounded loop header) + (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx,l: f" br label %loop_{ctx[l][1:]}\nloop_{ctx[l][1:]}:"), + (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void, name="l"), UPat(name="c"))), lambda ctx,l,c: + f" br i1 {ctx[c]}, label %loop_{ctx[l][1:]}, label %loop_exit_{ctx[l][1:]}\nloop_exit_{ctx[l][1:]}:"), + # range (UPat(Ops.RANGE, name="r"), lambda ctx,r: f" br label %loop_entry_{range_str(r)}\n" @@ -113,11 +118,6 @@ base_rewrite = PatternMatcher([ f" br label %loop_latch_{range_str(r)}\n" f"loop_exit_{range_str(r)}:"), - # loop - (UPat(Ops.LOOP, name="l"), lambda ctx,l: f" br label %loop_{ctx[l][1:]}\nloop_{ctx[l][1:]}:"), - (UPat(Ops.END, src=(UPat(), UPat(Ops.LOOP, name="l"), UPat(name="c"))), lambda ctx,l,c: - f" br i1 {ctx[c]}, label %loop_{ctx[l][1:]}, label %loop_exit_{ctx[l][1:]}\nloop_exit_{ctx[l][1:]}:"), - # if (UPat(Ops.IF, name="x"), lambda ctx,x: f" br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"), (UPat(Ops.ENDIF, name="x"), lambda ctx,x: f" br label %ifskip_{ctx[x.src[0]][1:]}\nifskip_{ctx[x.src[0]][1:]}:"), diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index f425e5407e..710e43bc2e 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -116,6 +116,9 @@ string_rewrite = PatternMatcher([ # simple (UPat(Ops.BUFFER, name="x"), lambda ctx, x: [] if x.addrspace == AddrSpace.REG else [ f".shared .align 16 .b8 local{x.arg.slot}[{x.max_numel()*x.dtype.itemsize}];", f"mov.u64 {ctx.r[x]}, local{x.arg.slot}[0];"]), + (UPat(Ops.RANGE, dtypes.void, name="l"), lambda ctx, l: f"WAITLOOP_{ctx.uops.index(l)}:"), + (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void, name="l"), UPat(name="c"))), lambda ctx, l, c: + f"@{ctx.r[c]} bra WAITLOOP_{ctx.uops.index(l)};"), (UPat(Ops.RANGE, name="r"), lambda ctx, r: [ f"mov.u32 {ctx.r[r]}, -1;", f"bra END_{ctx.r[r][1:]};", @@ -125,9 +128,6 @@ string_rewrite = PatternMatcher([ ctx.code_for_op[Ops.ADD](ctx.r[r], ctx.r[r], "1", dtypes.int, ctx.types[dtypes.int]), ctx.code_for_op[Ops.CMPLT](ctx.r[x], ctx.r[r], ctx.r[r.src[0]], dtypes.int, ctx.types[dtypes.int]), f"@{ctx.r[x]} bra LOOP_{ctx.r[r][1:]};"]), - (UPat(Ops.LOOP, name="l"), lambda ctx, l: f"WAITLOOP_{ctx.uops.index(l)}:"), - (UPat(Ops.END, src=(UPat(), UPat(Ops.LOOP, name="l"), UPat(name="c"))), lambda ctx, l, c: - f"@{ctx.r[c]} bra WAITLOOP_{ctx.uops.index(l)};"), (UPat(Ops.IF, name="x"), lambda ctx, x: f"@!{ctx.r[x.src[0]]} bra IF_{ctx.r[x.src[0]][1:]}_{ctx.uops.index(x)};"), (UPat(Ops.ENDIF, name="x"), lambda ctx, x: f"IF_{ctx.r[x.src[0].src[0]][1:]}_{ctx.uops.index(x.src[0])}:"), (UPat(Ops.WMMA, name="x"), lambda ctx, x: list(render_wmma(ctx, x))), @@ -214,6 +214,7 @@ class PTXRenderer(Renderer): prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None), Ops.CONST: ("const", None), Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"), Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None)) + if u.op is Ops.RANGE and u.dtype == dtypes.void: prefix = None # loop headers don't have a register if prefix: r[u] = ssa(prefix, u, dtype) l: str|list[str]|None = string_rewrite.rewrite(u, ctx=self) diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index ab43fb40b7..193726db98 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -48,7 +48,6 @@ class PythonProgram: st = time.perf_counter() warp = list(itertools.product(*[range(x) for x in local_size[::-1]])) warp_size = len(warp) - void_ops = {Ops.END, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.STORE, Ops.LOOP} for idxs in itertools.product(*[range(x) for x in global_size[::-1]]): values: dict[UOp, Any] = {} pbufs: list[memoryview] = list(bufs) @@ -57,12 +56,12 @@ class PythonProgram: i = 0 while i < len(self.uops): u = self.uops[i] - src_values = [values[v] for v in u.src if v.op not in void_ops] - src_dtypes = [v.dtype for v in u.src if v.op not in void_ops] + src_values = [values[v] for v in u.src if v.dtype is not dtypes.void] + src_dtypes = [v.dtype for v in u.src if v.dtype is not dtypes.void] if getenv("TRACE"): print(i, u.op, u.dtype, u.arg, src_values, src_dtypes) if u.op is Ops.END: if len(u.src) == 3: - # conditional backedge on LOOP: jump back while the condition is true + # conditional backedge on a loop: jump back while the condition is true if values[u.src[2]][0]: i = self.uop_to_index[u.src[1]] else: i += 1 else: i = self.uop_to_index[u.src[1]] @@ -75,7 +74,7 @@ class PythonProgram: exec_masks.pop() i += 1 continue - if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP, Ops.LOOP): + if u.op in (Ops.BARRIER, Ops.SINK, Ops.NOOP, Ops.GROUP) or (u.op is Ops.RANGE and u.dtype == dtypes.void): # in the python emulator, the warp is always in sync i += 1 continue diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 33561e6562..c1f36ecc2b 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -76,7 +76,7 @@ class Ops(FastEnum): # ** 5 -- control flow / consts / custom ** # control flow ops - BARRIER = auto(); RANGE = auto(); LOOP = auto(); IF = auto(); END = auto(); ENDIF = auto(); WAIT = auto() + BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto(); WAIT = auto() # const. CONST = auto() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b50ca97a0f..c362774a91 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -112,7 +112,7 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None: # here are the dtype production rules, eventually this will go in UOp as a recursive property match op: case Ops.STORE | Ops.CALL | Ops.LINEAR | Ops.SINK | Ops.PROGRAM | Ops.SOURCE | \ - Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | Ops.LOOP | \ + Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | \ Ops.TUPLE | Ops.FUNCTION | Ops.CUSTOM_FUNCTION | Ops.REWRITE_ERROR: # always void return dtypes.void @@ -346,7 +346,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # some ops init the shape case Ops.GETADDR: return () - case Ops.BIND | Ops.RANGE | Ops.SPECIAL | Ops.LOOP: return () + case Ops.BIND | Ops.RANGE | Ops.SPECIAL: return () case Ops.BINARY: return (len(self.arg),) case Ops.BUFFER: if len(self.src): return self.src[0].as_shape @@ -595,7 +595,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs): return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs) @staticmethod - def loop(axis_id:int, *arg): return UOp(Ops.LOOP, src=(), arg=(axis_id,)+arg) + def loop(axis_id:int, *arg): return UOp(Ops.RANGE, dtypes.void, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.LOOP)+arg) @staticmethod def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name) @staticmethod @@ -1025,7 +1025,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.WHERE and dtypes.is_int(self.dtype): return min(self.src[1].vmin, self.src[2].vmin), max(self.src[1].vmax, self.src[2].vmax) # NOTE: returned UOp is assumed to be CONST if self.op is Ops.PARAM and self.arg.vmin_vmax is not None: return self.arg.vmin_vmax - if self.op in (Ops.RANGE, Ops.SPECIAL): return 0, (self.src[0]-1).vmax + if self.op in (Ops.RANGE, Ops.SPECIAL) and self.dtype is not dtypes.void: return 0, (self.src[0]-1).vmax if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value if self.op is Ops.STACK: return min(x.vmin for x in self.src), max(x.vmax for x in self.src) if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index ab6d31693b..6087f3a301 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -34,8 +34,8 @@ def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str: renderer = PatternMatcher([ (UPat(Ops.PARAM, name="x"), lambda x: x.arg.name if x.arg.name is not None else f"p{x.arg.slot}"), (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg), + (UPat(Ops.RANGE, dtypes.void, name="x"), lambda x: f"loop{x.arg[0]}"), (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"), - (UPat(Ops.LOOP, name="x"), lambda x: f"loop{x.arg[0]}"), (UPat(Ops.CONST, name="x"), lambda x: str(x.arg)), (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"), (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]), diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index b85789fd5a..984f6faaf1 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -74,17 +74,15 @@ spec_shared = PatternMatcher([ # CAST (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: isinstance(x.arg, DType)), - # RANGE can be in the big graph now + # RANGE can be in the big graph now. a void RANGE is a bound-less loop header, the arg is an axis id like RANGE (UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \ all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)), (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(dtypes.is_int(y.dtype) for y in x.src[1:]) or None), - # LOOP is a bound-less loop header, the arg is an axis id like RANGE but without an AxisType - (UPat(Ops.LOOP, dtypes.void, name="l"), lambda l: isinstance(l.arg, tuple) and all(isinstance(ra, int) for ra in l.arg)), # END closes RANGEs (UPat(Ops.END, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(u.op is Ops.RANGE for u in x.src[1:]) or None), - # a LOOP-ended END requires a trailing bool condition for the backedge (loop again while true) - (UPat(Ops.END, src=(UPat(), UPat(Ops.LOOP), UPat(dtype=dtypes.bool))), lambda: True), + # a loop-ended END requires a trailing bool condition for the backedge (loop again while true) + (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void), UPat(dtype=dtypes.bool))), lambda: True), # PARAM (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index c55b0e8454..854b30e17b 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -286,7 +286,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)), # only RANGE/IF/STORE/KERNEL have side effects (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+ - tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.LOOP, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE} + tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE} else y.src for y in x.src[1:]]))))), # after with 1 src is just src[0] (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s), diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index db36db4d9f..59f7199974 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -46,7 +46,7 @@ from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphE from tinygrad.dtype import dtypes, AddrSpace uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B", - Ops.RANGE: "#c8a0e0", Ops.LOOP: "#dd88cc", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", + Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", Ops.INDEX: "#D8F9E4", Ops.STACK: "#D8F9E4", Ops.WMMA: "#efefc0", Ops.MULTI: "#f6ccff", Ops.INS: "#eec4ff", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",