diff --git a/test/unit/test_encodings.py b/test/unit/test_encodings.py index 535551db3f..f30bb642b1 100644 --- a/test/unit/test_encodings.py +++ b/test/unit/test_encodings.py @@ -3,7 +3,9 @@ from tinygrad.renderer.x86 import X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM from tinygrad.uop import X86Ops, Ops from tinygrad.uop.ops import UOp from tinygrad.dtype import dtypes, DType +from tinygrad.helpers import SPEC +@unittest.skipIf(SPEC > 1, "x86 spec not supported in full_spec") class TestEncodingsX86(unittest.TestCase): # NOTE: x86 supports a single displacement as memory address and index without base memory address # these have no use cases so they aren't supported diff --git a/test/unit/test_isa_schedule.py b/test/unit/test_isa_schedule.py new file mode 100644 index 0000000000..570e0e20e9 --- /dev/null +++ b/test/unit/test_isa_schedule.py @@ -0,0 +1,47 @@ +import unittest +from tinygrad.uop.ops import UOp, Ops, dtypes, graph_rewrite +from tinygrad.renderer.isa import IselContext +from tinygrad.renderer.x86 import X86Renderer + +class TestX86Schedule(unittest.TestCase): + def schedule(self, x:UOp) -> list[UOp]: + x = graph_rewrite(x, X86Renderer().pre_isel_matcher) + x = graph_rewrite(x, X86Renderer().isel_matcher, IselContext(x), bottom_up=True) + + def test_hide_latency(self): + buf = UOp(Ops.DEFINE_GLOBAL, dtypes.float32.ptr(), arg=0) + load1 = buf.index(UOp.const(dtypes.int32, 1), ptr=True).load() + load2 = buf.index(UOp.const(dtypes.int32, 1), ptr=True).load() + const = UOp.const(dtypes.float32, 1) + # short path, cheap alu + add = load1 + const + # long path, expensive alu + fmadd = UOp.alu(Ops.MULACC, load2, const, const) + # unify the paths + n = self.schedule(add + fmadd) + # load2 should be picked first as it has a longer path + + # in-order core can't issue ops with dependencies between them in a single cycle + def test_issue_io(self): pass + + # out-of-order core can issue ops with dependencies between them in a single cycle + def test_issue_ooo(self): pass + + # if micro ops > issue width can issue this cycle if no other micro ops were issued + def test_issue_width_empty_cycle(self): pass + + # if micro ops were issued this cycle and issue width can't fit micro ops then they can't be issued this cycle + def test_issue_width_non_empty_cycle(self): pass + + # test cycles advance and no op is issued until stall clears + def test_stall(self): pass + + # test reg pressure + def test_reg_pressure(self): pass + + # test you can issue x whose unit was reserved for y but x's unit end cycle <= y's unit start cycle + def test_resource_cycles_no_intersection(self): pass + + # now test x's unit end cycle > y's unit start cycle, can still issue x if ooo + def test_resource_cycles_intersection(self): pass + diff --git a/tinygrad/codegen/late/schedule.py b/tinygrad/codegen/late/schedule.py new file mode 100644 index 0000000000..e294deccab --- /dev/null +++ b/tinygrad/codegen/late/schedule.py @@ -0,0 +1,160 @@ +from tinygrad.uop.ops import UOp, AllOps +from tinygrad.renderer.isa import Register +from dataclasses import dataclass +import math + +# this is an execution unit +@dataclass +class Unit: pass + +# this is a group of execution units, an op can execute in any of the units +@dataclass +class Resource: + units: tuple[Unit, ...] + # size of the reservation station, micro-ops go here if their operands aren't ready or there isn't space in the resource + # -1 is for unified reservation station + # 0 is for in-order core + # 1 is for in-order units in out-of-order core + buffer_size: int = -1 + +# op scheduling info +@dataclass +class OpInfo: + latency: int # minimum delay added to the dependency chain + # resources used, includes the cycle when the unit is reserved and the cycle when the unit is released. one unit is reserved per resource + resources: tuple[tuple[Resource, int, int], ...] + micro_ops: int = 1 # number of micro-ops issued + +# info about the whole processor +@dataclass +class MachineInfo: + issue_width: int # number of micro-ops that can be issued per cycle + mop_buffer_size: int # number of micro-ops that can be buffered (this is the minimum between the size of the reorder buffer, + # entries in register file and size of the unified reservation station), for an in-order core this number is 0 + +class MachineScheduler: + def __init__(self, sink:UOp, mach_info: MachineInfo, op_info: dict[AllOps, OpInfo]): + self.op_info, self.mach_info = op_info, mach_info + self.consumers = sink.get_consumer_map() + # path from all dependencies of x to x (exclusive) with longest latency + self.depth: dict[UOp, int] = {} + for x in self.consumers: self.depth[x] = max([self.depth[s] + op_info[s.op].latency for s in x.src], default=0) + # path from all dependents of x to x (exclusive) with longest latency + self.height: dict[UOp, int] = {} + for x,y in reversed(self.consumers.items()): self.height[x] = max([self.height[c] + op_info[c.op].latency for c in y], default=0) + # map from resource to total count + self.res_count = {res:0 for info in op_info.values() for res,_,_ in info.resources} + # map from unit to next cycle when it's free, used for hazard check + self.unit_ready = {unit:0 for res in self.res_count for unit in res.units} + + self.latency_factor = math.lcm(mach_info.issue_width, *[len(res.units) for res in self.res_count]) + + self.mop_factor = self.latency_factor // mach_info.issue_width + # map from scheduled uop to cycle it was scheduled at, init with uops that aren't instructions + self.sched = {x:0 for x in self.consumers if not x.src} + # map from uop whose dependencies have all been scheduled to cycle in which all its operands are ready, used for hazard check + self.pending = {x:0 for x in self.sched if set(x.src).issubset(self.sched)} + # map from register set to amount of live regs in that set + self.reg_set: dict[tuple[Register, ...], int] = {} + # the current cycle in the timeline + self.cycle: int = 0 + # micro-ops issued in the current cycle + self.cycle_mops: int = 0 + # total micro-ops issued + self.total_mops: int = 0 + # total amount of latency scheduled, longest path so far + self.expected_latency: int = 0 + # the critical resource, oversubscribed + self.crit_res: Resource|None = None + + # total scheduled latency, stalls can cause cycle > expected, out-of-order can cause cycle < expected + @property + def sched_latency(self): return max(self.expected_latency, self.cycle) + @property + def crit_count(self): return self.total_mops * self.mop_factor if self.crit_res is None else self.res_count[self.crit_res] + # avoid x if it increases register pressure above limit, favor x if it reduces pressure above limit + def check_reg_pressure(self, x:UOp) -> int: + new_reg_set = self.reg_set.copy() + # if s was defined in the same block as x and x is its last use then s register is free + for s in x.src: + if isinstance(s.arg, Register) and set(self.consumers[s]) - set(self.sched) == {x} and s.ranges == x.ranges: new_reg_set[s.arg.cons] -= 1 + if isinstance(x.arg, Register): new_reg_set[x.arg.cons] += 1 + # difference in pressure above limit, any reduction or increase below limit is ignored + return sum(max(new_reg_set[r], len(r)) - max(self.reg_set[r], len(r)) for r in new_reg_set) + # avoid x if it uses an oversubscribed resource TODO: why does llvm accumulate this? + def check_res_pressure(self, x:UOp) -> int: return next((end for res,_,end in self.op_info[x.op].resources if res is self.crit_res), 0) + # avoid x if it's in the critical path and a predecessor was issued recently, only relevant for out-of-order as otherwise x isn't ready + def check_lower_bound_latency(self, x:UOp) -> int: return max(self.depth[x] - self.sched_latency, 0) + # favor x according to its remaining latency chain + def check_height(self, x:UOp) -> int: return -self.height[x] + + def pick(self) -> UOp|None: + # check whether this op can be issued this cycle + def _is_ready(x:UOp) -> bool: + # check issue width can fit new micro ops unless nothing has been issued this cycle + # in that case an expensive op with micro ops > issue width can be issued, but in multiple cycles + if self.cycle_mops > 0 and self.cycle_mops + self.op_info[x.op].micro_ops > self.mach_info.issue_width: return False + # these checks are skipped for out-of-order cores as then x can still be dispatched this cycle regardless of hazards + if self.mach_info.mop_buffer_size == 0: + # data hazard (operands not ready) check + if self.pending[x] < self.cycle: return False + # structural hazard (resources not available) check + if any(self.cycle < min(self.unit_ready[u] for u in res.units) for res,_,_ in self.op_info[x.op].resources): return False + return True + # pick the best according to heuristics + return min([x for x in self.pending if _is_ready(x)], key=lambda k: (self.check_reg_pressure(k), self.check_res_pressure(k), + self.check_lower_bound_latency(k), self.check_height(k)), default=None) + + def bump_cycle(self, next_cycle:int): + dec_mops = self.mach_info.issue_width * (next_cycle - self.cycle) + self.cycle_mops = 0 if self.cycle_mops <= dec_mops else self.cycle_mops - dec_mops + self.cycle = next_cycle + + def update(self, x:UOp|None): + next_cycle = self.cycle + if x is not None: + # add x and the current cycle to the schedule + # TODO: this prob shouldnt be a max + self.sched[x] = max(self.pending.pop(x), self.cycle) + # add consumers whose dependencies have all been scheduled to pending, and the first cycle when all its operands are ready + for v in self.consumers[x]: + if set(v.src).issubset(self.sched): self.pending[v] = max(self.sched[s] + self.op_info[s.op].latency for s in v.src) + + if self.mach_info.mop_buffer_size == 0: assert self.pending[x] <= next_cycle + # when is mop_buffer_size == 1? + elif self.mach_info.mop_buffer_size == 1: next_cycle = max(next_cycle, self.pending[x]) + # if this is an in-order resource in out-of-order core account for likely stall cycles + elif any(res.buffer_size == 1 for res,_,_ in self.op_info[x.op].resources): next_cycle = max(next_cycle, self.pending[x]) + + self.total_mops += self.op_info[x.op].micro_ops + # if this threshold is hit the resource is less critical than mop issue + if self.crit_res is not None and self.total_mops * self.mop_factor - self.res_count[self.crit_res] >= self.latency_factor: self.crit_res = None + # update resources + for res,start,end in self.op_info[x.op].resources: + self.res_count[res] += self.latency_factor // len(res.units) * (end - start) + if self.res_count[res] > self.crit_count: self.crit_res = res + + # update the cycle when unit in resource is released by x, only relevant for in-order + if self.mach_info.mop_buffer_size == 0: + #next_cycle = max(next_cycle, min(self.unit_ready[u] for res,_,_ in self.op_info[x.op].resources for u in res.units)) + for res,_,end in self.op_info[x.op].resources: + unit = min([u for u in res.units], key=lambda k: self.unit_ready[k]) + # TODO: when is unit_ready ever greater for in-order? + self.unit_ready[unit] = max(self.unit_ready[unit], next_cycle + end) + + self.expected_latency = max(self.expected_latency, self.depth[x]) + # if a stall occured, bump until stall clears + if next_cycle > self.cycle: self.bump_cycle(next_cycle) + + self.cycle_mops += self.op_info[x.op].micro_ops + while self.cycle_mops >= self.mach_info.issue_width: + next_cycle += 1 + self.bump_cycle(next_cycle) + + # if this threshold is hit the resource isn't deemed critical anymore + if self.crit_res is not None and not (self.crit_count - (self.latency_factor * self.sched_latency) >= self.latency_factor): self.crit_res = None + + def schedule(self) -> list[UOp]: + # TODO: check acyclic latency for ooo + while self.pending: self.update(self.pick()) + return list(self.sched) diff --git a/tinygrad/renderer/isa.py b/tinygrad/renderer/isa.py index dfee220942..f8a544faf1 100644 --- a/tinygrad/renderer/isa.py +++ b/tinygrad/renderer/isa.py @@ -1,11 +1,12 @@ from __future__ import annotations from tinygrad.renderer import Renderer from dataclasses import dataclass, field -from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, UPat, Ops +from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, UPat, Ops, AllOps from tinygrad.codegen import line_rewrite from tinygrad.codegen.late.linearizer import linearize +from tinygrad.codegen.late.schedule import MachineScheduler, MachineInfo, OpInfo from tinygrad.uop.spec import type_verify -from tinygrad.helpers import SPEC, DEBUG +from tinygrad.helpers import SPEC, DEBUG, getenv import itertools def print_uop_asm(uops:list[UOp]): @@ -47,6 +48,8 @@ class ISARenderer(Renderer): pre_isel_matcher: PatternMatcher isel_matcher: PatternMatcher post_regalloc_matcher: PatternMatcher + mach_info: MachineInfo + op_info: dict[AllOps, OpInfo] def two_address(self, x:UOp) -> int|None: raise NotImplementedError("arch specific") def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific") @@ -58,7 +61,8 @@ class ISARenderer(Renderer): sink = graph_rewrite(sink, self.isel_matcher, ctx=isel_ctx, name="instruction selection", bottom_up=True) # TODO: remove, annoying needed for noops sink = graph_rewrite(sink, isel_fixup, name="instruction selection fixup") - lst = linearize(sink) + if getenv("MACHINE_SCHEDULER"): lst = MachineScheduler(sink, self.mach_info, self.op_info).schedule() + else: lst = linearize(sink) if DEBUG >= 8: print_uop_asm(lst) regalloc_ctx = RegallocContext(lst, self, isel_ctx.stack_size) lst = line_rewrite(lst, pm_regalloc, regalloc_ctx) diff --git a/tinygrad/renderer/x86.py b/tinygrad/renderer/x86.py index 3d0d71af4c..6c094a7277 100644 --- a/tinygrad/renderer/x86.py +++ b/tinygrad/renderer/x86.py @@ -5,8 +5,28 @@ from tinygrad.uop import Ops, X86Ops, GroupOp, X86GroupOp from tinygrad.uop.ops import UOp, UPat, PatternMatcher from tinygrad.renderer.isa import Register, ISARenderer, IselContext from tinygrad.codegen.late.regalloc import assign +from tinygrad.codegen.late.schedule import OpInfo, Resource, Unit from tinygrad.helpers import getenv, CPU_COUNT +# ***** X86 scheduling info, specific to a processor generation ***** +# zen 4, this is the default scheduling model +zen4_agu0, zen4_agu1, zen4_agu2 = Unit(), Unit(), Unit() +zen4_lsu0, zen4_lsu1, zen4_lsu2 = Unit(), Unit(), Unit() +zen4_flp0, zen4_flp1, zen4_flp2 = Unit(), Unit(), Unit() +zen4_flp3, zen4_flp4, zen4_flp5 = Unit(), Unit(), Unit() +zen4_agus = Resource((zen4_agu0, zen4_agu1, zen4_agu2)) +zen4_load = Resource((zen4_lsu0, zen4_lsu1, zen4_lsu2)) +zen4_store = Resource((zen4_lsu0, zen4_lsu1)) +zen4_add = Resource((zen4_flp2, zen4_flp3)) +load_lat = 4 # assumes an l1 cache +# TODO: spends 3 cycles in agu if dtype <= 16 +zen4_op_info = { +X86Ops.MOV: OpInfo(load_lat+1, ((zen4_agus, 0, 1), (zen4_load, 1, 2))), +X86Ops.MOVm: OpInfo(1, ((zen4_agus, 0, 1), (zen4_store, 1, 3))), +**{x: OpInfo(3, ((zen4_add, 0, 1),)) for x in (X86Ops.VADDSS, X86Ops.VADDPS, X86Ops.VSUBSS, X86Ops.VSUBPS)}, +**{x: OpInfo(3, ((zen4_add, 0, 1),)) for x in (X86Ops.VADDSD, X86Ops.VADDPD, X86Ops.VSUBSD, X86Ops.VSUBPD)}, +} + # ***** X86 legalization ***** extra_matcher = PatternMatcher([