from __future__ import annotations import itertools from typing import cast from tinygrad.uop.ops import Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, identity_element from tinygrad.uop.ops import axis_colors, axis_to_pos from tinygrad.device import Buffer from tinygrad.dtype import dtypes from tinygrad.helpers import colored, getenv, DEBUG, NOOPT, round_up, prod, merge_dicts, get_single_element, flatten from tinygrad.helpers import ALLOW_TF32, count, Context from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check from tinygrad.codegen.simplify import pm_flatten_range from tinygrad.renderer import Renderer, TensorCore split_targets = {AxisType.UPCAST: (AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK), AxisType.UNROLL: (AxisType.REDUCE, AxisType.GROUP_REDUCE), AxisType.LOCAL: (AxisType.GLOBAL, AxisType.WEAK), AxisType.GROUP_REDUCE: (AxisType.REDUCE,)} class Scheduler: def __init__(self, ast:UOp, ren:Renderer): self.ast, self.ren = ast, ren self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else [] self.tensor_core:TensorCore|None = None self.opt_range = count(start=max([x.arg[0] for x in self.rngs], default=0)+1) @property def rngs(self): # always in order by axistype. void RANGEs are loops, not opt axes. the DEVICE axis is launched, not an opt axis 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 and u.arg[-1] is not AxisType.DEVICE], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1]) @property def shape_len(self) -> int: return len(self.rngs) @property def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs] @property def axis_types(self) -> list[AxisType]: return [x.arg[-1] for x in self.rngs] def copy(self) -> Scheduler: ret = Scheduler(self.ast, self.ren) ret.applied_opts = self.applied_opts[:] ret.tensor_core = self.tensor_core return ret def get_optimized_ast(self, name_override:str|None=None) -> UOp: if name_override is not None: name = name_override else: k_type = "r" if self.reduceop is not None else "E" special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg) special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops] name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())]) self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range") return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts)), tag=1) def _output_rngs(self) -> list[UOp]: return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END]) def _globalizable_rngs(self) -> list[UOp]: ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.WEAK] # exclude any output ranges from global that don't appear in all BUFFERIZE for x in self.ast.toposort(): if x.op is Ops.STAGE: ret = [r for r in ret if r in x.ranges] return ret def convert_loop_to_global(self) -> None: if not self.ren.has_local: return globalizible_rngs = self._globalizable_rngs() rng = [x.replace(arg=x.arg[0:-1]+(AxisType.GLOBAL,)) if x in globalizible_rngs else x for x in self.rngs] self.ast = self.ast.substitute(dict(zip(self.rngs, rng))) def colors(self) -> list[str]: output_rngs = self._output_rngs() globalizible_rngs = self._globalizable_rngs() ret = [] for x,r in zip(self.axis_types, self.rngs): if r not in output_rngs and x == AxisType.WEAK: ret.append("BLACK") elif r not in globalizible_rngs and x == AxisType.WEAK: ret.append("white") else: ret.append(axis_colors[x]) return ret def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())]) def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng:UOp|None=None): check(rng.arg[-1] in split_targets[new_type], f"{new_type} is from {split_targets[new_type]}, not {rng.arg[-1]}") if (old_sz:=rng.src[0].divides(amount)) is None: raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}") new_rng = UOp.range(amount, next(self.opt_range), new_type, dtype=rng.dtype) if input_new_rng is None else input_new_rng replaced_rng = rng.replace(src=(old_sz,)) sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng) self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}") return replaced_rng, new_rng def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type] def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type] @property def reduce_axes(self) -> list[int]: red = {r for u in self.ast.backward_slice if u.op is Ops.REDUCE for s in u.src[1:] for r in s.ranges} return [i for i,r in enumerate(self.rngs) if r in red] def upcast_size(self): return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL)) @property def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK) \ if isinstance(s:=self.full_shape[i], int) and s > 1] @property def unrollable_dims(self) -> list[int]: return [i for i in self.reduce_axes if self.axis_types[i] in (AxisType.GROUP_REDUCE, AxisType.REDUCE) \ and isinstance(s:=self.full_shape[i], int) and s > 1] def apply_opt(self, opt:Opt, append_opt:bool=True): if opt.op is OptOps.TC: rng = UOp(Ops.NOOP) else: check(type(opt.axis) is int and 0 <= opt.axis < self.shape_len, f"invalid axis on {opt.axis=} {opt.op=} {self.shape_len=}") rng = self.rngs[cast(int, opt.axis)] ret = None if opt.op is OptOps.SPLIT: check(isinstance(opt.arg, tuple) and len(opt.arg) in (2, 3), f"split arg is (amt, target) or (amt, target, top), not {opt.arg}") amt, new_type, top = (*cast(tuple, opt.arg), False)[0:3] check(type(amt) is int and (amt == 0 or amt > 1) and isinstance(new_type, AxisType) and new_type in split_targets and isinstance(top, bool), f"invalid split arg {opt.arg}") check(not top or new_type is AxisType.GROUP_REDUCE, "top is only for group reduce") if new_type in (AxisType.LOCAL, AxisType.GROUP_REDUCE): check(self.ren.has_local, "locals needed for opt") check(rng.arg[-1] in split_targets[new_type], f"{new_type} is from {split_targets[new_type]}, not {rng.arg[-1]}") if amt == 0: amt = int(rng.vmax+1) if new_type is AxisType.UNROLL: check(amt <= 32, "don't unroll more than 32") if new_type is AxisType.UPCAST: check(self.ren.target.device == "DSP" or amt <= 16, "don't upcast more than 16") # prevents METAL compiler hangs if self.reduceop is not None and (new_type is AxisType.GROUP_REDUCE or self.group_for_reduces): upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]) smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize check(smem_sz <= self.ren.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.ren.shared_max}") if self.reduceop is not None and new_type is AxisType.GROUP_REDUCE: # We currently dont support a group within another rudece, TODO: fix if-contexts reduce = [u for u in self.ast.backward_slice if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0] check(not any(u.arg[-1] in (AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE) for u in reduce.ranges), "cannot have a GROUP_REDUCE inside another reduce") ret = self.shift_to(rng, amt, new_type, top=top) elif opt.op is OptOps.TC: check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps check(opt.axis is not None and opt.axis >= 0, "tensor core opts must have an axis") check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg") check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.ren.tensor_cores), "tensor core opts must have valid tc_select") check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt") check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid") try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt) except ValueError as e: raise KernelOptError(str(e)) check(ret is not None, "no tensor core available") elif opt.op is OptOps.PADTO: check(type(opt.arg) is int and opt.arg > 1, f"padto arg is a multiple > 1, not {opt.arg}") check(rng.src[0].op is Ops.CONST, "only pad const axes") # TODO: upcasted is only wrong for a range pinned in WMMA tc_upcast_axes check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL, AxisType.WARP}, "cannot pad upcasted or warp") new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg)) check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work") replaced_rng = UOp.range(new_sz, *rng.arg, dtype=rng.dtype) replaces = {rng:replaced_rng} valid = replaced_rng < rng.vmax+1 for b in self.bufs: if rng in (i:=b.src[1]).ranges: replaces[b] = b.replace(src=(b.src[0], i.get_idx().valid(valid&i.get_valid()))) for r in self.reduceops: if any(rng in y.ranges for y in r.src[1:]): replaces[r] = r.replace(src=(valid.where(r.src[0], UOp.const(identity_element(r.arg[0], r.dtype), r.dtype)),)+r.src[1:]) self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}") ret = replaced_rng elif opt.op is OptOps.SWAP: check(type(opt.arg) is int and 0 <= opt.arg < self.shape_len, f"invalid swap axis on {opt.arg=} {self.shape_len=}") altrng:UOp = self.rngs[cast(int, opt.arg)] check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals") self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1])), altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]))}, name=f"swap {rng.arg[:-1]} {altrng.arg[:-1]}", walk=True) else: raise KernelOptError(f"unsupported opt {opt.op}") if append_opt: self.applied_opts.append(opt) return ret def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]: if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore") reduceop = reduceops[0] if reduceop.arg[0] is Ops.ADD: mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0] if mul.op is not Ops.MUL: return None in0, in1 = mul.src for tc in self.ren.tensor_cores if tc_select == -1 else [self.ren.tensor_cores[tc_select]]: if self.ren.target.device in ("CUDA", "NV") and tc.dtype_in == dtypes.float and not ALLOW_TF32: continue if tc.dtype_in == in0.dtype and tc.dtype_in == in1.dtype and tc.dtype_out == reduceop.dtype: # tensor cores have three ranges. X, Y, and REDUCE in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True) in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0], reverse=True) red_ranges = sorted(UOp.sink(*reduceop.src[1:]).ranges, key=lambda x: x.arg[0], reverse=True) if DEBUG >= 3: print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}", f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}") if not len(in0_ranges) or not len(in1_ranges) or not len(red_ranges): continue # pick ranges # NOTE: in1 and in0 are switched because tc.dims is (N, M, K) axis_choices = list(itertools.product(in1_ranges, in0_ranges, red_ranges)) if not (axis < len(axis_choices)): continue axes = list(axis_choices[axis]) if any(a.arg[-1] is AxisType.REDUCE for a in axes[:2]): raise KernelOptError("tensor core X/Y axes can't be REDUCE") # do optimizations and save the ranges ast, warp, ne = self.ast, UOp.range(tc.threads, -1, AxisType.WARP), {} try: for i,a in enumerate(axes): if (a.vmax+1) % tc.dims[i] != 0: if opt_level < 2: raise KernelOptError("tc padding requires opt_level >= 2") axes[i] = self.apply_opt(Opt(OptOps.PADTO, self.rngs.index(a), tc.dims[i]), append_opt=False) # PADTO might fail # we create the warp as a whole thing, in case some of these ranges are moved/removed later for c in tc.axis_coords(): d = "nmk".index(c[0]) if c in tc.frag_c[0]: axes[d], ne[c] = self.shift_to(axes[d], 2, AxisType.LOCAL, input_new_rng=warp//2**tc.frag_c[0].index(c)%2) else: axes[d], ne[c] = self.shift_to(axes[d], 2, AxisType.UNROLL if d == 2 else AxisType.UPCAST) except KernelOptError: self.ast = ast continue if use_tensor_cores != 2: reduceop = get_single_element([x for x in self.reduceops if axes[2] in UOp.sink(*x.src[1:]).ranges]) gate, mul = (r0.src[0], r0.src[1]) if (r0:=reduceop.src[0]).op is Ops.WHERE else (None, r0) if mul.op is Ops.CAST: mul = mul.src[0] ins = mul.src if gate is None else tuple(gate.where(x, UOp.const(0, x.dtype)) for x in mul.src) srcs = [x.substitute({ne[a]: ne[b] for a,b in rl.items()}, walk=True) for x,rl in zip(ins, tc.relabel())] # get upcast axes for the tensor cores base_upcast_axes = [ne[c].arg[0] for c in tc.base_upcast_axes()] upcast_cnt = [len(f[1]) for f in (tc.frag_a, tc.frag_b, tc.frag_c)] # each operand upcasts its first upcast_cnt axes, the axes only A or B upcast are size 1 so the operands broadcast tc_upcast_axes = tuple([tuple([(a, 2 if j < cnt else 1) for j,a in enumerate(base_upcast_axes[:max(cnt, *upcast_cnt[:2])])]) for cnt in upcast_cnt]) # construct the op # TODO: remove tc_upcast_axes from the arg tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const((0.0,)*2**upcast_cnt[2], tc.dtype_out), tc.dims, tc.threads, tc_upcast_axes=tc_upcast_axes) # preserve extra reduces reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x not in [ne[c] for c in ne if c[0] == "k"]] if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, src=(tc_uop,)+tuple(reduce_ranges), arg=(Ops.ADD, 0)) self.ast = self.ast.substitute({reduceop: tc_uop}) self.tensor_core = tc return axes return None # helpers for hand_coded_optimizations @property def reduceops(self) -> list[UOp]: return [x for x in self.ast.backward_slice if x.op is Ops.REDUCE] @property def reduceop(self) -> UOp|None: return red[0] if (red:=self.reduceops) else None @property def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1] @property def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL)) @property def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE)) def args_from_ast(ast:UOp, dname:str) -> tuple[list[Buffer], dict[str, int]]: glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM and x.arg.slot >= 0], key=lambda x: x.arg.slot) return [Buffer(dname, x.max_numel(), x.dtype) for x in glbls], {k.expr:int(k.vmax+k.vmin)//2 for k in ast.variables()} def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp: if ast.tag is not None: return ast k = Scheduler(ast, ren) k.convert_loop_to_global() if ast.arg is not None and ast.arg.opts_to_apply is not None: for opt in ast.arg.opts_to_apply: k.apply_opt(opt) elif beam >= 1: from tinygrad.codegen.opt.search import beam_search rawbufs, var_vals = args_from_ast(ast, ren.target.device) # beam search may open devices with Context(ALLOW_DEVICE_USAGE=1): k = beam_search(k, rawbufs, var_vals, beam, bool(getenv("BEAM_ESTIMATE", 1))) elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()): from tinygrad.codegen.opt.heuristic import hand_coded_optimizations # NOTE: hand_coded_optimizations doesn't support multiblock opts yet if not any(u.op is Ops.STAGE for u in ast.backward_slice): k = hand_coded_optimizations(k) return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)