start new scheduler

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

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
+

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)

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)

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([