sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

add x86 loops (#17179)

Browse Source 
* add x86 loops

* lint
This commit is contained in:
Raine
2026-07-24 07:26:09 -07:00
committed by GitHub
parent c9b60caf8c
commit fcb1d3fc9f
2 changed files with 29 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_wait_loop.py
-3
View File
@@ -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)
tinygrad/renderer/isa/x86.py
+29 -9
View File
@@ -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