From fcb1d3fc9f47ecb0720dcf63aee3595dd976e51c Mon Sep 17 00:00:00 2001 From: Raine Date: Fri, 24 Jul 2026 10:26:09 -0400 Subject: [PATCH] add x86 loops (#17179) * add x86 loops * lint --- test/backend/test_wait_loop.py | 3 --- tinygrad/renderer/isa/x86.py | 38 ++++++++++++++++++++++++++-------- 2 files changed, 29 insertions(+), 12 deletions(-) diff --git a/test/backend/test_wait_loop.py b/test/backend/test_wait_loop.py index 8a133a7587..83e0543e61 100644 --- a/test/backend/test_wait_loop.py +++ b/test/backend/test_wait_loop.py @@ -3,7 +3,6 @@ from tinygrad import Tensor, UOp from tinygrad.device import Device, Buffer, BufferSpec from tinygrad.dtype import AddrSpace, dtypes from tinygrad.engine.realize import run_linear -from tinygrad.renderer.isa.x86 import X86Renderer from tinygrad.uop.ops import Ops, KernelInfo def wait_loop_kernel(C:UOp) -> UOp: @@ -82,7 +81,6 @@ def loop_in_loop_kernel(C:UOp) -> UOp: return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=())) -@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "loops are not supported in X86") class TestWaitLoop(unittest.TestCase): def test_wait_loop(self): c = Tensor.empty(1, dtype=dtypes.int) @@ -109,7 +107,6 @@ class TestWaitLoop(unittest.TestCase): self.assertEqual(c.item(), 12) @unittest.skipUnless(Device.DEFAULT in ("CPU", "AMD", "NV"), "need proper uncached=True handling") -@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "loops are not supported in X86") class TestVolatileLoops(unittest.TestCase): def test_async_wait_ext(self): sig_buf = Buffer(Device.DEFAULT, 1, dtypes.int, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True) diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index daf44709bd..a19a2fbdc5 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -13,7 +13,7 @@ from tinygrad.helpers import getenv, CPU_COUNT, unwrap, Target class X86Ops(FastEnum): # NOTE: X86Ops with i suffix are variants that take an immediate, m suffix are variants that can write to memory instead of read from # these aren't real instructions, DEFINE is a register placeholder that defines a register without emitting an instruction - FRAME_INDEX = auto(); LABEL = auto(); DEFINE = auto() + FRAME_INDEX = auto(); LABEL = auto(); DEFINE = auto(); LOOP_CMP = auto() # index LEA = auto() # register / memory / immediate moves @@ -328,6 +328,7 @@ def _xmm_sz_m(x: UOp) -> X86Ops: def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None: # register placeholders with real registers if x.arg is X86Ops.DEFINE and x.tag is not None: return None + if x.arg is X86Ops.LOOP_CMP: return None # this is an immediate if x.arg is X86Ops.FRAME_INDEX: return None # no register definition @@ -353,6 +354,9 @@ isel_matcher = PatternMatcher([ # range is lowered to acc, cmp, jmp after regalloc (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])), (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None), + # really all a backedge END is is an IF with a tag referencing the RANGE start label + (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"), + lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])), # **** Op -> X86Op **** # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc # so regalloc builds the prologue/epilogue naturally @@ -562,22 +566,37 @@ pre_regalloc_matcher = PatternMatcher([ # TODO: control flow should be overhauled so that this isn't necessary def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]: loop_label = "_".join(str(i) for i in x.arg[:-1]) - acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:]) label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}") - cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0])) - jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}") - ctx.loop_label[acc] = loop_label - return (acc, [acc, label, cmp, jump_out]) + # loop, cmp on backedge all we need is a jmp tag + if x.dtype is dtypes.void: return (label, [label]) + else: + acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:]) + cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CONST else X86Ops.CMP, src=(acc, x.src[0])) + jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}") + ctx.loop_label[acc] = loop_label + return (acc, [acc, label, cmp, jump_out]) + +def lower_end(ctx, x:UOp) -> tuple[UOp, list[UOp]]: + end_label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}") + jmp = UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}") + inc = x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)) + return (inc, [inc, jmp, end_label]) + +def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]: + cond = x.replace(op=x.tag, src=x.src[:2]) + jmp = isel_matcher.rewrite(UOp(Ops.IF, src=(cond,))) + return (jmp.src[0], [jmp.src[0], jmp.replace(tag=x.src[3].tag)]) # final rewrite to match the isa spec post_regalloc_matcher = PatternMatcher([ # rewrite FRAME_INDEX to IMM now that the stack size is known (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=x.const_like(ctx.stack_size + x.tag), [nx])), + # expand the cmp here so we can preserve rng src edge to get label from ctx + (UPat(Ops.INS, arg=X86Ops.LOOP_CMP, name="x"), lower_loop), # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound - (UPat(Ops.RANGE, name="x"), lambda ctx,x: lower_range(ctx, x)), + (UPat(Ops.RANGE, name="x"), lower_range), # rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target - (UPat(Ops.END, name="x"), lambda ctx,x: (jmp:=UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}"), - [x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),)), jmp, UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")])), + (UPat(Ops.END, name="x"), lower_end), # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]), [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg in X86GroupOp.TwoAddress else None), @@ -845,6 +864,7 @@ class X86Renderer(ISARenderer): binary = bytearray() for u in uops: if u.op is not Ops.INS or u.arg is X86Ops.DEFINE: continue + if u.arg is X86Ops.LOOP_CMP: continue if u.arg is X86Ops.LABEL: targets[u.tag] = len(binary) continue