mirror of
https://github.com/dragonpilot/dragonpilot.git
synced 2026-08-21 08:03:42 +08:00
dragonpilot v0.9.8
This commit is contained in:
Comma Device
@@ -0,0 +1,727 @@
|
||||
from __future__ import annotations
|
||||
import itertools, functools
|
||||
from dataclasses import dataclass
|
||||
from collections import defaultdict
|
||||
from typing import Optional, cast, Final, DefaultDict, Callable, Sequence
|
||||
from enum import Enum, auto
|
||||
|
||||
from tinygrad.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, print_uops, type_verify, resolve, Variable, sint, \
|
||||
graph_rewrite, track_rewrites, view_left
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.renderer import Renderer, TensorCore, ProgramSpec
|
||||
from tinygrad.dtype import ImageDType
|
||||
from tinygrad.helpers import all_same, colored, ansilen, dedup, getenv, prod, round_up, all_int, to_function_name, diskcache_put
|
||||
from tinygrad.helpers import DEBUG, TC_OPT, USE_TC, AMX
|
||||
from tinygrad.shape.shapetracker import ShapeTracker
|
||||
from tinygrad.shape.view import strides_for_shape
|
||||
from tinygrad.codegen.linearize import linearize_uop
|
||||
from tinygrad.codegen.uopgraph import full_graph_rewrite
|
||||
from tinygrad.codegen.lowerer import rewrite_shapetracker_with_index, get_contraction
|
||||
|
||||
class OptOps(Enum):
|
||||
TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
|
||||
GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702
|
||||
def __lt__(self, x:OptOps): return self.value < x.value
|
||||
|
||||
class KernelOptError(Exception): pass
|
||||
|
||||
def check(cond:bool, msg:str=""):
|
||||
if not cond: raise KernelOptError(msg)
|
||||
|
||||
@dataclass(frozen=True, order=True)
|
||||
class Opt:
|
||||
op: OptOps
|
||||
axis: Optional[int] = None
|
||||
amt: Optional[int] = None
|
||||
def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, amt={self.amt})"
|
||||
def real_axis(self, k:Kernel):
|
||||
if self.axis is None: return -1
|
||||
if self.op is OptOps.UNROLL: return k.first_reduce+self.axis
|
||||
if self.op in {OptOps.GROUP, OptOps.GROUPTOP}: return k.first_reduce+k.group_for_reduces+self.axis
|
||||
return self.axis
|
||||
|
||||
@dataclass
|
||||
class TensorCoreOptions:
|
||||
axes: tuple[int, ...] # the location of the original N and M axes if still in the shape
|
||||
axes_exist: tuple[bool, ...] # true if the original N and M axes are still in the shape
|
||||
axis_pads: tuple[tuple[int, int], ...]
|
||||
def fix_axes(self, removed_axis:int): # adjust the TC axes if necesssary when a dimension is removed
|
||||
axes, axes_exist = list(self.axes), list(self.axes_exist)
|
||||
for tc_dim in [i for i in range(2) if axes_exist[i]]:
|
||||
if removed_axis < axes[tc_dim]: axes[tc_dim] -= 1
|
||||
elif removed_axis == axes[tc_dim]: axes_exist[tc_dim] = False
|
||||
self.axes, self.axes_exist = tuple(axes), tuple(axes_exist)
|
||||
|
||||
class Kernel:
|
||||
def __init__(self, ast:UOp, opts:Optional[Renderer]=None):
|
||||
if ast.op is Ops.SINK: self.ast = ast
|
||||
|
||||
self.opts = opts if opts is not None else Device[Device.DEFAULT].renderer
|
||||
try: uop_sts_map = verify_ast(self.ast)
|
||||
except AssertionError as e:
|
||||
print("INVALID AST")
|
||||
print(self.ast)
|
||||
raise e
|
||||
|
||||
self.reduceops = [x for x in self.ast.toposort if x.op is Ops.REDUCE_AXIS]
|
||||
|
||||
self.vars: list[Variable] = self.ast.variables()
|
||||
# NOTE: this requires a specific order with the [::-1], this is likely a bug
|
||||
self.bufs: list[UOp] = [x for x in self.ast.toposort if x.op in GroupOp.Buffer][::-1]
|
||||
|
||||
# get earlybufs, before any reduceops
|
||||
earlybufs: list[UOp] = [x for reduceop in self.reduceops for x in reduceop.src[0].toposort if x.op in GroupOp.Buffer]
|
||||
self.full_buf_index: int = self.bufs.index(earlybufs[0]) if earlybufs else 0
|
||||
# NOTE: full_shape can be wrong if there's a tree of reduces
|
||||
|
||||
# create new shapetrackers inside this kernel, we will permute them
|
||||
self.sts: list[ShapeTracker] = [x.st_arg for x in self.bufs]
|
||||
|
||||
# add the shapetrackers for each reduce
|
||||
# we use this to track which axes are reduced in each reduce
|
||||
for x in self.reduceops:
|
||||
self.sts.append(uop_sts_map[x])
|
||||
self.sts.append(uop_sts_map[x.src[0]])
|
||||
|
||||
# move all reduce axes to the end
|
||||
reduce = list(enumerate(zip(self.full_shape, self.output_shape)))
|
||||
permute = tuple([i for i,(s,n) in reduce if not resolve(s != n)] + [i for i,(s,n) in reduce if resolve(s != n)])
|
||||
self.reshape_and_permute(None, permute)
|
||||
|
||||
# parameters for optimization
|
||||
self.applied_opts: list[Opt] = []
|
||||
self.group_for_reduces: int = 0
|
||||
self.upcasted: int = 0
|
||||
self.local_dims: int = 0
|
||||
self.tensor_core: Optional[TensorCore] = None
|
||||
self.tensor_core_opts: Optional[TensorCoreOptions] = None
|
||||
self.use_tensor_cores: int = 0
|
||||
self.dont_use_locals: bool = False
|
||||
|
||||
# group simplifies
|
||||
self.simplify_ones()
|
||||
self.simplify_merge_adjacent()
|
||||
|
||||
def copy(self):
|
||||
ret = type(self).__new__(type(self))
|
||||
|
||||
# base linearizer params
|
||||
ret.opts, ret.ast = self.opts, self.ast
|
||||
|
||||
# things downstream of the AST
|
||||
ret.reduceops, ret.vars, ret.bufs, ret.full_buf_index = self.reduceops, self.vars, self.bufs, self.full_buf_index
|
||||
ret.sts = self.sts[:len(ret.bufs)+len(ret.reduceops)*2] # NOTE: must redo the local buffers with TC in beam
|
||||
|
||||
# parameters for optimizations
|
||||
ret.applied_opts, ret.group_for_reduces, ret.upcasted, ret.local_dims, ret.dont_use_locals = \
|
||||
self.applied_opts[:], self.group_for_reduces, self.upcasted, self.local_dims, self.dont_use_locals
|
||||
ret.tensor_core, ret.tensor_core_opts, ret.use_tensor_cores = self.tensor_core, self.tensor_core_opts, self.use_tensor_cores
|
||||
|
||||
return ret
|
||||
|
||||
@property
|
||||
def membufs(self) -> list[UOp]: return dedup([x.src[0] for x in self.bufs if x.op in {Ops.LOAD, Ops.STORE}])
|
||||
|
||||
# TODO: these need more tests or it might silently be no-op
|
||||
def float4_axis(self, i:int): return [x-self.first_upcast for x in self.sts[i].unit_stride_axes() if x >= self.first_upcast and self.sts[i].shape[x]%4 == 0] # noqa: E501
|
||||
|
||||
def upcasted_axis(self, i:int) -> list[tuple[int, Optional[sint], bool]]:
|
||||
upcasted_shape, upcasted_stride = self.sts[i].shape[self.first_upcast:], self.sts[i].real_strides()[self.first_upcast:]
|
||||
assert all_int(upcasted_shape), f"cannot upcast a symbolic amount {upcasted_shape=}"
|
||||
return list(zip(upcasted_shape, upcasted_stride,
|
||||
[x!=y for x,y in zip(self.sts[0].shape[self.first_upcast:], self.full_shape[self.first_upcast:])]))
|
||||
|
||||
@property
|
||||
def first_reduce(self) -> int:
|
||||
return [resolve(x!=y) for x,y in zip(self.sts[0].shape[:self.first_upcast]+(0,), self.full_shape[:self.first_upcast]+(1,))].index(True)
|
||||
|
||||
@property
|
||||
def first_upcast(self) -> int: return self.shape_len-self.upcasted
|
||||
|
||||
@property
|
||||
def reduceop(self) -> UOp|None: return self.reduceops[0] if len(self.reduceops) > 0 else None
|
||||
|
||||
@property
|
||||
def output_shape(self) -> tuple[sint, ...]: return self.sts[0].shape
|
||||
|
||||
@property
|
||||
def full_shape(self) -> tuple[sint, ...]: return self.sts[self.full_buf_index].shape
|
||||
|
||||
@property
|
||||
def full_unupcasted_shape(self) -> tuple[sint, ...]: return self.full_shape[:self.first_upcast]
|
||||
|
||||
@property
|
||||
def shape_len(self) -> int: return len(self.sts[0].shape)
|
||||
|
||||
@property
|
||||
def global_dims(self) -> int: return self.first_reduce-self.local_dims
|
||||
|
||||
# there's eight chunks of the shape
|
||||
# blue -- global dims
|
||||
# cyan -- local dims (warp ones first)
|
||||
# *** self.first_reduce
|
||||
# green -- reduce-local dims
|
||||
# red -- reduce loops
|
||||
# *** self.upcasted
|
||||
# purple -- reduce upcasted
|
||||
# yellow -- normal upcasted dimensions
|
||||
def colors(self) -> list[str]:
|
||||
# first non local non reduce dims are global (blue)
|
||||
colors = ["blue"] * self.global_dims if not self.dont_use_locals else ["BLUE"] * self.global_dims
|
||||
# after global are local_dims; warp ones used in tensor cores must be closest to first_reduce (cyan)
|
||||
colors += ["cyan"] * self.local_dims
|
||||
# between first_reduce and first_reduce + group_for_reduces, they are late upcasted (green)
|
||||
colors += ["green"] * self.group_for_reduces
|
||||
# between first_reduce + group_for_reduces and upcasted, they are reduce (red)
|
||||
colors += ["red"] * (self.first_upcast - (self.first_reduce + self.group_for_reduces))
|
||||
# upcasted dimensions are reduce (magenta) or normal (yellow)
|
||||
colors += ["magenta" if self.full_shape[i] != self.sts[0].shape[i] else "yellow" for i in range(self.first_upcast, self.shape_len)]
|
||||
assert len(colors) == self.shape_len, "colors size mismatch"
|
||||
return colors
|
||||
|
||||
def colored_shape(self, pad:Optional[int]=None, dense=False) -> str:
|
||||
shape_strs = [(s if dense else f"{s:4d}") if isinstance(s, int) else s.render() for s in self.full_shape]
|
||||
ret = ' '.join(colored(s, color) for s,color in zip(shape_strs, self.colors()))
|
||||
if pad: ret += ' '*(pad-ansilen(ret))
|
||||
return ret
|
||||
|
||||
# ******************** base simplifiers ********************
|
||||
|
||||
# apply reshape and permute to all shapetrackers
|
||||
def reshape_and_permute(self, new_shape_fxn:Optional[Callable[[tuple[sint, ...]], Sequence[sint]]], axis:Optional[Sequence[int]]):
|
||||
def reshape(st:ShapeTracker): return st.reshape(tuple(new_shape_fxn(st.shape))) if new_shape_fxn is not None else st
|
||||
def permute(st:ShapeTracker): return st.permute(tuple(axis)) if axis is not None else st
|
||||
self.sts = [permute(reshape(st)) for st in self.sts]
|
||||
|
||||
# drops the final dimension
|
||||
def upcast(self):
|
||||
check(self.full_shape[-1] != 1, "can't upcast a dimension with size 1")
|
||||
self.upcasted += 1
|
||||
|
||||
# axis : the axis to pull from
|
||||
# amount : the amount to take
|
||||
# top : if you want to pull that amount from the top
|
||||
# insert_before : place to insert the new stuff
|
||||
def shift_to(self, axis, amount, top=False, insert_before=None):
|
||||
if insert_before is None: insert_before = self.shape_len
|
||||
move_axis = axis if top else axis+1
|
||||
if move_axis < insert_before: insert_before += 1
|
||||
self.reshape_and_permute(
|
||||
lambda x: x[0:axis] + (((amount, x[axis]//amount) if top else (x[axis]//amount, amount)) if x[axis] > 1 else (1,1)) + x[axis+1:],
|
||||
[i for i in range(insert_before) if i != move_axis] + [move_axis] + [i for i in range(insert_before, self.shape_len+1) if i != move_axis])
|
||||
|
||||
# ******************** complex simplifiers ********************
|
||||
|
||||
def simplify_ones(self) -> bool:
|
||||
# remove places where the shape is all ones
|
||||
# TODO: this should be factored in to multi shape stride
|
||||
if self.shape_len == 0: return False
|
||||
all_ones = [s==1 for s in self.full_shape]
|
||||
self.local_dims -= sum(all_ones[self.first_reduce-self.local_dims:self.first_reduce])
|
||||
self.upcasted -= sum(all_ones[self.first_upcast:]) # TODO: no necessary since upcasted axis can't be un-upcasted
|
||||
self.reshape_and_permute(lambda shape: [x for i,x in enumerate(shape) if not all_ones[i]], None)
|
||||
return any(all_ones)
|
||||
|
||||
def simplify_merge_adjacent(self):
|
||||
if self.shape_len == 0: return
|
||||
shapes, strides = [x.shape for x in self.sts], [x.real_strides() for x in self.sts]
|
||||
|
||||
# if it's an image, insert fake strides such that this fusion doesn't happen across image axes
|
||||
if isinstance(self.membufs[0].dtype, ImageDType):
|
||||
base_shape = self.membufs[0].dtype.shape
|
||||
if shape_idx_groups := get_contraction(self.output_shape, base_shape):
|
||||
special_strides: tuple[sint, ...] = tuple()
|
||||
for i,g in enumerate(shape_idx_groups):
|
||||
shape_piece = tuple(self.output_shape[x] for x in g)
|
||||
assert prod(shape_piece) == base_shape[i], f"get_contraction was wrong? {shape_piece} != {base_shape[i]}"
|
||||
special_strides += strides_for_shape(shape_piece)
|
||||
# adding the fake image shape
|
||||
shapes.append(self.output_shape)
|
||||
strides.append(special_strides)
|
||||
|
||||
# merge dimensions if we can, multi _merge_dims
|
||||
# NOTE: this does not always preserve the reduce dimension
|
||||
# TODO: move this into shapetracker, with tests!
|
||||
# TODO: how does this work with multi-reduce?
|
||||
rets = [[(s[0], st[0])] for s,st in zip(shapes, strides)]
|
||||
for i in range(1, len(shapes[0])):
|
||||
can_merge = []
|
||||
for s,st,ret in zip(shapes, strides, rets):
|
||||
# TODO: added the always mergeability of 1s, is this right? if so, add to shapetracker in the 1 case
|
||||
si, sti, last_st = s[i], st[i], ret[-1][1]
|
||||
can_merge.append((sti is not None) and ((sti != 0 and last_st == si*sti) or (sti == 0 and last_st == 0)))
|
||||
# more can merge than this
|
||||
mergeable = all(can_merge) and i != self.first_reduce
|
||||
for j,(s,st) in enumerate(zip(shapes, strides)):
|
||||
if mergeable: rets[j][-1] = (rets[j][-1][0] * s[i], st[i])
|
||||
else: rets[j].append((s[i], st[i]))
|
||||
|
||||
# do the reshapes
|
||||
for i,x in enumerate(rets[:len(self.sts)]): self.sts[i] = self.sts[i].reshape(tuple([y[0] for y in x]))
|
||||
|
||||
# ******************** high level optimizers ********************
|
||||
|
||||
def _create_tc_opts(self, reduceop:UOp, tc:TensorCore, axis:int, opt_level:int) -> Optional[TensorCoreOptions]:
|
||||
has_cast = tc.dtype_in != tc.dtype_out
|
||||
if has_cast and not (reduceop.src[0].op is Ops.CAST and reduceop.src[0].dtype == tc.dtype_out): return None
|
||||
|
||||
mul_op = reduceop.src[0].src[0] if has_cast else reduceop.src[0]
|
||||
if mul_op.op is not Ops.MUL: return None
|
||||
|
||||
def buf_index(src:UOp) -> Optional[int]:
|
||||
# TODO: apply tc even if the sources are not from LOAD
|
||||
if src.op is Ops.LOAD and src.dtype == tc.dtype_in: return self.bufs.index(src)
|
||||
try:
|
||||
if opt_level >= 1 and src.op is Ops.CAST and src.dtype == tc.dtype_in: return self.bufs.index(src.src[0])
|
||||
except ValueError: return None
|
||||
return None
|
||||
if (buf0:=buf_index(mul_op.src[0])) is None or (buf1:=buf_index(mul_op.src[1])) is None: return None
|
||||
|
||||
buf0_strides, buf1_strides = self.sts[buf0].real_strides(), self.sts[buf1].real_strides()
|
||||
axis_buf0 = [(i,self.full_shape[i],buf1_strides[i]) for i,s in enumerate(buf0_strides[:self.first_reduce]) if s == 0]
|
||||
axis_buf1 = [(i,self.full_shape[i],buf0_strides[i]) for i,s in enumerate(buf1_strides[:self.first_reduce]) if s == 0]
|
||||
if not (axis_buf0 and axis_buf1 and ((self.shape_len-self.first_reduce) == 1 or (opt_level >= 1))): return None
|
||||
|
||||
axis_choices = list(itertools.product(axis_buf0, axis_buf1, range(self.first_reduce, self.shape_len)))
|
||||
if not (axis < len(axis_choices)): return None
|
||||
|
||||
s0, s1, s2 = axis_choices[-(axis+1)][0][0], axis_choices[-(axis+1)][1][0], axis_choices[-(axis+1)][2] # s0 is n, s1 is m, s2 is k
|
||||
axis_pads = tuple((x, tc.dims[i]) for i, x in enumerate([s0, s1, s2]) if resolve(self.full_shape[x]%tc.dims[i] != 0))
|
||||
if axis_pads and (opt_level < 2): return None
|
||||
if DEBUG >= 3: print("TENSOR CORES", axis_buf0, axis_buf1, tc)
|
||||
return TensorCoreOptions(axes=(s0, s1, s2), axes_exist=(True, True), axis_pads=axis_pads)
|
||||
|
||||
def _apply_tc_opt(self, use_tensor_cores:int, axis:int, opt_level:int) -> bool:
|
||||
if use_tensor_cores and self.reduceop is not None and self.reduceop.arg[0] is Ops.ADD:
|
||||
for tc in self.opts.tensor_cores:
|
||||
tensor_core_opts = [self._create_tc_opts(reduceop, tc, axis, opt_level) for reduceop in self.reduceops]
|
||||
# can only fuse reduces with the same tc options
|
||||
assert all_same(tensor_core_opts)
|
||||
if tensor_core_opts[0] is None: continue
|
||||
# tensor core -- unroll the reduce dim, upcast input and local the correct thread pattern
|
||||
self.tensor_core_opts = tc_opts = tensor_core_opts[0]
|
||||
|
||||
# attempt to pad the tensor axes that require it
|
||||
try:
|
||||
for axis, dim in tc_opts.axis_pads: self.apply_opt(Opt(OptOps.PADTO, axis, dim), append_opt=False) # PADTO might fail
|
||||
except KernelOptError: continue
|
||||
for tc_dim, amt in tc.reduce_axes: self.apply_opt(Opt(OptOps.UNROLL,tc_opts.axes[2]-self.first_reduce,amt), append_opt=False)
|
||||
for opt in tc.opts_seq:
|
||||
if opt == "UP":
|
||||
for tc_dim, amt in tc.early_upcast_axes: self.apply_opt(Opt(OptOps.UPCAST,tc_opts.axes[tc_dim],amt), append_opt=False)
|
||||
elif opt == "LC":
|
||||
for tc_dim, amt in tc.threads: self.apply_opt(Opt(OptOps.LOCAL,tc_opts.axes[tc_dim],amt), append_opt=False)
|
||||
self.tensor_core = tc
|
||||
self.use_tensor_cores = use_tensor_cores # TC=2 will do the shape ops without the WMMA
|
||||
return True
|
||||
return False
|
||||
|
||||
def apply_tensor_cores(self, use_tensor_cores=1, extra_opts:Optional[list[Opt]]=None, axis:int=0, tc_opt:Optional[int]=None) -> bool:
|
||||
""" Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false.
|
||||
Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N).
|
||||
|
||||
Keyword arguments:
|
||||
use_tensor_cores -- controls how tensor cores are applied (default 1)
|
||||
0: will disable any tensor core matching
|
||||
1: enable tensor cores
|
||||
2: apply tensor core shape but don't use UOp.WMMA
|
||||
extra_opts -- additional Opt's to apply after the tensor core instead of the hand-coded additional Opt's (default None)
|
||||
tc_opt -- controls which kinds of kernels may be eligible for tensor cores application (default 2 during BEAM, 0 otherwise)
|
||||
0: applies to only kernels with a single reduce axis and direct UOps.LOAD into Ops.MUL
|
||||
1: allows kernels with multiple reduce axes and also multiplication of UOps.CAST'd buffers
|
||||
2: allows kernels with M, N, K axes that are not multiples of the tensor core dimensions by applying padding those axes as needed
|
||||
"""
|
||||
if tc_opt is None: tc_opt = TC_OPT.value
|
||||
if not self.opts.tensor_cores and use_tensor_cores != 2: return False
|
||||
try: # check TC first and apply hand-coded opts if successful
|
||||
self.apply_opt(Opt(OptOps.TC, axis, tc_opt))
|
||||
|
||||
if (tc_opts:=self.tensor_core_opts) is not None:
|
||||
if extra_opts is not None:
|
||||
for opt in extra_opts: self.apply_opt(opt)
|
||||
else:
|
||||
if (self.opts.device == "CLANG" and AMX): return True # skip hand-coded TC opts if AMX, upcasting will make kernel slower
|
||||
# hand-coded TC opts
|
||||
for tc_dim in [tc_dim for tc_dim in [1,0] if tc_opts.axes_exist[tc_dim]]: # attempt to upcast M and N
|
||||
szs = [sz for sz in [5,4,3,2] if self.full_shape[tc_opts.axes[tc_dim]] % sz == 0]
|
||||
if szs: self.apply_opt(Opt(OptOps.UPCAST, tc_opts.axes[tc_dim], szs[0]))
|
||||
|
||||
if tc_opts.axes_exist[0] and (szs := [sz for sz in [4,2] if self.full_shape[tc_opts.axes[0]] % sz == 0]): # attempt to local N
|
||||
self.apply_opt(Opt(OptOps.LOCAL, tc_opts.axes[0], szs[0]))
|
||||
return True
|
||||
except KernelOptError:
|
||||
return False
|
||||
|
||||
def apply_opt(self, opt:Opt, append_opt:bool=True):
|
||||
if self.dont_use_locals: check(opt.op not in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}, "not using locals")
|
||||
|
||||
if opt.op is OptOps.TC:
|
||||
check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: things like PADTO might be fine
|
||||
check(opt.axis is not None and opt.amt is not None, "tensor core opts must have an axis and amt")
|
||||
check((use_tensor_cores:=USE_TC.value) == 2 or len(self.opts.tensor_cores) > 0, "must have tensor cores or TC=2")
|
||||
check(self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), cast(int, opt.amt)), "no tensor core available")
|
||||
self.applied_opts.append(opt)
|
||||
return
|
||||
|
||||
axis = opt.real_axis(self)
|
||||
check(axis < len(self.full_shape), "invalid axis")
|
||||
|
||||
if opt.op is OptOps.SWAP: amt = cast(int, opt.amt) # amt is an axis in the SWAPs
|
||||
elif opt.amt is not None:
|
||||
amt = opt.amt if opt.amt != 0 else self.full_shape[axis]
|
||||
check(isinstance(amt, int) and amt != 1, "shift/padto of amt 1 or Node is meaningless")
|
||||
if opt.op is not OptOps.PADTO: check(self.full_shape[axis] % amt == 0, "no longer valid shift")
|
||||
else: amt = -1
|
||||
|
||||
if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP} or \
|
||||
(self.group_for_reduces and opt.op not in {OptOps.NOLOCALS, OptOps.PADTO})):
|
||||
acc_sz = self.reduceop.dtype.itemsize
|
||||
upcast_sz = prod([a for a,b in zip(self.full_shape[self.first_upcast:], self.sts[0].shape[self.first_upcast:]) if a == b])
|
||||
local_sz = prod(self.full_shape[self.first_reduce-self.local_dims:self.first_reduce+self.group_for_reduces])
|
||||
smem_sz = amt*acc_sz*upcast_sz*local_sz
|
||||
check(smem_sz <= self.opts.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.opts.shared_max}")
|
||||
|
||||
if opt.op is OptOps.LOCAL: # cyan
|
||||
check(self.opts.has_local, "target does not support local")
|
||||
check(axis < self.global_dims, "local is for globals")
|
||||
self.shift_to(axis, amt, insert_before=self.first_reduce)
|
||||
self.local_dims += 1
|
||||
elif opt.op in {OptOps.GROUP, OptOps.GROUPTOP}: # green
|
||||
check(self.opts.has_local and self.opts.has_shared, "target does not support local or shared mem")
|
||||
check(self.first_reduce + self.group_for_reduces <= axis < self.first_upcast, "must be reduce axis to group")
|
||||
check(not self.tensor_core, "can't group with tensor cores")
|
||||
check(len(reduce_axes:=[i for r in self.reduceops for i in r.axis_arg]) == len(set(reduce_axes)), "can't group with parallel reduces")
|
||||
self.shift_to(axis, amt, top=(opt.op is OptOps.GROUPTOP), insert_before=self.first_reduce + self.group_for_reduces)
|
||||
self.group_for_reduces += 1
|
||||
elif opt.op is OptOps.UNROLL: # purple
|
||||
check(axis < self.first_upcast, "can't upcasted already upcasted")
|
||||
check(amt <= 32, "don't unroll more than 32")
|
||||
# TODO: fix upcast_count to put purples before yellows. broken because of METAL tensor cores
|
||||
#upcast_count = sum(x == y for x,y in zip(self.full_shape[-self.upcasted:], self.output_shape[-self.upcasted:])) if self.upcasted else 0
|
||||
#self.shift_to(axis, amt, insert_before=None if upcast_count == 0 else self.shape_len-upcast_count)
|
||||
if self.full_shape[axis] == amt and axis == self.first_reduce: self.local_dims += 1 # first_reduce will ++, so offset loss in simplify_ones
|
||||
if self.full_shape[axis] == amt and axis < self.first_reduce+self.group_for_reduces: self.group_for_reduces -= 1 # fully unrolling a GROUP
|
||||
self.shift_to(axis, amt, insert_before=None)
|
||||
self.upcast()
|
||||
elif opt.op is OptOps.UPCAST: # yellow
|
||||
check(axis < self.first_reduce, "upcast is for non-reduce")
|
||||
check(not (self.tensor_core and self.global_dims <= axis < self.global_dims+len(self.tensor_core.threads)), "can't upcast TC locals")
|
||||
check(amt <= 16, "don't upcast more than 16")
|
||||
self.shift_to(axis, amt, insert_before=None)
|
||||
self.upcast()
|
||||
elif opt.op is OptOps.NOLOCALS:
|
||||
check(self.opts.has_local and not self.dont_use_locals, "NOLOCALS is meaningless if target does not support local or already not using locals")
|
||||
check(self.local_dims == 0 and self.group_for_reduces == 0, "can't have no locals with locals")
|
||||
self.dont_use_locals = True
|
||||
elif opt.op is OptOps.SWAP:
|
||||
check(axis < amt < self.global_dims, f"swap is only for globals with axis < amt, getting {amt=}, {axis=}, {self.global_dims=}")
|
||||
permute = list(range(self.shape_len))
|
||||
permute[axis], permute[amt] = permute[amt], permute[axis]
|
||||
self.reshape_and_permute(None, tuple(permute))
|
||||
elif opt.op is OptOps.PADTO:
|
||||
check(not self.vars, "does not work with symbolic shape")
|
||||
check(axis < self.first_upcast, "cannot pad upcasted")
|
||||
# ok to pad SUM if all parent ALU ops have f(0) = 0
|
||||
if (r:=self.reduceop) is not None and self.first_reduce <= axis: check(r.arg[0] is Ops.ADD and can_pad(r, {}, set()), f"cannot pad {r}")
|
||||
padded = False
|
||||
for i,st in enumerate(self.sts):
|
||||
if (s:=st.shape[axis]) == 1: continue # reduced
|
||||
check(s > amt//4, f"pad adds more than quadruple the work {st.shape[axis]=} > {amt//4=}")
|
||||
if (ru := round_up(cast(int, s), amt) - s):
|
||||
# pad right seems to be faster
|
||||
self.sts[i] = st.pad(((0,0),) * axis + ((0,ru),) + ((0,0),) * (len(st.shape)-axis-1))
|
||||
padded = True
|
||||
check(padded, "nothing was padded")
|
||||
|
||||
if append_opt: self.applied_opts.append(opt)
|
||||
if self.simplify_ones() and self.tensor_core_opts:
|
||||
self.tensor_core_opts.fix_axes(axis) # fix up axes in TC opts if required after simplify_ones()
|
||||
|
||||
def required_optimizations(self) -> Kernel:
|
||||
if isinstance(self.membufs[0].dtype, ImageDType):
|
||||
unit_stride_axes_mul_4 = [i for i in self.sts[0].unit_stride_axes(ignore_valid=True) if self.sts[0].shape[i]%4 == 0]
|
||||
assert unit_stride_axes_mul_4, f"needs a unit stride axis in {self.bufs[0]}"
|
||||
if all(x < self.first_upcast for x in unit_stride_axes_mul_4): self.apply_opt(Opt(OptOps.UPCAST, unit_stride_axes_mul_4[0], 4))
|
||||
return self
|
||||
|
||||
def hand_coded_optimizations(self) -> Kernel:
|
||||
self.required_optimizations()
|
||||
|
||||
# should use matvec - TODO: adjust/tune based on the wide vs tall/large vs small mat
|
||||
MV_BLOCKSIZE, MV_THREADS_PER_ROW, MV_ROWS_PER_THREAD = getenv("MV_BLOCKSIZE", 4), getenv("MV_THREADS_PER_ROW", 8), getenv("MV_ROWS_PER_THREAD", 4)
|
||||
if self.opts.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \
|
||||
self.reduceop is not None and self.reduceop.arg[0] is Ops.ADD and len(self.full_shape) >= 2 and self.opts.has_shared and \
|
||||
(mulop:=self.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.LOAD and mulop.src[1].op is Ops.LOAD:
|
||||
st0, st1 = self.sts[self.bufs.index(mulop.src[0])], self.sts[self.bufs.index(mulop.src[1])]
|
||||
strides0, strides1 = st0.real_strides(), st1.real_strides()
|
||||
def has_expanded_axis(shape, strides): return any(resolve(s > 1) and not resolve(st != 0) for s,st in zip(shape,strides))
|
||||
if strides0[self.first_reduce] == 1 and not (has_expanded_axis(st0.shape, strides0) and has_expanded_axis(st1.shape, strides1)):
|
||||
for global_idx in range(self.global_dims):
|
||||
if self.full_shape[self.first_reduce]%MV_THREADS_PER_ROW == 0 and self.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
|
||||
if DEBUG >= 3:
|
||||
print(f"MATVEC: {self.full_shape=} {self.first_reduce=} {strides0=} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
|
||||
if MV_THREADS_PER_ROW > 1: self.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
|
||||
if MV_BLOCKSIZE > 1: self.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
|
||||
if MV_ROWS_PER_THREAD > 1: self.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
|
||||
return self
|
||||
|
||||
if self.opts.has_local and self.opts.has_shared and all_int(self.sts[0].shape[:self.first_reduce]):
|
||||
# are we grouping? (requires local shape support)
|
||||
if not self.float4_axis(0) and self.first_reduce <= 2 and self.first_reduce + 1 <= self.shape_len and prod(self.sts[0].shape[:self.first_reduce]) <= 2048: # noqa: E501
|
||||
# TODO: use 1024 if it's allowed in a smarter way
|
||||
for sz in ([256, 16] if prod(self.sts[0].shape[:self.first_reduce]) <= 32 else [16]):
|
||||
if all(st.shape[self.first_reduce] % sz == 0 or st.shape[self.first_reduce] == 1 for st in self.sts):
|
||||
try: # may fail due to excessive smem usage
|
||||
self.apply_opt(Opt(OptOps.GROUPTOP, 0, sz))
|
||||
break
|
||||
except KernelOptError: pass
|
||||
|
||||
# upcast float4 images
|
||||
for buf_index,buf in enumerate(self.bufs):
|
||||
unit_stride_axes_mul_4 = [i for i in self.sts[buf_index].unit_stride_axes(ignore_valid=True) if self.sts[buf_index].shape[i]%4 == 0]
|
||||
if buf.src[0].dtype.__class__ is ImageDType:
|
||||
#assert len(unit_stride_axes_mul_4) >= 1, f"needs a unit stride axis in {self.bufs[buf_index]}"
|
||||
if len(unit_stride_axes_mul_4) and all(x < self.first_upcast for x in unit_stride_axes_mul_4):
|
||||
if unit_stride_axes_mul_4[0] < self.first_reduce:
|
||||
self.apply_opt(Opt(OptOps.UPCAST, unit_stride_axes_mul_4[0], 4))
|
||||
else:
|
||||
self.apply_opt(Opt(OptOps.UNROLL, unit_stride_axes_mul_4[0]-self.first_reduce, 4))
|
||||
|
||||
# no more opt if we are grouping
|
||||
if self.group_for_reduces: return self
|
||||
|
||||
# **** below this line need to be optional and benchmarked ****
|
||||
|
||||
# TODO: doing extra upcasts with images doesn't work for some reason (maybe has to do with to_image_idx)
|
||||
# to trigger the above bug, remove prod(self.full_shape[self.first_upcast:]) from the below
|
||||
# expression and run test/test_ops.py with IMAGE=2
|
||||
# if there are small dims with lots of valid masks, upcast them (they might be from Tensor.stack)
|
||||
# this can be made much smarter
|
||||
to_upcast: list[int] = []
|
||||
# upcast leading axes first (hack-ish for winograd; we actually want to upcast masked axes with low stride first)
|
||||
for axis in range(self.first_reduce):
|
||||
# we might want to be able to split axes that are masked, or refuse to merge them in simplify_merge_adjacent
|
||||
# for now skip upcasting here if there is a symbolic axis
|
||||
if isinstance(self.full_shape[axis], int) and self.full_shape[axis] <= 7 and any(st.axis_is_masked(axis) for st in self.sts) and \
|
||||
prod(self.full_shape[self.first_upcast:]) * prod(self.full_shape[j] for j in to_upcast) * self.full_shape[axis] <= 7 * 7:
|
||||
if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
|
||||
to_upcast.append(axis)
|
||||
for axis in to_upcast[::-1]: self.apply_opt(Opt(OptOps.UPCAST, axis, 0))
|
||||
|
||||
# potentially do more upcasts of non reduce axes based on a heuristic
|
||||
upcasted_axis = set()
|
||||
while resolve(prod(self.sts[0].shape[:self.first_reduce]) >= 1024):
|
||||
xb_choices = []
|
||||
for axis, upcast_amount in itertools.product(range(self.first_reduce), [3,4]): # consider all the non reduce axes, and a 3 or 4 reduce
|
||||
# if we haven't upcasted it, it's not symbolic, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
|
||||
if axis not in upcasted_axis and isinstance(self.full_shape[axis], int) and self.full_shape[axis]%upcast_amount == 0 and any(st.views[-1].strides[axis] == 0 and not any(x[1] == 0 for x in self.upcasted_axis(buf_index)) for buf_index, st in enumerate(self.sts)): # noqa: E501
|
||||
xb_choices.append((sum(st.views[-1].strides[axis]>0 for st in self.sts), sum(st.views[-1].strides[axis] for st in self.sts), axis, upcast_amount)) # noqa: E501
|
||||
if xb_choices:
|
||||
xb_choices = sorted(xb_choices)
|
||||
if DEBUG >= 4: print(f"float4 merging axis : {xb_choices}")
|
||||
self.apply_opt(Opt(OptOps.UPCAST, xb_choices[0][2], xb_choices[0][3]))
|
||||
upcasted_axis.add(xb_choices[0][2])
|
||||
else: break
|
||||
|
||||
# if last dim is small(ish) and it's a reduce dim, upcast the reduce (loop unrolling). no simplify needed since it's just an upcast.
|
||||
if self.first_reduce < self.first_upcast and (prod(self.full_shape[self.first_upcast:]) <= 4 or not any(r for _,_,r in self.upcasted_axis(self.full_buf_index))) and (self.upcasted == 0 or prod(self.full_shape[-self.upcasted:]) < 64): # noqa: E501
|
||||
if isinstance(s:=self.full_unupcasted_shape[-1], int) and s <= 32: # NOTE: cannot loop unroll symbolic axis
|
||||
self.apply_opt(Opt(OptOps.UNROLL, len(self.full_unupcasted_shape)-1-self.first_reduce, 0))
|
||||
# if it's small, upcast a second reduce dimension too
|
||||
if self.first_reduce < self.first_upcast and s <= 3 and isinstance(s2:=self.full_unupcasted_shape[-1], int) and s2 <= 3:
|
||||
self.apply_opt(Opt(OptOps.UNROLL, len(self.full_unupcasted_shape)-1-self.first_reduce, 0))
|
||||
else:
|
||||
for splits in [4]:
|
||||
if self.full_unupcasted_shape[-1]%splits == 0:
|
||||
self.apply_opt(Opt(OptOps.UNROLL, len(self.full_unupcasted_shape)-1-self.first_reduce, splits))
|
||||
break
|
||||
|
||||
# if nothing at all is upcasted and it's easy to, do an upcast
|
||||
# TODO: this is breaking the tests
|
||||
for splits in [4]:
|
||||
if self.upcasted == 0 and self.full_unupcasted_shape and self.full_unupcasted_shape[-1] % splits == 0:
|
||||
self.apply_opt(Opt(OptOps.UPCAST, len(self.full_unupcasted_shape)-1, splits))
|
||||
|
||||
# **** local groups ****
|
||||
|
||||
if self.opts.has_local:
|
||||
if getenv("NOLOCALS") and self.local_dims == 0 and not self.group_for_reduces:
|
||||
self.apply_opt(Opt(OptOps.NOLOCALS))
|
||||
else:
|
||||
# prioritize making expand axes local
|
||||
local_axis_ranking = [(any(self.sts[buf_index].views[-1].strides[axis] == 0 for buf_index in range(len(self.sts))), axis) for axis in range(len(self.full_shape[:self.first_reduce]))] # noqa: E501
|
||||
to_local: list[tuple[int, int]] = []
|
||||
for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
|
||||
local_size = prod(sz for _, sz in to_local)
|
||||
local_sz: Optional[int] = next((x for x in ([32] * (axis == 0) + [16, 8, 4, 3, 2]) if self.full_shape[axis] % x == 0 and local_size * x <= 128), None) # noqa: E501
|
||||
if local_sz is not None: to_local.append((axis, local_sz))
|
||||
deleted_shape = 0
|
||||
for axis, local_sz in sorted(to_local[:3]):
|
||||
axis = axis - deleted_shape
|
||||
will_delete_shape = local_sz == self.full_shape[axis]
|
||||
self.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
|
||||
if will_delete_shape: deleted_shape += 1
|
||||
|
||||
return self
|
||||
|
||||
# **** kernel outputs ****
|
||||
|
||||
kernel_cnt: Final[DefaultDict[str, int]] = defaultdict(int)
|
||||
@functools.cached_property
|
||||
def name(self) -> str:
|
||||
# kernel name (before late upcast)
|
||||
kernel_type = "r" if self.reduceop is not None else ("C" if all(x.op is Ops.SINK or x.op in GroupOp.Buffer for x in self.ast.toposort) else "E")
|
||||
suffix = colored('_', 'BLACK').join([colored(x.render() if isinstance(x, UOp) else str(x), c) for x,c in zip(self.full_shape, self.colors())])
|
||||
name = kernel_type + (f"{len(self.ast.src)}" if len(self.ast.src) > 1 else "") + "_" + suffix
|
||||
|
||||
# name the function something unique
|
||||
Kernel.kernel_cnt[(function_name := to_function_name(name))] += 1
|
||||
num = f"n{Kernel.kernel_cnt[function_name]-1}" if Kernel.kernel_cnt[function_name] > 1 else ""
|
||||
return name + colored(num, 'BLACK')
|
||||
|
||||
def get_optimized_ast(self) -> UOp:
|
||||
@functools.lru_cache(None)
|
||||
def fixup_ast(op:UOp) -> UOp:
|
||||
ret = op.replace(src=tuple(fixup_ast(x) for x in op.src))
|
||||
if op.op in GroupOp.Buffer and op in self.bufs:
|
||||
st_uop = self.sts[self.bufs.index(op)].to_uop()
|
||||
return ret.replace(src=(st_uop,)) if op.op is Ops.VALID else ret.replace(src=(ret.src[0], st_uop, *ret.src[2:]))
|
||||
if op.op is Ops.SINK: return ret.replace(arg = KernelInfo(self.local_dims, self.upcasted, self.dont_use_locals))
|
||||
if op.op is Ops.REDUCE_AXIS:
|
||||
reduce_idx = len(self.bufs) + self.reduceops.index(op) * 2
|
||||
|
||||
def reduced_axes(start, stop):
|
||||
return tuple(i for i in range(start, stop) if resolve(self.sts[reduce_idx].shape[i] != self.sts[reduce_idx + 1].shape[i]))
|
||||
axes = reduced_axes(self.first_reduce + self.group_for_reduces, self.shape_len)
|
||||
grouped_axes = reduced_axes(self.first_reduce, self.first_reduce + self.group_for_reduces)
|
||||
|
||||
if (tc := self.tensor_core) and (self.use_tensor_cores == 1 or self.use_tensor_cores == 3):
|
||||
def fix_st(st: ShapeTracker, wd_pattern, tcd_pattern):
|
||||
st = ShapeTracker.from_shape(st.shape) # st needs to be contiguous
|
||||
wd, warp_dims = self.global_dims, tuple(sz for _, sz in tc.threads)
|
||||
tcd, tcd_dims = self.first_upcast, tuple(sz for _, sz in tc.reduce_axes + tc.early_upcast_axes)
|
||||
|
||||
assert st.shape[wd:wd+len(warp_dims)] == warp_dims, f"warp dims wrong: {st.shape[wd:wd+len(warp_dims)]=} != {warp_dims=}"
|
||||
assert st.shape[tcd:tcd+len(tcd_dims)] == tcd_dims, f"tcd dims wrong: {st.shape[tcd:tcd+len(tcd_dims)]=} != {tcd_dims=}"
|
||||
assert tc.expanded_shape is not None
|
||||
|
||||
new_shape = st.shape[:tcd] + tc.expanded_shape + st.shape[tcd+len(tcd_dims):] # expand the tcd
|
||||
permaxis = list(range(wd)) + [y + (wd if x == 0 else tcd) for x,y in wd_pattern] + list(range(wd+len(warp_dims),tcd)) + \
|
||||
[y + (wd if x == 0 else tcd) for x,y in tcd_pattern] + list(range(tcd+len(tc.expanded_shape),len(new_shape)))
|
||||
return st.reshape(new_shape).permute(tuple(permaxis)).reshape(st.shape).simplify()
|
||||
|
||||
srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
|
||||
for i, tc_pattern in enumerate([tc.st1_pattern, tc.st2_pattern]):
|
||||
if tc_pattern: srcs[i] = srcs[i].view(fix_st(srcs[i].st_arg if srcs[i].op is Ops.LOAD else srcs[i].src[0].st_arg, *tc_pattern))
|
||||
|
||||
if self.use_tensor_cores == 3: # for TC=3, emulate the warp addressing with locals
|
||||
local_shape = tuple(1 if i >= self.first_reduce and i < self.first_upcast else s for i, s in enumerate(self.full_shape))
|
||||
st = store_st = ShapeTracker.from_shape(local_shape)
|
||||
local_buffer = UOp(Ops.DEFINE_LOCAL, tc.dtype_in.ptr(size=st.real_size(), local=True), (), (f"temp{i + 1}", st.real_size()))
|
||||
if tc_pattern: store_st = fix_st(store_st, *tc_pattern)
|
||||
local_store = UOp.store(local_buffer, store_st.to_uop(), srcs[i])
|
||||
srcs[i] = UOp(Ops.LOAD, tc.dtype_in, (local_buffer, st.to_uop(), local_store))
|
||||
|
||||
tc_reduce_axes = tuple(self.first_upcast + ax for ax, _ in tc.reduce_axes)
|
||||
if self.use_tensor_cores == 1: # real WMMA, use CONTRACT/UNROLL to get the vectorization right
|
||||
upcast_axes = tuple(tuple((self.first_upcast + ax, sz) for ax, sz in up) for up in tc.upcast_axes)
|
||||
wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, prod(sz for _, sz in tc.threads), upcast_axes, tc_reduce_axes)
|
||||
wmma_sz = [prod(x[1] for x in l) for l in upcast_axes]
|
||||
wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(wmma_sz[2]), src=(
|
||||
UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(wmma_sz[0]), src=(srcs[0],), arg=upcast_axes[0]),
|
||||
UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(wmma_sz[1]), src=(srcs[1],), arg=upcast_axes[1]),
|
||||
UOp.const(tc.dtype_out.vec(wmma_sz[2]), 0.0)), arg=wmma_arg)
|
||||
tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=upcast_axes[2])
|
||||
|
||||
else: # for TC=3 MUL/SUM instead of WMMA
|
||||
tc_uop = UOp(Ops.REDUCE_AXIS, tc.dtype_out, ((srcs[0] * srcs[1]).cast(tc.dtype_out),), (Ops.ADD, tc_reduce_axes))
|
||||
|
||||
new_reduce_axes = tuple(i for i in axes if i not in tc_reduce_axes)
|
||||
return ret.replace(src=(tc_uop,), arg=(Ops.ADD, new_reduce_axes)) if new_reduce_axes else tc_uop
|
||||
|
||||
ret = ret.replace(arg = (op.arg[0], axes))
|
||||
if self.group_for_reduces and grouped_axes:
|
||||
local_shape = (1,) * self.global_dims + self.full_shape[self.global_dims:self.global_dims+self.local_dims] + \
|
||||
tuple([self.full_shape[i] if self.sts[reduce_idx].shape[i] != self.sts[reduce_idx+1].shape[i] else 1 \
|
||||
for i in range(self.first_reduce, self.first_reduce+self.group_for_reduces)]) + \
|
||||
(1,) * (self.shape_len - self.upcasted - self.group_for_reduces - self.first_reduce) + tuple([x[0] for x in self.upcasted_axis(0)])
|
||||
st_uop = ShapeTracker.from_shape(local_shape).to_uop()
|
||||
local_size = st_uop.arg.real_size()
|
||||
local_buffer = UOp(Ops.DEFINE_LOCAL, op.dtype.ptr(local_size, local=True), (), (f"temp{self.reduceops.index(op)+1}", local_size))
|
||||
local_load = UOp(Ops.LOAD, op.dtype, (local_buffer, st_uop, UOp.store(local_buffer, st_uop, ret)))
|
||||
grouped_reduce = UOp(Ops.REDUCE_AXIS, op.dtype, (local_load,), arg=(op.arg[0], grouped_axes))
|
||||
if op is self.reduceops[-1]: return grouped_reduce
|
||||
st_uop = ShapeTracker.from_shape(tuple([1 if i in grouped_axes else a for i,a in enumerate(local_shape)])).to_uop()
|
||||
return UOp(Ops.LOAD, op.dtype, (local_buffer, st_uop, UOp.store(local_buffer, st_uop, grouped_reduce)))
|
||||
|
||||
return ret
|
||||
|
||||
return graph_rewrite(fixup_ast(self.ast), view_left)
|
||||
|
||||
# **** this is the lowerer ****
|
||||
|
||||
@track_rewrites()
|
||||
def linearize(self) -> Kernel:
|
||||
modified_ast = self.get_optimized_ast()
|
||||
|
||||
if DEBUG >= 3:
|
||||
print(self.name)
|
||||
if getenv("RAWAST"): print(self.ast)
|
||||
print(modified_ast)
|
||||
print(self.applied_opts)
|
||||
verify_ast(modified_ast)
|
||||
|
||||
self.uops:list[UOp] = linearize_uop(full_graph_rewrite(rewrite_shapetracker_with_index(modified_ast, self.opts), self.opts))
|
||||
if DEBUG >= 5: print_uops(self.uops)
|
||||
return self
|
||||
|
||||
def to_program(self, name_override:Optional[str]=None) -> ProgramSpec:
|
||||
self.linearize()
|
||||
src = self.opts.render(name:=to_function_name(ansiname:=(name_override if name_override is not None else self.name)), self.uops)
|
||||
|
||||
if getenv("RUN_PROCESS_REPLAY"):
|
||||
from test.external.process_replay.helpers import get_process_replay_ctx
|
||||
diskcache_put("kernel_process_replay", str(id(self)), (self.ast, self.opts, self.applied_opts, name, *get_process_replay_ctx(), src))
|
||||
|
||||
# group non-local bufs by the op type (LOAD or STORE) and the buffer arg. take the max access of that buffer in bytes
|
||||
# TODO: these max and min don't work on symbolic, and results are very wrong.
|
||||
mem_bytes = sum(max(x.src[0].dtype.itemsize * x.st_arg.real_size() for x in group)
|
||||
for _, group in itertools.groupby([x for x in self.ast.toposort if x.op in GroupOp.Buffer and x.src[0].op is Ops.DEFINE_GLOBAL],
|
||||
key=lambda x: (x.op, x.src[0].arg)))
|
||||
return ProgramSpec(ansiname, src, self.opts.device, self.uops, mem_estimate=mem_bytes,
|
||||
global_size=[1,1,1] if self.opts.has_local else None, local_size=[1,1,1] if self.opts.has_local else None)
|
||||
|
||||
# the living definition of intermediate UOps
|
||||
|
||||
def _assert_valid_uop(uop:UOp, st:ShapeTracker, sts:dict[UOp, ShapeTracker]) -> None:
|
||||
if not uop.has_st or uop in sts: return
|
||||
# restore globals from the two stage reduce
|
||||
if uop.op is Ops.LOAD and uop.src[0].op is Ops.DEFINE_LOCAL:
|
||||
_assert_valid_uop(local_reduce:=uop.src[2].src[2], uop.st_arg, sts)
|
||||
sts[uop] = sts[local_reduce]
|
||||
return
|
||||
for x in uop.src: _assert_valid_uop(x, st, sts)
|
||||
# only reduceuop is allowed to change shape, limited to turning n to 1
|
||||
if uop.op in {Ops.REDUCE_AXIS, Ops.WMMA}: st = ShapeTracker.from_shape(sts[uop.src[0]].reduce(uop.axis_arg))
|
||||
# movementops are pushed to VIEW
|
||||
elif uop.op is Ops.VIEW:
|
||||
assert len(uop.src) == 0, f"can't swizzle in kernel yet {uop}"
|
||||
st = uop.arg
|
||||
# everything else inherits shape
|
||||
else:
|
||||
if len(src_sts:=[sts[x] for x in uop.src if x in sts]) == 0: return None
|
||||
st = src_sts[0]
|
||||
if not all_same(shapes:=[x.shape for x in src_sts]):
|
||||
if all_same(sizes:=[prod(x) for x in shapes]): raise AssertionError(f"found implicit reshape {shapes}")
|
||||
raise AssertionError(f"found implicit expand {sizes} {shapes}")
|
||||
sts[uop] = st
|
||||
|
||||
def verify_ast(ast:UOp) -> dict[UOp, ShapeTracker]:
|
||||
assert ast.op is Ops.SINK and all(x.op is Ops.STORE for x in ast.src), "must be SINK"
|
||||
assert all_same([x.st_arg.size for x in ast.src]), "outputs must be exactly the same size"
|
||||
sts: dict[UOp, ShapeTracker] = {}
|
||||
for out in ast.src: _assert_valid_uop(out, out.st_arg, sts)
|
||||
shape_dims = [sorted(dedup(dims)) for dims in zip(*[x.shape for x in sts.values()])]
|
||||
assert all(len(x) == 1 or (len(x) == 2 and x[0] == 1) for x in shape_dims), f"shapes must have either 1 or n in each dimension, {shape_dims}"
|
||||
type_verify(list(sts))
|
||||
return sts
|
||||
@@ -0,0 +1,224 @@
|
||||
from __future__ import annotations
|
||||
import collections, heapq
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.ops import type_verify, UOp, Ops, PatternMatcher, UPat, graph_rewrite, GroupOp
|
||||
from tinygrad.dtype import dtypes, PtrDType
|
||||
from tinygrad.helpers import dedup, flatten, partition
|
||||
|
||||
DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST, *GroupOp.Block}
|
||||
|
||||
def disp(y:UOp) -> str:
|
||||
if y.op is Ops.BLOCKSTART: return "w"+disp(y.src[0])
|
||||
if y.op is Ops.IF: return f'IF{id(y)}'
|
||||
if y.op is Ops.RANGE: return str(y.arg)
|
||||
return "<NONE>"
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BasicBlock:
|
||||
ctx: tuple[UOp, ...]
|
||||
lst: tuple[UOp, ...]
|
||||
end: UOp|None = None
|
||||
def __lt__(self, o:BasicBlock): return tuple(x.tuplize for x in self.ctx+self.lst) < tuple(x.tuplize for x in o.ctx+o.lst)
|
||||
def __repr__(self):
|
||||
return f"{(str(disp(self.end))+' ') if self.end is not None else ''}"+\
|
||||
f"{[disp(y) for y in self.ctx]} {len(self.lst)}" + "\n" + '\n'.join([str(x.op) for x in self.lst])
|
||||
|
||||
def append_to_block(ctx:tuple[dict[UOp, tuple[UOp, ...]], dict[UOp, list[UOp]]], x:UOp):
|
||||
block_ctxs, children = ctx
|
||||
in_this_block = set(x.arg.lst)
|
||||
|
||||
# collections to build
|
||||
new_srcs: list[UOp] = []
|
||||
to_append: list[UOp] = []
|
||||
old_blocks: dict[tuple[UOp, ...], UOp] = {}
|
||||
new_blocks: dict[tuple[UOp, ...], list[UOp]] = {}
|
||||
|
||||
for u in x.src:
|
||||
if u.op is Ops.BLOCK:
|
||||
# merge sibling blocks. NOTE: blocks must only have one output source
|
||||
assert u.arg.ctx not in old_blocks, "sibiling should never have been created"
|
||||
old_blocks[u.arg.ctx] = u
|
||||
elif u.op not in DONT_PLACE_IN_BLOCK and set(children[u]).issubset(in_this_block):
|
||||
# if it can go in blocks and all its children are in the block, we add it to the block
|
||||
if (block_ctx:=block_ctxs[u]) == x.arg.ctx:
|
||||
# if it's the same context, we place the UOp in this block and append the parents to its srcs
|
||||
new_srcs.extend(u.src)
|
||||
to_append.append(u)
|
||||
else:
|
||||
# if it's a different context, we create a new block with this UOp
|
||||
new_blocks.setdefault(block_ctx, []).append(u)
|
||||
else:
|
||||
# otherwise, we keep it in the srcs
|
||||
new_srcs.append(u)
|
||||
if len(to_append) == 0 and len(new_blocks) == 0: return None
|
||||
|
||||
for rng,lst in new_blocks.items():
|
||||
srcs = flatten(y.src for y in lst)
|
||||
if (old_block:=old_blocks.pop(rng, None)) is not None:
|
||||
# NOTE: order shouldn't matter here
|
||||
srcs.extend(old_block.src)
|
||||
lst.extend(old_block.arg.lst)
|
||||
new_block = UOp(Ops.BLOCK, dtypes.void, tuple(dedup(srcs)), BasicBlock(rng, tuple(lst)))
|
||||
lrng = list(rng)
|
||||
for r in rng[::-1]:
|
||||
if r not in x.arg.ctx and r.op is not Ops.BLOCKSTART:
|
||||
lrng.remove(r)
|
||||
new_block = UOp(Ops.BLOCKEND, src=(new_block,),
|
||||
arg=BasicBlock(tuple(lrng), (UOp(Ops.ENDIF if r.op is Ops.IF else Ops.ENDRANGE, src=(r,)),), r))
|
||||
new_srcs.append(new_block)
|
||||
return UOp(Ops.BLOCK, dtypes.void, tuple(dedup(list(old_blocks.values())+new_srcs)), BasicBlock(x.arg.ctx, tuple(to_append)+x.arg.lst))
|
||||
|
||||
make_basic_blocks = PatternMatcher([
|
||||
(UPat(Ops.SINK, name="x"), lambda x: UOp(Ops.BLOCK, src=x.src, arg=BasicBlock((), (x,)))),
|
||||
(UPat(Ops.BLOCK, name="x"), append_to_block),
|
||||
])
|
||||
|
||||
def block_merge(ctx, x:UOp):
|
||||
# ctx is children here
|
||||
if x.op is Ops.BLOCKEND:
|
||||
# if it's a BLOCKEND, see if we are done with placement. if all the children of the range are in here
|
||||
in_this_block = set(x.arg.lst)
|
||||
if len([y for y in ctx[x.arg.end] if y not in in_this_block]) == 0:
|
||||
# find the parent block that has the BLOCKSTART in the ctx
|
||||
parent_blocks = [y for y in x.src if y.op is Ops.BLOCK and UOp(Ops.BLOCKSTART, src=(x.arg.end,)) in y.arg.ctx]
|
||||
assert len(parent_blocks) <= 1, "should never have two parent blocks"
|
||||
if len(parent_blocks) == 1:
|
||||
parent_block = parent_blocks[0]
|
||||
# range needs DEFINE_ACC to be before the range (never in DEFINE_ACC for if)
|
||||
early_ops, late_ops = partition(x.arg.lst, lambda y: y.op is Ops.DEFINE_ACC and x.arg.end in y.src)
|
||||
return UOp(Ops.BLOCK, dtypes.void, tuple(y for y in x.src if y is not parent_block)+parent_block.src,
|
||||
BasicBlock(tuple(y for y in x.arg.ctx if y is not x.arg.end), tuple(early_ops)+parent_block.arg.lst+tuple(late_ops)))
|
||||
|
||||
new_srcs: list[UOp] = []
|
||||
to_append: list[UOp] = []
|
||||
new_ctx = x.arg.ctx
|
||||
placed = set()
|
||||
for u in x.src:
|
||||
if u.op is Ops.BLOCK and (tuple(u.arg.ctx) == tuple(x.arg.ctx) or (x.arg.end is not None and x.arg.end in u.arg.ctx)):
|
||||
# NOTE: this can't appear in srcs twice or it would be a BLOCKFORK
|
||||
new_ctx += tuple(y for y in u.arg.ctx if y not in x.arg.ctx)
|
||||
new_srcs.extend(u.src)
|
||||
to_append.extend(u.arg.lst)
|
||||
elif u.op is Ops.BLOCKFORK and x.src.count(u) == u.arg: # block fork appears # of times in srcs
|
||||
if u not in placed:
|
||||
new_srcs.extend(u.src)
|
||||
placed.add(u)
|
||||
else:
|
||||
# keep it in srcs
|
||||
new_srcs.append(u)
|
||||
if len(to_append) == 0 and len(placed) == 0: return None
|
||||
return UOp(x.op, dtypes.void, tuple(new_srcs), BasicBlock(tuple(sorted(new_ctx, key=lambda x: x.tuplize)), tuple(to_append)+x.arg.lst, x.arg.end))
|
||||
|
||||
pm_block_merge = PatternMatcher([(UPat((Ops.BLOCKEND, Ops.BLOCK), name="x"), block_merge),])
|
||||
|
||||
# NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed
|
||||
def block_reorder(in_block:UOp):
|
||||
in_this_block = set(in_block.arg.lst)
|
||||
local_children: collections.defaultdict[UOp, list[UOp]] = collections.defaultdict(list)
|
||||
in_degree: collections.defaultdict[UOp, int] = collections.defaultdict(int)
|
||||
priorities:dict[UOp, int] = {}
|
||||
|
||||
# get local children and assign priorities
|
||||
for u in reversed(in_block.arg.lst):
|
||||
for s in u.src:
|
||||
if s in in_this_block:
|
||||
local_children[s].append(u)
|
||||
in_degree[u] += 1
|
||||
# put loads in the beginning of the block and prevent priority inversion
|
||||
priorities[u] = min([-1000 if u.op is Ops.LOAD else 0] + [priorities[x] for x in local_children[u]])
|
||||
|
||||
# placement queue
|
||||
queue:list[tuple[int, tuple, UOp]] = []
|
||||
def push(u:UOp): heapq.heappush(queue, (priorities[u], u.tuplize, u))
|
||||
|
||||
# place the first ones that don't have deps
|
||||
for u in in_block.arg.lst:
|
||||
if u not in in_degree: push(u)
|
||||
|
||||
newlst = []
|
||||
while queue:
|
||||
_,_,x = heapq.heappop(queue)
|
||||
newlst.append(x)
|
||||
for u in local_children[x]:
|
||||
in_degree[u] -= 1
|
||||
if in_degree[u] == 0: push(u)
|
||||
|
||||
assert len(newlst) == len(in_block.arg.lst), f"len mismatch {len(newlst)} != {len(in_block.arg.lst)}"
|
||||
return in_block.replace(arg=BasicBlock(in_block.arg.ctx, tuple(newlst)))
|
||||
|
||||
def linearize_uop(sink:UOp, skip_check:bool=not __debug__) -> list[UOp]:
|
||||
assert sink.op is Ops.SINK, f"sink isn't sink, it's {sink.op}"
|
||||
|
||||
# get children and all block contexts
|
||||
temp_block_ctxs: dict[UOp, list[UOp]] = {}
|
||||
children: dict[UOp, list[UOp]] = {}
|
||||
for u in sink.toposort:
|
||||
this_block_ctx: list[UOp] = []
|
||||
for s in u.src:
|
||||
# save children
|
||||
children.setdefault(s, []).append(u)
|
||||
# compute block ctx
|
||||
if s.op in {Ops.RANGE, Ops.IF}: this_block_ctx.append(s)
|
||||
# don't flow (fully) through assign and store
|
||||
elif s.op is Ops.STORE:
|
||||
# ugh, deal with non-reduce locals. probably wrong
|
||||
if isinstance(s.src[0].dtype, PtrDType) and s.src[0].dtype.local:
|
||||
idx_context, store_context = temp_block_ctxs[s.src[0]], temp_block_ctxs[s]
|
||||
this_block_ctx += [x for x in store_context if x not in idx_context and x.op is Ops.RANGE]
|
||||
elif s.op is Ops.ASSIGN:
|
||||
# flow though assign, but remove the ranges used in the assign
|
||||
assert s.src[0].op is Ops.DEFINE_ACC
|
||||
this_block_ctx += [x for x in temp_block_ctxs[s.src[1]] if x not in s.src[0].src[1:]]
|
||||
else:
|
||||
# flow though everything else
|
||||
this_block_ctx += temp_block_ctxs[s]
|
||||
temp_block_ctxs[u] = sorted(dedup(this_block_ctx), key=lambda x: x.tuplize)
|
||||
|
||||
# make final block_ctxs, add BLOCKSTART to block_ctxs for IF and RANGE
|
||||
block_ctxs: dict[UOp, tuple[UOp, ...]] = {}
|
||||
for u in sink.toposort:
|
||||
block_ctxs[u] = ((UOp(Ops.BLOCKSTART, src=(u,)),) + tuple(temp_block_ctxs[u])) if u.op in {Ops.IF, Ops.RANGE} else tuple(temp_block_ctxs[u])
|
||||
|
||||
# TODO: there's probably a clever way to remove this while loop
|
||||
while 1:
|
||||
sink = graph_rewrite(sink, make_basic_blocks, ctx=(block_ctxs, children))
|
||||
|
||||
# add BLOCKFORK (slow!)
|
||||
block_parent_count = collections.Counter(flatten([x.src for x in sink.toposort if x.op is Ops.BLOCK]))
|
||||
non_block_parents = set(flatten([x.src for x in sink.toposort if x.op is not Ops.BLOCK]))
|
||||
forks = {u:UOp(Ops.BLOCKFORK, src=(UOp(Ops.BLOCK, src=u.src, arg=BasicBlock(block_ctxs[u], (u,))),), arg=child_count)
|
||||
for u,child_count in block_parent_count.items() if u.op not in DONT_PLACE_IN_BLOCK and child_count > 1 and u not in non_block_parents}
|
||||
|
||||
if not len(forks): break
|
||||
sink = sink.substitute(forks)
|
||||
|
||||
# combine matching BLOCKENDS
|
||||
blockends_to_arg: dict[UOp, list[UOp]] = {}
|
||||
for be in sink.toposort:
|
||||
if be.op is Ops.BLOCKEND: blockends_to_arg.setdefault(be.arg.end, []).append(be)
|
||||
new_forks = {}
|
||||
for k,v in blockends_to_arg.items():
|
||||
# NOTE: if any BLOCKEND is the parent of any other with the same arg, this algo fails
|
||||
if len(v) > 1:
|
||||
out = UOp(Ops.BLOCKFORK, src=(UOp(Ops.BLOCKEND, src=tuple(flatten(x.src for x in v)),
|
||||
arg=BasicBlock(tuple(dedup(flatten([y.arg.ctx for y in v]))), v[0].arg.lst, k)),), arg=len(v))
|
||||
for u in v: new_forks[u] = out
|
||||
sink = sink.substitute(new_forks)
|
||||
|
||||
# reorder ops in block for speed
|
||||
sink = sink.substitute({u:newu for u in sink.toposort if u.op is Ops.BLOCK and (newu:=block_reorder(u)) is not u})
|
||||
|
||||
# final rewrite to merge all blocks into one
|
||||
sink = graph_rewrite(sink, pm_block_merge, ctx=children)
|
||||
|
||||
# there should just be one block left, with a few parents with 0 srcs
|
||||
assert sink.op is Ops.BLOCK
|
||||
_uops = sorted(dedup(sink.src), key=lambda x: x.tuplize)
|
||||
assert all(len(x.src) == 0 and x.op not in {Ops.BLOCK, Ops.BLOCKSTART, Ops.BLOCKEND, Ops.BLOCKFORK} for x in _uops)
|
||||
_uops += sink.arg.lst
|
||||
|
||||
# sanity checks (NOTE: these can cause things to be skipped in BEAM)
|
||||
if not skip_check: type_verify(_uops)
|
||||
|
||||
# strip the SINK
|
||||
return _uops[:-1]
|
||||
@@ -0,0 +1,133 @@
|
||||
# the job of the lowerer is to do indexing
|
||||
import functools, itertools, operator
|
||||
from dataclasses import dataclass
|
||||
from typing import cast
|
||||
from tinygrad.dtype import dtypes, PtrDType
|
||||
from tinygrad.ops import KernelInfo, UOp, Ops, graph_rewrite, PatternMatcher, UPat, sint, identity_element, sint_to_uop
|
||||
from tinygrad.renderer import Renderer
|
||||
from tinygrad.helpers import all_int, prod, partition, flatten
|
||||
|
||||
# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
|
||||
def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
|
||||
acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
|
||||
try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
|
||||
except ValueError: return None
|
||||
return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
|
||||
|
||||
# ***** indexing *****
|
||||
|
||||
def _limit_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
|
||||
# TODO: symbolic shape
|
||||
if not all_int(dims): return dims
|
||||
while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
|
||||
for i,m in enumerate(max_sizes):
|
||||
if dims[i] * dims[i+1] <= m:
|
||||
dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
|
||||
break
|
||||
else: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
|
||||
return dims
|
||||
|
||||
def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
|
||||
if reverse: dims = dims[::-1]
|
||||
limited = _limit_dims(dims, max_sizes) if max_sizes is not None else dims
|
||||
ret = raw_idxs = [UOp(Ops.SPECIAL, dtypes.int, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)]
|
||||
if limited != dims:
|
||||
ret = []
|
||||
if (contraction:=get_contraction(dims, limited)) is None: raise AssertionError(f"get_contraction should not be None {dims=} {limited=}")
|
||||
for idx, contraction_group in zip(raw_idxs, contraction):
|
||||
for c in contraction_group[:-1]:
|
||||
ret.append(idx % dims[c])
|
||||
idx //= dims[c]
|
||||
ret.append(idx)
|
||||
return ret[::-1] if reverse else ret
|
||||
|
||||
@dataclass
|
||||
class IndexContext:
|
||||
idxs: list[UOp]
|
||||
ridxs: list[UOp]
|
||||
acc_num: int = 0
|
||||
|
||||
def get_index(ast:UOp, opts:Renderer) -> IndexContext:
|
||||
ki = ast.arg if isinstance(ast.arg, KernelInfo) else KernelInfo()
|
||||
# NOTE: assumes the shape is <global dims> <local dims> <group_for_reduces> <reduces> <upcasts/unrolls>
|
||||
full_shape = ast.full_shape
|
||||
first_upcasted = len(full_shape)-ki.upcasted
|
||||
# if there's no reduce, this is first_upcasted. assumes reduces are at the end
|
||||
first_reduce = min([first_upcasted]+flatten(x.axis_arg for x in ast.toposort if x.op is Ops.REDUCE_AXIS))
|
||||
local_loads = [x for x in ast.toposort if x.op is Ops.LOAD and x.src[0].op is Ops.DEFINE_LOCAL]
|
||||
# NOTE: sum up the reduced axes looking across all local loads, yields the number of grouped reduces
|
||||
group_for_reduces = sum([any(l.st_arg.shape[i]!=ast.src[0].st_arg.shape[i] for l in local_loads) for i in range(first_reduce,first_upcasted)])
|
||||
global_dims = first_reduce-ki.local_dims
|
||||
|
||||
if opts.has_local:
|
||||
if ki.dont_use_locals:
|
||||
assert ki.local_dims == 0, "can't use locals if there's no local dims"
|
||||
idxs = get_grouped_dims("idx", full_shape[:global_dims], opts.global_max, reverse=True)
|
||||
else:
|
||||
# define indexes for GPU-like execution
|
||||
idxs = get_grouped_dims("gidx", full_shape[:global_dims], opts.global_max, reverse=True) + \
|
||||
get_grouped_dims("lidx", full_shape[global_dims:first_reduce+group_for_reduces], opts.local_max)
|
||||
else:
|
||||
# all loops are RANGES
|
||||
idxs = [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(0), sint_to_uop(g)), i) for i,g in enumerate(full_shape[:first_reduce])]
|
||||
|
||||
# reduce loops
|
||||
idxs += [UOp(Ops.RANGE, dtypes.int, (sint_to_uop(0), sint_to_uop(g)), i)
|
||||
for i,g in enumerate(full_shape[first_reduce+group_for_reduces:first_upcasted], start=first_reduce+group_for_reduces)]
|
||||
|
||||
# upcast loops
|
||||
for i,g in enumerate(full_shape[first_upcasted:], start=first_upcasted):
|
||||
assert isinstance(g, int), "needs to be int to upcast/unroll"
|
||||
idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(g), tuple(range(g))),), ((i,g),)))
|
||||
|
||||
# late indexes (group for reduce)
|
||||
ridxs = idxs[:]
|
||||
for a in range(first_reduce, first_reduce+group_for_reduces):
|
||||
ridxs[a] = UOp(Ops.RANGE, dtypes.int, (sint_to_uop(0), sint_to_uop(full_shape[a])), 1000+a)
|
||||
|
||||
return IndexContext(idxs, ridxs)
|
||||
|
||||
# ***** lowering (given index) *****
|
||||
|
||||
def lower_reduce_axis(ctx: IndexContext, x: UOp):
|
||||
# NOTE: always using ridxs is fine here
|
||||
reduce_range, reduce_expand = partition([ctx.ridxs[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
|
||||
assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
|
||||
alu_op: Ops = x.arg[0]
|
||||
ret = x.src[0]
|
||||
if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
|
||||
ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis))
|
||||
ret = functools.reduce(lambda x,y: x.alu(alu_op, y), [ret.gep(i) for i in range(ret.dtype.count)])
|
||||
if not len(reduce_range): return ret
|
||||
# create ACC and assign
|
||||
acc = UOp(Ops.DEFINE_ACC, x.dtype, (x.const_like(identity_element(alu_op, x.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,))
|
||||
ctx.acc_num += 1
|
||||
return acc.assign(acc.alu(alu_op, ret))
|
||||
|
||||
def lower_load_store(ctx: IndexContext, x: UOp):
|
||||
idx, valid = x.st_arg.to_indexed_uops(ctx.ridxs if x.op is Ops.LOAD and x.src[0].op is Ops.DEFINE_LOCAL else ctx.idxs)
|
||||
buf = x.src[0]
|
||||
if x.op is Ops.LOAD:
|
||||
barrier = (UOp(Ops.BARRIER, dtypes.void, (x.src[2],)),) if x.src[0].op is Ops.DEFINE_LOCAL else ()
|
||||
return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + barrier)
|
||||
# NOTE: only store the local reduceop in the threads that are actually doing the reduce
|
||||
if cast(PtrDType, x.src[0].dtype).local and x.src[2].op is Ops.ASSIGN:
|
||||
reduce_input = x.src[2].src[1].src[1] if x.src[2].src[1].src[1] is not x.src[2].src[0] else x.src[2].src[1].src[0]
|
||||
store_back = reduce_input.op is Ops.LOAD and cast(PtrDType, reduce_input.src[0].dtype).local
|
||||
else: store_back = False
|
||||
# NOTE: If we're storing the reduced value back into each thread, need to zero-out the reduced axes
|
||||
if store_back: idx, _ = x.st_arg.to_indexed_uops([u.const_like(0) if u in x.src[2].src else u for u in ctx.idxs])
|
||||
if (not cast(PtrDType, x.src[0].dtype).local) or store_back:
|
||||
for oidx, ridx in zip(ctx.idxs, ctx.ridxs):
|
||||
if oidx is not ridx: valid = valid * oidx.eq(0)
|
||||
return UOp(Ops.STORE, dtypes.void, (buf.index(idx, valid), x.src[2]))
|
||||
|
||||
pm_lowerer = PatternMatcher([
|
||||
(UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis),
|
||||
(UPat(Ops.VALID, src=(UPat(Ops.VIEW),), name="x"), lambda ctx,x: x.st_arg.to_indexed_uops(ctx.idxs)[1]),
|
||||
# rewrite LOAD/STORE VIEW to LOAD/STORE with indexed
|
||||
(UPat((Ops.LOAD, Ops.STORE), src=(UPat(), UPat(Ops.VIEW)), allow_any_len=True, name="x"), lower_load_store),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"), UPat.const(dtypes.bool, True))), lambda b, idx: b.index(idx)),
|
||||
])
|
||||
|
||||
def rewrite_shapetracker_with_index(ast:UOp, opts:Renderer) -> UOp: return graph_rewrite(ast, pm_lowerer, ctx=get_index(ast, opts))
|
||||
@@ -0,0 +1,256 @@
|
||||
import math
|
||||
from tinygrad.dtype import dtypes, DType
|
||||
from tinygrad.helpers import polyN
|
||||
from tinygrad.ops import UOp
|
||||
|
||||
TRANSCENDENTAL_SUPPORTED_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64)
|
||||
|
||||
def _lazy_map_numbers(x:UOp, inf:UOp, _inf:UOp, nan:UOp, ratio:UOp):
|
||||
"""replace inf -> inf, -inf -> _inf, nan -> nan, otherwise -> ratio"""
|
||||
return x.ne(math.inf).where(x.ne(x).where(nan, x.ne(-math.inf).where(ratio, _inf)), inf)
|
||||
|
||||
# *** helper functions for bit manipulation ***
|
||||
def mantissa_bits(d:DType) -> int: return dtypes.finfo(d)[1]
|
||||
def exponent_bias(d:DType) -> int: return {dtypes.float64: 1023, dtypes.float32: 127, dtypes.float16: 15}[d]
|
||||
def exponent_mask(d:DType) -> int: return {dtypes.float64: 2047, dtypes.float32: 255, dtypes.float16: 31}[d]
|
||||
|
||||
# **** utils ****
|
||||
def shr(x:UOp, y:int) -> UOp: return x // (2**y)
|
||||
def shl(x:UOp, y:int) -> UOp: return x * (2**y)
|
||||
|
||||
def rintk(d:UOp) -> UOp:
|
||||
"""round d:float to int away from 0"""
|
||||
out_dtype = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype]
|
||||
return (d + (d<0.0).where(d.const_like(-0.5), d.const_like(0.5))).cast(out_dtype)
|
||||
|
||||
def pow2if(q:UOp, float_dtype:DType):
|
||||
"""cast(2^q, float_dtype) where q is any integer in the range of [-126, 127]"""
|
||||
out_dtype = {dtypes.int64: dtypes.float64, dtypes.int32: dtypes.float32, dtypes.int16: float_dtype}[q.dtype]
|
||||
return shl(q + exponent_bias(out_dtype), mantissa_bits(out_dtype)).bitcast(out_dtype)
|
||||
|
||||
def ilogb2k(d:UOp) -> UOp:
|
||||
"""calculate the integer part of log2(d), where d is normalized fp value in the range of [0, +inf)."""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
dint = d.bitcast({dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}[d.dtype])
|
||||
# -1 <= ilog2bk(d) <= 128
|
||||
return (shr(dint, mantissa_bits(d.dtype)) & exponent_mask(d.dtype)) - exponent_bias(d.dtype)
|
||||
|
||||
def ldexp3k(d:UOp, e:UOp) -> UOp:
|
||||
"""d*2^e. e is a number obtained by casting an integer in the range [-127, 127] to a float. d is any float number."""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES and e.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
cast_map = {dtypes.float64: dtypes.int64, dtypes.float32: dtypes.int32, dtypes.float16: dtypes.int16}
|
||||
m1 = d.bitcast(cast_map[d.dtype])
|
||||
m2 = shl(e.cast(cast_map[d.dtype]), mantissa_bits(d.dtype))
|
||||
return (m1 + m2).bitcast(d.dtype).cast(d.dtype)
|
||||
|
||||
def ldexp2k(d:UOp, e:UOp) -> UOp:
|
||||
"""d*2^e. much faster than ldexp3k but risky. d > 0 and d is not denormal."""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES and e.dtype in (dtypes.int16, dtypes.int32, dtypes.int64)
|
||||
return (d * pow2if(shr(e, 1), d.dtype)) * pow2if(e - shr(e, 1), d.dtype)
|
||||
|
||||
def frexp(v:UOp) -> tuple[UOp, UOp]:
|
||||
"""frexp(v) -> (mantissa, exponent) assuming v != 0"""
|
||||
assert v.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
# m1 = masks for mantissa, m2 = masks to normalize the mantissa.
|
||||
m1 = {dtypes.float64: 0x000FFFFFFFFFFFFF, dtypes.float32: 0x807FFFFF, dtypes.float16: 0x83FF}[v.dtype]
|
||||
m2 = {dtypes.float64: 0x3FE0000000000000, dtypes.float32: 0x3F000000, dtypes.float16: 0x3800}[v.dtype]
|
||||
bits = v.bitcast({dtypes.float64: dtypes.uint64, dtypes.float32: dtypes.uint32, dtypes.float16: dtypes.uint16}[v.dtype])
|
||||
exponent = shr(bits, mantissa_bits(v.dtype)) & exponent_mask(v.dtype)
|
||||
# Set the exponent bits appropriately to normalize the mantissa into the range of [0.5, 1.0).
|
||||
mantissa = ((bits & m1) | m2).bitcast(v.dtype)
|
||||
exp = exponent - exponent_bias(v.dtype) + 1
|
||||
return mantissa, exp
|
||||
|
||||
# *** reduction algorithms for sine ***
|
||||
def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
|
||||
"""
|
||||
Performs Payne-Hanek Reduction: computes the remainder of `d` modulo pi/2 for the values `d` where
|
||||
39800.0 <= d <= +Inf
|
||||
Returns a tuple of `(r, q)`:
|
||||
- `r`[d.dtype] is the reminder value corresponding to `round_to_nearest(x % pi/2)`.
|
||||
- `q`[int32] is an integer, and q % 4 is corresponding to the quadrant of the original angle `d`.
|
||||
"""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
# https://stackoverflow.com/questions/30463616/payne-hanek-algorithm-implementation-in-c/30465751#30465751
|
||||
# 190 bits of 2/pi for Payne-Hanek style argument reduction
|
||||
two_over_pi_f = [0x00000000, 0x28be60db, 0x9391054a, 0x7f09d5f4, 0x7d4d3770, 0x36d8a566, 0x4f10e410]
|
||||
|
||||
intermediate_dtype = dtypes.float32 if d.dtype == dtypes.float16 else d.dtype
|
||||
|
||||
f, e = frexp(d)
|
||||
ia = (f.cast(intermediate_dtype) * 4.294967296e9).cast(dtypes.uint64)
|
||||
# extract 96 relevant bits of 2/pi based on magnitude of argument
|
||||
i = shr(e.cast(dtypes.uint64), 5)
|
||||
e = e.cast(dtypes.int32) & 31
|
||||
offset = 32 - e
|
||||
|
||||
def _take(an:UOp, offset:int, count:int=0) -> UOp:
|
||||
"""an = two_over_pi_f[i+offset]"""
|
||||
if count+offset < len(two_over_pi_f) - 1:
|
||||
an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
|
||||
return an
|
||||
def _shl_lazy(x, y): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
|
||||
def _shr_lazy(x, y): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
|
||||
|
||||
a = [_take(UOp.const(dtypes.uint32, 0), i) for i in range(4)]
|
||||
# (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
|
||||
# Note: e >= 1 for all numbers d >= 1.0. assume e != 0
|
||||
hi = _shl_lazy(a[0], e) | _shr_lazy(a[1], offset)
|
||||
mi = _shl_lazy(a[1], e) | _shr_lazy(a[2], offset)
|
||||
lo = _shl_lazy(a[2], e) | _shr_lazy(a[3], offset)
|
||||
|
||||
def _hp_mul(x:UOp, y:UOp) -> UOp: return x.cast(dtypes.uint64) * y.cast(dtypes.uint64)
|
||||
# compute x * 2/pi
|
||||
p = shl(_hp_mul(ia, hi), 32) + _hp_mul(ia, mi) + shr(_hp_mul(ia, lo), 32)
|
||||
|
||||
# round quotient to nearest
|
||||
q = shr(p, 62).cast(dtypes.int32)
|
||||
p = p & 0x3fffffffffffffff
|
||||
r = (p.cast(intermediate_dtype) * (3.4061215800865545e-19)).cast(d.dtype)
|
||||
|
||||
# if fraction >= 0.5, r -= pi/2, q += 1
|
||||
return (f<0.5).where(r, r - math.pi/2), (f<0.5).where(q, q + 1)
|
||||
|
||||
def cody_waite_reduction(d:UOp) -> tuple[UOp, UOp]:
|
||||
"""
|
||||
Performs Cody-Waite Reduction: computes the reminder of `d` modulo pi/2 for the values `d` where
|
||||
0 <= abs(d) <= 39800.0
|
||||
Returns a tuple of `(r, q)`, where the output format is the same as that of `payne_hanek_reduction`.
|
||||
"""
|
||||
def _reduce_d(x:UOp, q:UOp):
|
||||
# https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefdp.c#L789-L823
|
||||
if x.dtype == dtypes.float64:
|
||||
# https://github.com/shibatch/sleef/blob/f6d8a841fbfddd26ce712834d4da220cd76048fb/src/common/misc.h#L77
|
||||
PI_A, PI_B, PI_C, PI_D = 3.1415926218032836914, 3.1786509424591713469e-08, 1.2246467864107188502e-16, 1.2736634327021899816e-24
|
||||
d = qdh * -PI_A + x
|
||||
d = q * -PI_A + d
|
||||
d = qdh * -PI_B + d
|
||||
d = q * -PI_B + d
|
||||
d = qdh * -PI_C + d
|
||||
d = q * -PI_C + d
|
||||
d = (qdh + q) * -PI_D + d
|
||||
elif x.dtype == dtypes.float16:
|
||||
# [FIXME] when reducing `d`, FP16 needs FP32 precision to achieve 1.0 ULP precision.
|
||||
d = _reduce_d(x.cast(dtypes.float32), q.cast(dtypes.float32)).cast(dtypes.float16)
|
||||
else:
|
||||
# https://github.com/shibatch/sleef/blob/4e08851f59fc2b545f9c393c6a23dfd311a26308/src/libm/sleefsp.c#L464-L503
|
||||
d = q * -3.1414794921875 + x
|
||||
d = q * -0.00011315941810607910156 + d
|
||||
d = q * -1.9841872589410058936e-09 + d
|
||||
d = q * -1.2154201256553420762e-10 + d
|
||||
return d
|
||||
|
||||
m_1_pi = 0.318309886183790671537767526745028724
|
||||
qdh = (d * (m_1_pi / 2.0**24)).cast(dtypes.int64).cast(d.dtype) * (2.0**24)
|
||||
quadrant = rintk(d * m_1_pi -qdh) if d.dtype == dtypes.float64 else rintk(d * m_1_pi)
|
||||
return _reduce_d(d, quadrant.cast(d.dtype)), quadrant.cast(dtypes.int32)
|
||||
|
||||
# *** approximate sine on small angle. ***
|
||||
def trig_poly(d:UOp, coeff32, coeff64): return d * (polyN(d*d, coeff64) if d.dtype == dtypes.float64 else polyN(d*d, coeff32))
|
||||
# approximate sine on [-pi/2, pi/2]
|
||||
def sin_poly(d:UOp) -> UOp:
|
||||
return trig_poly(d, [2.6083159809786593541503e-06, -0.0001981069071916863322258, 0.00833307858556509017944336, -0.166666597127914428710938, 1.0],
|
||||
[-7.97255955009037868891952e-18, 2.81009972710863200091251e-15, -7.64712219118158833288484e-13, 1.60590430605664501629054e-10,
|
||||
-2.50521083763502045810755e-08, 2.75573192239198747630416e-06, -0.000198412698412696162806809, 0.00833333333333332974823815,
|
||||
-0.166666666666666657414808, 1.0])
|
||||
|
||||
def _ifand(q:UOp, n:int): return (q & n).ne(0)
|
||||
|
||||
def sin_poly_small(d:UOp, q:UOp) -> UOp:
|
||||
r = sin_poly(d)
|
||||
return r * _ifand(q, 1).where(r.const_like(-1), r.const_like(1))
|
||||
|
||||
def sin_poly_large(d:UOp, q:UOp) -> UOp:
|
||||
r = sin_poly(d + _ifand(q, 1).where(d.const_like(math.pi / 2), d.const_like(0)))
|
||||
return r * _ifand(q, 2).where(r.const_like(-1), r.const_like(1))
|
||||
|
||||
# *** toplevel functions for xsin/xlog2/xexp2 ***
|
||||
|
||||
def xsin(d:UOp, fast:bool=False, switch_over:float=30.0) -> UOp:
|
||||
"""
|
||||
Implements a 1.0 ULP approximation for Ops.SIN.
|
||||
- fast=True assumes x <= switch_over.
|
||||
- switch_over is the threshold for switching to payne_hanek_reduction.
|
||||
"""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
# mask +-inf/nan as zero
|
||||
x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
|
||||
# x_sign = sign(x)
|
||||
x_sign = x.ne(0).where((x<0).where(x.const_like(-1), x.const_like(1)), x.const_like(0))
|
||||
x_abs = x * x_sign
|
||||
r, q = (cody_waite_reduction if fast else payne_hanek_reduction)(x_abs)
|
||||
if fast: result = sin_poly_small(r, q)
|
||||
else:
|
||||
# Payne Hanek Reduction assumes abs(x) >= pi/4, so for smaller values, use cody_waite_reduction.
|
||||
r_small, q_small = cody_waite_reduction(x_abs)
|
||||
result = (x_abs<switch_over).where(sin_poly_small(r_small, q_small), sin_poly_large(r, q))
|
||||
# adjusts the sign for abs(x)
|
||||
result = result * x_sign
|
||||
# sin(Inf) = NaN, sin(-Inf) = NaN, sin(NaN) = NaN
|
||||
return _lazy_map_numbers(d, d.const_like(math.nan), d.const_like(math.nan), d.const_like(math.nan), result)
|
||||
|
||||
def xexp2(d:UOp) -> UOp:
|
||||
"""
|
||||
Implements a 1.0 ULP approximation for Ops.EXP2
|
||||
- Paper: https://arxiv.org/pdf/2001.09258
|
||||
"""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
# mask +=inf/nan as zero.
|
||||
x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
|
||||
q = rintk(x)
|
||||
# s = d - round(d)
|
||||
s = x - q.cast(x.dtype)
|
||||
# a polynomial approximation with 13 non-zero terms in the range of [−(log 2)/2,(log 2)/2].
|
||||
if d.dtype == dtypes.float64:
|
||||
u = polyN(s, [0.4434359082926529454e-9, 0.7073164598085707425e-8, 0.1017819260921760451e-6, 0.1321543872511327615e-5, 0.1525273353517584730e-4,
|
||||
0.1540353045101147808e-3, 0.1333355814670499073e-2, 0.9618129107597600536e-2, 0.5550410866482046596e-1, 0.2402265069591012214e+0,
|
||||
0.6931471805599452862e+0, 0.1000000000000000000e+1])
|
||||
else: u = polyN(s, [0.1535920892e-3, 0.1339262701e-2, 0.9618384764e-2, 0.5550347269e-1, 0.2402264476e+0, 0.6931471825e+0, 1.0])
|
||||
u = ldexp2k(u, q) # u*2^q
|
||||
upper, lower = {dtypes.float64: (1024, -2000), dtypes.float32: (128, -150), dtypes.float16: (23, -22)}[d.dtype]
|
||||
# Replace x >= upper with +inf
|
||||
u = (d >= upper).where(d.const_like(math.inf), u)
|
||||
# Replace x < lower with zero.
|
||||
u = (d<lower).where(d.const_like(0.0), u)
|
||||
# exp2(NaN) = NaN
|
||||
return d.ne(d).where(d.const_like(math.nan), u)
|
||||
|
||||
def xlog2(d:UOp) -> UOp:
|
||||
"""
|
||||
Implements a 1.0 ULP approximation for Ops.LOG2
|
||||
Paper: https://arxiv.org/pdf/2001.09258 5.5
|
||||
"""
|
||||
assert d.dtype in TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
# TODO: float16 denormal need float32 to achieve precision
|
||||
if d.dtype == dtypes.float16: return xlog2(d.cast(dtypes.float32)).cast(dtypes.float16)
|
||||
FLT_MIN = d.const_like(1e-6 if d.dtype == dtypes.float16 else 1e-4)
|
||||
is_denormal = d<FLT_MIN
|
||||
a = is_denormal.where(d * (2 ** 64), d)
|
||||
|
||||
e = ilogb2k(a * (1.0 / 0.75)).cast(a.dtype)
|
||||
m = ldexp3k(a, -e)
|
||||
e = is_denormal.where(e - 64, e)
|
||||
|
||||
x = (m - 1.0) / (m + 1.0)
|
||||
x2 = x * x
|
||||
if d.dtype == dtypes.float64:
|
||||
t = polyN(x2, [0.2211941750456081490e+0, 0.2200768693152277689e+0, 0.2623708057488514656e+0, 0.3205977477944495502e+0,
|
||||
0.4121985945485324709e+0, 0.5770780162997058982e+0, 0.96179669392608091449])
|
||||
s_hi, s_lo = e+x*2.885390081777926774, e.const_like(0)
|
||||
else:
|
||||
t = polyN(x2, [0.4374550283e+0, 0.5764790177e+0, 0.9618012905120])
|
||||
s_hi, s_lo = e+x*2.8853900432586669922, x*3.2734474483568488616e-08
|
||||
r = t * (x * x2) + (s_hi + s_lo)
|
||||
|
||||
# log2(Inf) = Inf
|
||||
r = d.ne(math.inf).where(r, r.const_like(math.inf))
|
||||
# log2(x) = NaN for x < 0
|
||||
r = (d<-0.0).where(r.const_like(math.nan), r)
|
||||
# log2(0) = -Inf, but we will compare using the value of y because 1e-200==0 is true.
|
||||
# log2_zero = the value of unmasked xlog2(0.0).
|
||||
log2_zero = {dtypes.float64: -1087, dtypes.float32: -191, dtypes.float16: -79}[d.dtype]
|
||||
r = r.ne(log2_zero).where(r, r.const_like(-math.inf))
|
||||
# log2(NaN) = NaN
|
||||
r = d.ne(d).where(r.const_like(math.nan), r)
|
||||
# log2(-0.0) = -Inf. In certain devices like PTX, x == -0.0 won't be true. so making reciprocal.
|
||||
return d.reciprocal().ne(-math.inf).where(r, r.const_like(-math.inf))
|
||||
@@ -0,0 +1,513 @@
|
||||
from __future__ import annotations
|
||||
from typing import Optional, TYPE_CHECKING, Any, DefaultDict, Callable
|
||||
import functools, itertools, operator
|
||||
from collections import defaultdict
|
||||
from tinygrad.dtype import dtypes, ImageDType, PtrDType
|
||||
from tinygrad.ops import UOp, Ops, UPat, PatternMatcher, symbolic_flat, symbolic_simple
|
||||
from tinygrad.ops import graph_rewrite, split_uop, uop_given_valid, parse_valid, is_increasing, simplify_valid, GroupOp
|
||||
from tinygrad.helpers import DEBUG, getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, partition, all_same
|
||||
from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, TRANSCENDENTAL_SUPPORTED_DTYPES
|
||||
|
||||
if TYPE_CHECKING: from tinygrad.renderer import Renderer
|
||||
|
||||
# ***** float4/image store handling *****
|
||||
|
||||
def fold_expanded(ex, buf):
|
||||
if buf.dtype.base != dtypes.float and buf.dtype.base != dtypes.half and not isinstance(buf.dtype, ImageDType): return None
|
||||
new_srcs = dedup(list(ex.src))
|
||||
old_new_srcs = new_srcs[:]
|
||||
is_load, is_image = new_srcs[0].op is Ops.LOAD, isinstance(buf.dtype, ImageDType)
|
||||
|
||||
# first, extract all the relevant offsets
|
||||
offsets_rootsrc: DefaultDict[Any, dict] = defaultdict(dict)
|
||||
for i,s in enumerate(new_srcs):
|
||||
idx = s.src[0].src[1]
|
||||
if s.dtype.count != 1 or (is_image and idx.dtype.count == 2): continue
|
||||
if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
|
||||
elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
|
||||
else: root_src, arg = idx, 0
|
||||
# add gates for gated
|
||||
if len(s.src[0].src) == 3: root_src = (s.src[0].src[2], root_src)
|
||||
assert arg not in offsets_rootsrc[root_src], f"{offsets_rootsrc[root_src][arg]} != {i} with {len(s.src)} sources"
|
||||
offsets_rootsrc[root_src][arg] = i
|
||||
|
||||
# then rewrite everything we can
|
||||
lengths = [4] if is_image else ([8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else ([16,8,4,2] if AMX else [4,2]))
|
||||
used: set[tuple[UOp, UOp]] = set()
|
||||
for rootsrc, offsets in offsets_rootsrc.items():
|
||||
for o in offsets:
|
||||
for fold_length in lengths:
|
||||
if all((rootsrc,o+i) not in used and o+i in offsets for i in range(fold_length)):
|
||||
load_1 = new_srcs[offsets[o]]
|
||||
new_src = list(load_1.src)
|
||||
oidx = new_src[0].src[1]
|
||||
if oidx.divides(fold_length) is None: continue
|
||||
if is_image:
|
||||
# for images, we rewrite the index. it must evenly divide 4 from the above check
|
||||
new_src[0] = buf.index(
|
||||
UOp(Ops.VECTORIZE, dtypes.int.vec(2), ((oidx // 4) % buf.dtype.shape[1], (oidx // (4*buf.dtype.shape[1])))),
|
||||
rootsrc[0] if isinstance(rootsrc, tuple) else None)
|
||||
else:
|
||||
# for non image, we upcast the index pointer
|
||||
new_src[0] = new_src[0].cast(new_src[0].dtype.base.vec(fold_length).ptr(size=new_src[0].dtype.size//fold_length,
|
||||
local=new_src[0].dtype.local))
|
||||
# generate the folded new_srcs
|
||||
if is_load:
|
||||
new_load = UOp(Ops.LOAD, load_1.dtype.vec(fold_length), tuple(new_src))
|
||||
for i in range(fold_length): new_srcs[offsets[o+i]] = new_load.gep(i)
|
||||
else: # vectorize the store
|
||||
new_src[1] = UOp(Ops.VECTORIZE, new_src[1].dtype.vec(fold_length), tuple(new_srcs[offsets[o+i]].src[1] for i in range(fold_length)))
|
||||
for i in range(fold_length): new_srcs[offsets[o+i]] = UOp(Ops.STORE, dtypes.void, tuple(new_src)) if i == 0 else None
|
||||
used.update((rootsrc,o+i) for i in range(fold_length))
|
||||
|
||||
# dedup expand for LOAD
|
||||
if is_load and len(old_new_srcs) != len(ex.src): new_srcs = [new_srcs[old_new_srcs.index(s)] for s in ex.src]
|
||||
# remove Nones for STORE
|
||||
return UOp(ex.op, ex.dtype, tuple(x for x in new_srcs if x is not None), ex.arg) if len(used) else None
|
||||
|
||||
def fix_unfoldable_image_load(load:UOp, buf:UOp):
|
||||
if not isinstance(buf.dtype, ImageDType) or (oidx:=load.src[0].src[1]).dtype.count == 2: return None
|
||||
id4 = oidx % 4
|
||||
new_src = list(load.src)
|
||||
# TODO: copied logic from above
|
||||
new_src[0] = load.src[0].src[0].index(
|
||||
UOp(Ops.VECTORIZE, dtypes.int.vec(2), ((oidx // 4) % buf.dtype.shape[1], (oidx // (4*buf.dtype.shape[1])))),
|
||||
load.src[0].src[2] if len(load.src[0].src) == 3 else None)
|
||||
vec_load = UOp(Ops.LOAD, load.dtype.vec(4), tuple(new_src))
|
||||
return functools.reduce(lambda ret, i: id4.ne(i).where(ret, vec_load.gep(i)), range(4), load.const_like(float('nan')))
|
||||
|
||||
buf_idx_pat = UPat(Ops.INDEX, src=(UPat.var("buf"),), allow_any_len=True)
|
||||
float4_folding = PatternMatcher([
|
||||
(UPat(Ops.VECTORIZE, src=UPat(Ops.LOAD, src=(buf_idx_pat,), allow_any_len=True), name="ex"), fold_expanded),
|
||||
(UPat((Ops.BARRIER, Ops.SINK), src=UPat(Ops.STORE, src=(buf_idx_pat,), allow_any_len=True), name="ex"), fold_expanded),
|
||||
])
|
||||
|
||||
# ***** image load valid simplification *****
|
||||
|
||||
def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
|
||||
if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.const_like(0)
|
||||
if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx, valid)
|
||||
|
||||
# wait for it to be image indexed before running simplification
|
||||
if start_idx.dtype.count != 2: return None
|
||||
|
||||
# can drop valid if idx is out of bound when valid is False
|
||||
drop_stmt = []
|
||||
for stmt in split_uop(valid, Ops.AND):
|
||||
X, is_upper_bound, c = parse_valid(stmt)
|
||||
|
||||
# for X0 + X1 + ... >= 1, check if it's out of bound when Xi = 0 for all i
|
||||
if not is_upper_bound and c == 1 and all(u.op in GroupOp.Irreducible and u.vmin == 0 for u in split_uop(X, Ops.ADD)):
|
||||
testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), split_uop(X, Ops.ADD), idx)
|
||||
testidx = testidx.simplify()
|
||||
if testidx.gep(0).vmax < 0 or testidx.gep(1).vmax < 0:
|
||||
drop_stmt.append(stmt)
|
||||
continue
|
||||
|
||||
# if X <= c, check if it's out of bound when X = c+1
|
||||
# if X >= c, check if it's out of bound when X = c-1
|
||||
test_value = c + 1 if is_upper_bound else c - 1
|
||||
for i,b in zip(idx.src, (buf.dtype.shape[1], buf.dtype.shape[0])):
|
||||
if is_increasing(i):
|
||||
rw = i.substitute({X:X.const_like(test_value)}).simplify()
|
||||
if rw.vmin >= b or rw.vmax < 0:
|
||||
drop_stmt.append(stmt)
|
||||
break
|
||||
|
||||
if not drop_stmt and idx is start_idx: return None
|
||||
new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in split_uop(valid, Ops.AND) if s not in drop_stmt]) else None
|
||||
return buf.index(idx, new_valid)
|
||||
|
||||
# ***** optional patterns *****
|
||||
|
||||
powers_of_two = {2**i:i for i in range(64)}
|
||||
@functools.lru_cache(None)
|
||||
def get_late_rewrite_patterns(ops, force_transcendental=False):
|
||||
pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
|
||||
((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
|
||||
# rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
|
||||
if Ops.AND in ops:
|
||||
pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
|
||||
# rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
|
||||
if Ops.SHL in ops and Ops.SHR in ops:
|
||||
pat += [
|
||||
(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << powers_of_two[c.arg] if c.arg in powers_of_two else None),
|
||||
(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: x >> powers_of_two[c.arg] if c.arg in powers_of_two else None)
|
||||
]
|
||||
if Ops.NEG in ops:
|
||||
pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
|
||||
if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
|
||||
if Ops.MULACC in ops:
|
||||
pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
|
||||
return PatternMatcher(pat)
|
||||
|
||||
# ***** threefry *****
|
||||
|
||||
def threefry2x32(x: UOp, key: UOp):
|
||||
# split x into two uint32, since x in a uint64
|
||||
x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
|
||||
|
||||
rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
|
||||
key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
|
||||
ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
|
||||
xr = [x0 + ks[-1], x1 + ks[0]]
|
||||
for i in range(5):
|
||||
for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
|
||||
xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
|
||||
|
||||
return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
|
||||
|
||||
# ***** other math rewrite ****
|
||||
|
||||
def sigmoid_like(x:UOp, y:UOp): return (t:=(1/(x+1))) * (1-t) * y
|
||||
|
||||
# ***** main rewriter *****
|
||||
|
||||
def loop_collapse(compval, multconst, rng:UOp, acc:UOp, idx2=None,idx3=None,extra=None,vec=None,ne=None,
|
||||
add=UOp.const(dtypes.int, 0), mul:UOp=UOp.const(dtypes.int, 1)):
|
||||
if getenv("DISABLE_LOOP_COLLAPSE") or rng not in acc.src: return None # must be the right REDUCE
|
||||
loop_start, loop_end = rng.src
|
||||
if loop_start.arg != 0:
|
||||
# TODO: support and test this with other mul and loop_starts
|
||||
if DEBUG >= 1: print(f"WARNING, NOT FOLDING: mul:{mul.arg} loop_start:{loop_start.arg}")
|
||||
return None
|
||||
if idx2 is not None: add = add + idx2
|
||||
if idx3 is not None: add = add + idx3
|
||||
if vec is not None:
|
||||
# add, mul, loop_start, loop_end
|
||||
def dvec(x:UOp):
|
||||
if x.op is Ops.CONST: return UOp.const(x.dtype.vec(vec.dtype.count), x.arg)
|
||||
return UOp(Ops.VECTORIZE, x.dtype.vec(vec.dtype.count), src=(x,)*vec.dtype.count)
|
||||
add, mul, loop_start, loop_end = dvec(add), dvec(mul), dvec(loop_start), dvec(loop_end)
|
||||
if mul.vmin > 0 and ne is not None:
|
||||
comprange = UOp.minimum(loop_end, UOp.maximum((add-compval)//mul + (loop_end-loop_start), loop_start))
|
||||
elif mul.vmax < 0 and ne is None:
|
||||
comprange = UOp.minimum(loop_end, UOp.maximum((add-compval-mul)//mul + (loop_end-loop_start), loop_start))
|
||||
else:
|
||||
return None
|
||||
new_reduce_op = comprange.cast(multconst.dtype) * multconst
|
||||
# TODO: what does it mean to have the same numbered DEFINE_ACC with different ranges?
|
||||
new_acc = acc.replace(src=acc.src[0:1]+tuple(x for x in acc.src[1:] if x is not rng))
|
||||
ret = new_acc.assign(new_acc+new_reduce_op)
|
||||
if extra is not None: ret = ret + acc.assign(acc+extra)
|
||||
return ret
|
||||
|
||||
def index_collapse(idx:UOp,rng:UOp,buf:UOp,ld:UOp,acc:UOp,add=UOp.const(dtypes.int, 0),mul=UOp.const(dtypes.int, 1)):
|
||||
if rng not in acc.src: return None
|
||||
new_load = UOp.load(buf.index(add+mul*idx, (idx >= rng.src[0]) & (idx < rng.src[1])), dtype=ld.dtype)
|
||||
new_acc = acc.replace(src=acc.src[0:1]+tuple(x for x in acc.src[1:] if x is not rng))
|
||||
return new_acc.assign(new_acc+new_load)
|
||||
|
||||
# TODO: there's a lot shared with no_vectorized_wmma here
|
||||
def gep_through_wmma(gep:UOp, wmma:UOp):
|
||||
out_sz = prod(x[1] for x in wmma.arg[6][-1])
|
||||
wmma_idxs = gep.arg[::out_sz]
|
||||
for i in range(out_sz):
|
||||
if tuple(x-i for x in gep.arg[i::out_sz]) != wmma_idxs: return None
|
||||
tsrcs = []
|
||||
for s,sz in zip(wmma.src, wmma.arg[6]):
|
||||
src_args = []
|
||||
ssz = prod(x[1] for x in sz)
|
||||
for w in wmma_idxs: src_args += list(range((w//out_sz)*ssz, (w//out_sz)*ssz + ssz))
|
||||
tsrcs.append(s.gep(tuple(src_args)))
|
||||
return UOp(Ops.WMMA, gep.dtype, tuple(tsrcs), wmma.arg)
|
||||
|
||||
def no_vectorized_wmma(wmma:UOp):
|
||||
out_sz = prod(x[1] for x in wmma.arg[6][-1])
|
||||
if wmma.dtype.count == out_sz: return None
|
||||
tsrcs = []
|
||||
for s,sz in zip(wmma.src, wmma.arg[6]):
|
||||
ssz = prod(x[1] for x in sz)
|
||||
tsrcs.append([s.gep(tuple(range(grp, grp+ssz))) for grp in range(0, s.dtype.count, ssz)])
|
||||
wmmas = [UOp(Ops.WMMA, wmma.dtype.scalar().vec(out_sz), tsrc, wmma.arg) for tsrc in zip(*tsrcs)]
|
||||
wmma_ex = flatten([[e.gep(i) for i in range(out_sz)] for e in wmmas])
|
||||
return UOp(Ops.VECTORIZE, wmma.dtype, tuple(wmma_ex))
|
||||
|
||||
def reduce_collapse(acc:UOp, ret:UOp, alu:UOp):
|
||||
reduce_parented, reduce_unparented = partition(acc.src[1:], lambda x: x in ret.toposort)
|
||||
if len(reduce_unparented) == 0: return None
|
||||
new_acc = acc.replace(src=acc.src[0:1]+tuple(reduce_parented))
|
||||
ret = new_acc.assign(new_acc.alu(alu.op, ret))
|
||||
if alu.op is Ops.ADD:
|
||||
for r in reduce_unparented: ret = ret * (r.src[1]-r.src[0]).cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
|
||||
return ret
|
||||
|
||||
acc_pat, rng_pat = UPat(Ops.DEFINE_ACC, name="acc"), UPat(Ops.RANGE, name="rng")
|
||||
rng_aug = UPat.any(rng_pat, UPat.var("add")+rng_pat, UPat.var("mul")*rng_pat, UPat.var("add")+UPat.var("mul")*rng_pat)
|
||||
|
||||
index_load = UPat.var("buf").index(rng_aug).load(name="ld")
|
||||
|
||||
arange_augrng = UPat.any(rng_aug, rng_aug+UPat.var("idx2"), rng_aug+UPat.var("idx2")+UPat.var("idx3"), UPat(Ops.VECTORIZE, name="vec", src=rng_aug))
|
||||
arange_m = ((arange_augrng<UPat.cvar("compval"))!=UPat(Ops.CONST, name="ne", arg=True)).where(UPat.cvar("multconst"), UPat.const(None, 0))
|
||||
|
||||
# this is symbolic 2.0
|
||||
sym = symbolic_flat+PatternMatcher([
|
||||
# self ASSIGN is just self
|
||||
(UPat(Ops.ASSIGN, src=(UPat.var('x'), UPat.var('x'))), lambda x: x),
|
||||
# VECTORIZE/CONST, VECTORIZE/GEP
|
||||
(UPat(Ops.VECTORIZE, src=UPat(Ops.CONST), name="vec"), lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src))),
|
||||
(UPat(Ops.VECTORIZE, src=UPat(Ops.GEP, src=(UPat(name="x"),)), name="vec"), lambda vec,x: x.gep(tuple(y.arg[0] for y in vec.src))),
|
||||
# reorder ALU/VECTORIZE
|
||||
(UPat(GroupOp.ALU, src=(UPat(Ops.VECTORIZE, src=UPat(name='x')), UPat(Ops.VECTORIZE, src=UPat(name='y'))), name='alu'),
|
||||
lambda x,y,alu: UOp(Ops.VECTORIZE, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)),
|
||||
# VECTORIZE of a single element is just that element
|
||||
(UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
|
||||
# VECTORIZE void is SINK
|
||||
(UPat(Ops.VECTORIZE, dtype=dtypes.void, src=UPat(Ops.BARRIER, name='b')), lambda b: b),
|
||||
(UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp(Ops.SINK, dtypes.void, x.src)),
|
||||
# GEP/VECTORIZE, GEP/GEP, GEP/CONST, GEP/VCONST
|
||||
(UPat(Ops.GEP, src=(UPat(Ops.GEP, name='g2'),), name='g1'),
|
||||
lambda g1, g2: g2.src[0].gep(tuple(g2.arg[g1.arg[i]] for i in range(g1.dtype.count)))),
|
||||
(UPat(Ops.GEP, src=(UPat(Ops.VECTORIZE, name="vec"),), name="gep"),
|
||||
lambda gep, vec: UOp(Ops.VECTORIZE, gep.dtype, tuple(vec.src[i] for i in gep.arg)) if len(gep.arg) > 1 else vec.src[gep.arg[0]]),
|
||||
(UPat(Ops.GEP, src=(UPat.cvar("c", vec=False),), name="gep"), lambda gep, c: gep.const_like(c.arg)),
|
||||
(UPat(Ops.GEP, src=(UPat(Ops.VCONST, name="c"),), name="gep"), lambda gep, c: gep.const_like(tuple(c.arg[x] for x in gep.arg))),
|
||||
# push all GEPs through ALUs (fix arange stuff)
|
||||
(UPat(Ops.GEP, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu'),), name='gep'),
|
||||
lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg)),
|
||||
# push some GEPs through WMMAs
|
||||
(UPat(Ops.GEP, src=(UPat(Ops.WMMA, name="wmma"),), name="gep"), gep_through_wmma),
|
||||
# tensor core with a 0 input is acc
|
||||
(UPat(Ops.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc),
|
||||
(UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc),
|
||||
# tensor core cleanups
|
||||
(UPat.var("add") + UPat(Ops.WMMA, name="wmma"),
|
||||
lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
|
||||
# threefry + remove longs
|
||||
(UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32),
|
||||
(UPat.var('x', dtypes.uint32).cast(dtypes.uint64).cast(dtypes.uint32), lambda x: x), # cast there and back is noop (TODO: genericize)
|
||||
((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)), # cast does truncation
|
||||
(((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
|
||||
(((UPat.var('x', dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
|
||||
# hacks for threefry long removal when padded (TODO: genericize)
|
||||
(UPat.var('x', dtypes.uint32).cast(dtypes.uint64) * UPat.var('y').where(UPat.const(dtypes.uint64, 1<<32), UPat.const(dtypes.uint64, 0)),
|
||||
lambda x,y: y.where(x, UOp.const(dtypes.uint32, 0)).cast(dtypes.uint64) * (1<<32)),
|
||||
((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32),
|
||||
lambda x,y: y.where(x.cast(dtypes.uint32), UOp.const(dtypes.uint32, 0))),
|
||||
# arange loop folding
|
||||
(acc_pat.assign(UPat.any(arange_m, arange_m+UPat.var("extra"))+acc_pat), loop_collapse),
|
||||
# indexing, with cast or where
|
||||
(acc_pat.assign(UPat.var("idx").eq(UPat(Ops.RANGE, name="rng")).cast()*index_load+acc_pat), index_collapse),
|
||||
(acc_pat.assign(UPat.var("idx").eq(UPat(Ops.RANGE, name="rng")).where(index_load, UPat.const(None, 0.0))+acc_pat), index_collapse),
|
||||
# parentless reduce # TODO: add MUL
|
||||
(acc_pat.assign(UPat((Ops.ADD, Ops.MAX), src=[acc_pat, UPat.var("ret")], name="alu")), reduce_collapse),
|
||||
# ** self folding **
|
||||
(UPat(Ops.DEFINE_ACC, src=(UPat.var("x"),)), lambda x: x), # a DEFINE_ACC without ranges is a CONST
|
||||
(UPat(Ops.ASSIGN, src=(UPat.cvar(),UPat.var("x"))), lambda x: x), # an ASSIGN to a const is a NOOP
|
||||
# x!=0 -> (bool)x
|
||||
(UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
|
||||
# ** load/store folding **
|
||||
(UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)),
|
||||
(UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"), UPat.load(UPat(Ops.INDEX, name="index")))),
|
||||
lambda index, gate, alt: UOp.store(index.src[0].index(index.src[1], gate), alt)),
|
||||
# fold gated LOAD/STORE
|
||||
(UPat().index(UPat(), UPat.const(dtypes.bool, True)).named("idx"), lambda idx: idx.replace(src=idx.src[0:2])), # remove True
|
||||
(UPat().index(UPat(), UPat.const(dtypes.bool, False)).named("idx"), lambda idx: idx.const_like(0)), # False -> NULL pointer
|
||||
(UPat(Ops.LOAD, src=(UPat.const(None, 0),), allow_any_len=True, name="x"), lambda x: x.const_like(0)), # NULL pointer load loads 0
|
||||
(UPat(Ops.STORE, src=(UPat.const(None, 0),), allow_any_len=True), lambda: UOp(Ops.NOOP)), # NULL pointer store does nothing
|
||||
# remove NOOPs from SINK
|
||||
(UPat(Ops.SINK, name="root"),
|
||||
lambda root: UOp(Ops.SINK, root.dtype, a, root.arg) if len(a:=tuple(x for x in root.src if x.op is not Ops.NOOP)) != len(root.src) else None),
|
||||
# remove VECTORIZE from SINK/BARRIER
|
||||
(UPat(Ops.BARRIER, src=(UPat((Ops.VECTORIZE, Ops.SINK), name='sink'),)), lambda sink: UOp(Ops.BARRIER, dtypes.void, sink.src)),
|
||||
(UPat(Ops.SINK, name="root"),
|
||||
lambda root: UOp(Ops.SINK, root.dtype, tuple(flatten(x.src if x.op in {Ops.SINK, Ops.UNROLL} else (x,) for x in root.src)), root.arg)
|
||||
if any(x.op in {Ops.SINK, Ops.UNROLL} for x in root.src) else None),
|
||||
# stable sigmoid
|
||||
(UPat.var("x")*(((UPat.var("x")+1)*(UPat.var("x")+1)).reciprocal()), lambda x: sigmoid_like(x, x.const_like(1))),
|
||||
(UPat.var("x")*(((UPat.var("x")+1)*(UPat.var("x")+1)).reciprocal()*UPat.var("y")), sigmoid_like),
|
||||
(UPat.var("x")*(((UPat.var("x")+1)*(UPat.var("x")+1)*(UPat.var("x")+1)).reciprocal()), lambda x: sigmoid_like(x, (x+1).reciprocal())),
|
||||
])
|
||||
|
||||
# *** uop expander ***
|
||||
|
||||
def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int:
|
||||
idx, mul = 0, 1
|
||||
for axis,m in args[::-1]:
|
||||
idx += rpk[axis] * mul
|
||||
mul *= m
|
||||
return idx
|
||||
|
||||
def _choices_from_args(args:tuple[tuple[int, int], ...]) -> list[dict[int, int]]:
|
||||
return [dict(x) for x in itertools.product(*[zip(itertools.repeat(axis), range(m)) for axis,m in args])]
|
||||
|
||||
@functools.lru_cache(None)
|
||||
def _swizzle_args(cargs:tuple[tuple[int, int], ...], eargs:tuple[tuple[int, int], ...], exclude_args:tuple[int, ...]) -> list[int]:
|
||||
return [_expand_arg_to_idx(eargs, {**rpk, **{x:0 for x in exclude_args}} if exclude_args else rpk) for rpk in _choices_from_args(cargs)]
|
||||
|
||||
def do_expand(root:UOp):
|
||||
expands = [x for x in root.src if x.op is Ops.UNROLL]
|
||||
if len(expands) == 0: return None
|
||||
# NOTE: we 0 out the reduce axis for WMMA. in theory they should all be the same, but is this always correct?
|
||||
exclude_args = tuple(dedup(root.arg[-1] + tuple(y[0] for y in flatten(root.arg[-2])))) if root.op is Ops.WMMA else ()
|
||||
if all_same(expands_args:=[x.arg for x in expands]) and len(exclude_args) == 0:
|
||||
# if there's only one expand arg, it's okay to use it (optimization)
|
||||
expand_args = expands[0].arg
|
||||
else:
|
||||
# otherwise, we sort them and GEP
|
||||
expand_args = tuple(x for x in sorted(dedup(flatten(expands_args))) if x[0] not in exclude_args)
|
||||
expand_sz = prod([x[1] for x in expand_args])
|
||||
new_srcs = []
|
||||
for i,src in enumerate(root.src):
|
||||
if src.op is Ops.UNROLL:
|
||||
if root.op is Ops.IF and i == 0:
|
||||
# IF means OR on first arg to IF
|
||||
new_srcs.append(functools.reduce(operator.__or__, [src.src[0].gep(i) for i in range(expand_sz)]))
|
||||
elif expand_args == src.arg:
|
||||
# just remove the expand
|
||||
new_srcs.append(src.src[0])
|
||||
else:
|
||||
lst = _swizzle_args(expand_args, src.arg, exclude_args)
|
||||
# if the base dtype is > 1, put those at the end
|
||||
if src.dtype.count > 1: lst = flatten([[i*src.dtype.count+j for j in range(src.dtype.count)] for i in lst])
|
||||
new_srcs.append(src.src[0].gep(tuple(lst)))
|
||||
else:
|
||||
# non-UNROLL input
|
||||
if root.op is Ops.IF:
|
||||
# for the first arg of IF, just pass them through ignoring UNROLLS
|
||||
new_srcs.append(src)
|
||||
elif src.dtype.count > 1:
|
||||
# put any input dtype > 1 grouped together
|
||||
new_srcs.append(UOp(Ops.VECTORIZE,
|
||||
src.dtype.scalar().vec(expand_sz*src.dtype.count), tuple(src.gep(i) for i in range(src.dtype.count))*expand_sz))
|
||||
else:
|
||||
# repeat the arg
|
||||
new_srcs.append(src.broadcast(expand_sz))
|
||||
|
||||
new_arg = root.arg
|
||||
if root.op is Ops.GEP:
|
||||
assert root.dtype.count == 1
|
||||
# is this right?
|
||||
new_arg = tuple(range(root.arg[0], new_srcs[0].dtype.count, new_srcs[0].dtype.count // expand_sz))
|
||||
nsrc = UOp(root.op, root.dtype.scalar().vec(root.dtype.count*expand_sz), tuple(new_srcs), new_arg)
|
||||
return UOp(Ops.UNROLL, root.dtype, (nsrc,), expand_args)
|
||||
|
||||
def do_contract(con:UOp):
|
||||
ex = con.src[0]
|
||||
# CONTRACT without UNROLL repeats the element VECTORIZED
|
||||
if ex.op is not Ops.UNROLL: return UOp(Ops.VECTORIZE, con.dtype, con.src*con.dtype.count)
|
||||
# CONTRACT may remove several axes from UNROLL
|
||||
assert con.dtype.count == prod([x[1] for x in con.arg]), "dtype is wrong"
|
||||
idxs = []
|
||||
for rpk in _choices_from_args(new_ex_args:=tuple(x for x in ex.arg if x not in con.arg)):
|
||||
idxs += [_expand_arg_to_idx(ex.arg, {**rpk, **lrpk}) for lrpk in _choices_from_args(con.arg)]
|
||||
return UOp(Ops.UNROLL, con.dtype, (ex.src[0].gep(tuple(idxs)),), new_ex_args)
|
||||
|
||||
def no_vectorized_alu(alu):
|
||||
if alu.dtype.vcount == 1: return None
|
||||
alus = tuple(UOp(alu.op, alu.dtype.scalar(), tuple(s.gep(i) for s in alu.src), alu.arg) for i in range(alu.dtype.vcount))
|
||||
return UOp(Ops.VECTORIZE, alu.dtype, alus)
|
||||
|
||||
def create_gate(root:UOp) -> UOp|None:
|
||||
@functools.lru_cache(None)
|
||||
def _gate_srcs(u:UOp, gate:UOp) -> UOp:
|
||||
if u.op is Ops.BARRIER: return u
|
||||
if u.op is Ops.LOAD and u.src[-1].op is Ops.BARRIER:
|
||||
return UOp(u.op, u.dtype, u.src[:-1]+(UOp(Ops.IF, dtypes.void, (gate, u.src[-1])),), u.arg)
|
||||
return u if (replace_source:=tuple(_gate_srcs(x, gate) for x in u.src)) == u.src else UOp(u.op, u.dtype, replace_source, u.arg)
|
||||
idx = root.src[0]
|
||||
if idx.op is Ops.CAST: idx = idx.src[0]
|
||||
return None if idx.op is not Ops.INDEX or len(idx.src) == 2 or (ret:=_gate_srcs(root, idx.src[2])) is root else ret
|
||||
|
||||
expander = PatternMatcher([
|
||||
# double expand
|
||||
(UPat(Ops.UNROLL, name="outer", src=(UPat(Ops.UNROLL, name="inner"),)),
|
||||
lambda outer, inner: UOp(Ops.UNROLL, outer.dtype, (inner.src[0],), inner.arg+outer.arg)),
|
||||
# do expansion
|
||||
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.GEP, Ops.WMMA, Ops.LOAD, Ops.STORE, Ops.INDEX, Ops.ASSIGN,
|
||||
Ops.VECTORIZE, Ops.IF), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
|
||||
(UPat(Ops.CONTRACT, name="con"), do_contract),
|
||||
# vectorize DEFINE_ACC
|
||||
(UPat(Ops.VECTORIZE, src=UPat(Ops.DEFINE_ACC, name="acc"), name="v"), lambda acc,v: acc.replace(dtype=v.dtype)),
|
||||
# BARRIERs aren't actually expanded
|
||||
(UPat(Ops.BARRIER, src=(UPat(Ops.UNROLL, name="ex"),)),
|
||||
lambda ex: UOp(Ops.UNROLL, dtypes.void, (UOp(Ops.BARRIER, dtypes.void, ex.src),)*len(ex.src), ex.arg)),
|
||||
# empty UNROLL is NOOP
|
||||
(UPat(Ops.UNROLL, src=(UPat.var('x'),), arg=()), lambda x: x),
|
||||
# UNROLL GEP (needed for WMMA, generalize this) -> vectorized ALU
|
||||
(UPat(Ops.UNROLL, name="ex", src=tuple(UPat.var('x').gep(i)+UPat.var('y').gep(i) for i in range(256 if AMX else 8))),
|
||||
lambda ex,x,y: UOp(Ops.UNROLL, ex.dtype, tuple((x+y).gep(i) for i in range(256 if AMX else 8)), ex.arg)),
|
||||
])
|
||||
|
||||
def no_vectorized_load_store(ls:UOp):
|
||||
idx = ls.src[0]
|
||||
assert isinstance(idx.dtype, PtrDType)
|
||||
if idx.dtype.v == 1: return None
|
||||
tv = [UOp(ls.op, ls.dtype.scalar(), tuple(j.gep(i) for j in ls.src)) for i in range(idx.dtype.v)]
|
||||
return UOp(Ops.VECTORIZE, ls.dtype, tuple(tv))
|
||||
|
||||
def no_vectorized_acc(acc:UOp):
|
||||
if acc.dtype.count == 1: return None
|
||||
alus = tuple(UOp(acc.op, acc.dtype.scalar(),
|
||||
tuple(s.gep(i) if j == 0 else s for j,s in enumerate(acc.src)), acc.arg+(i,)) for i in range(acc.dtype.count))
|
||||
return UOp(Ops.VECTORIZE, acc.dtype, alus)
|
||||
|
||||
devectorize = PatternMatcher([
|
||||
# no ALU on vectorized dtypes
|
||||
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.ASSIGN, Ops.INDEX), name="alu"), no_vectorized_alu),
|
||||
(UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
|
||||
(UPat(Ops.DEFINE_ACC, name="acc"), no_vectorized_acc),
|
||||
(UPat((Ops.LOAD, Ops.STORE), name="ls"), no_vectorized_load_store),
|
||||
])
|
||||
|
||||
def delete_redundant_gates(buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
|
||||
if store_gate not in [gate.src[0] for gate in val.toposort if gate.op is Ops.IF]: return None
|
||||
# remove the gate from the index
|
||||
return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val)
|
||||
|
||||
load_store_indexing = PatternMatcher([
|
||||
# late fixup of unfoldable image loads
|
||||
(UPat(Ops.LOAD, src=(UPat.var("buf"), UPat()), allow_any_len=True, name="load"), fix_unfoldable_image_load),
|
||||
# simplify valid
|
||||
(UPat(Ops.AND, name="valid"), simplify_valid),
|
||||
# image load valid idx simplification
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat.var("valid"))), simplify_valid_load),
|
||||
# delete_redundant_gates (after expand)
|
||||
(UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
|
||||
UPat.var("val"))), delete_redundant_gates),
|
||||
])
|
||||
|
||||
migrate_indexing = PatternMatcher([
|
||||
# create gate MUST BE BEFORE expander
|
||||
(UPat(Ops.STORE, name="root"), create_gate),
|
||||
])
|
||||
|
||||
def move_mask(x:UOp, buf:UOp, idx:UOp, mask:UOp, cast:UOp|None=None) -> UOp:
|
||||
# this moves the mask from the indexing to the load/store op for rendering
|
||||
nidx = buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx)
|
||||
return UOp.load(nidx, x.const_like(0), mask, *x.src[1:], dtype=x.dtype) if x.op is Ops.LOAD else UOp.store(nidx, x.src[1], mask, *x.src[2:])
|
||||
|
||||
pm_render = PatternMatcher([
|
||||
# for rendering, we use explicit VECTORIZE
|
||||
(UPat(Ops.CONST, name='c'),
|
||||
lambda c: UOp(Ops.VECTORIZE, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None),
|
||||
(UPat(Ops.VCONST, name='c'), lambda c: UOp(Ops.VECTORIZE, c.dtype, tuple(UOp.const(c.dtype.scalar(), x) for x in c.arg))),
|
||||
(UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.VECTORIZE, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
|
||||
(UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
|
||||
# move masks of loads/stores
|
||||
(UPat((Ops.LOAD, Ops.STORE), src=(UPat.any(masked_index:=UPat(Ops.INDEX, src=(UPat(name="buf"), UPat(name="idx"), UPat(name="mask"))),
|
||||
masked_index.cast(None).named("cast")),), allow_any_len=True, name="x"), move_mask),
|
||||
# gate any stores that aren't gated with ifs
|
||||
(UPat(Ops.STORE, dtype=dtypes.void, src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="store"),
|
||||
lambda store: UOp(Ops.STORE, src=store.src[:2]+(UOp(Ops.IF, src=(store.src[2],)),))),
|
||||
])
|
||||
|
||||
# *** uop graph ***
|
||||
|
||||
def full_graph_rewrite(sink:UOp, opts:Optional[Renderer]=None) -> UOp:
|
||||
assert sink.op is Ops.SINK, f"sink isn't sink, it's {sink.op}"
|
||||
supported_ops = tuple(opts.code_for_op.keys()) if opts is not None else ()
|
||||
extra_matcher = opts.extra_matcher if opts is not None and opts.extra_matcher is not None else PatternMatcher([])
|
||||
|
||||
# initial symbolic + migrate indexing (remove this)
|
||||
sink = graph_rewrite(sink, sym+migrate_indexing)
|
||||
|
||||
# expand
|
||||
sink = graph_rewrite(sink, sym+expander)
|
||||
|
||||
# devectorize + load_store_indexing
|
||||
sink = graph_rewrite(sink, sym+(devectorize+float4_folding if opts is not None and opts.supports_float4 else devectorize)+load_store_indexing)
|
||||
|
||||
# final rules for the renderer (without sym)
|
||||
sink = graph_rewrite(sink, symbolic_simple+get_late_rewrite_patterns(supported_ops, TRANSCENDENTAL>=2)+pm_render+extra_matcher)
|
||||
return sink
|
||||
Reference in New Issue
Block a user