openpilot v0.11.1 release
date: 2026-06-04T09:49:56 master commit: c0ab3550eca2e9daf197c46b7e4b24aa9637cf2e
This commit is contained in:
Vehicle Researcher
@@ -0,0 +1,208 @@
|
||||
from typing import cast
|
||||
from dataclasses import replace
|
||||
import itertools
|
||||
from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
|
||||
from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
|
||||
from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo
|
||||
from tinygrad.uop.render import pyrender
|
||||
from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
|
||||
from tinygrad.renderer import Renderer, Estimates
|
||||
from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
# import all pattern matchers here
|
||||
from tinygrad.codegen.gpudims import pm_add_gpudims
|
||||
from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
|
||||
from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps
|
||||
from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
|
||||
from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize_buf_and_index, devectorize_alu, pm_reduce, \
|
||||
ReduceContext, correct_load_store, pm_render, pm_add_loads, pm_make_images
|
||||
from tinygrad.codegen.opt.postrange import apply_opts
|
||||
from tinygrad.codegen.late.gater import pm_move_gates_from_index
|
||||
from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
|
||||
from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges
|
||||
from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
|
||||
from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
|
||||
|
||||
def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
|
||||
if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
|
||||
if DEBUG >= 5: print(pyrender(ast))
|
||||
if SPEC: type_verify(ast, spec_tensor)
|
||||
|
||||
# preprocess
|
||||
sink = graph_rewrite(ast, pm_mops+pm_syntactic_sugar+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True)
|
||||
|
||||
# first we optimize
|
||||
if optimize:
|
||||
# collapse loads reduce (indexing by a tensor)
|
||||
sink = graph_rewrite(sink, pm_load_collapse, name="load collapse")
|
||||
|
||||
# split ranges
|
||||
sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
|
||||
|
||||
# symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
|
||||
sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
|
||||
|
||||
# optimize (schedule) the AST
|
||||
sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
|
||||
|
||||
# do postrange optimization, BEAM or hand_coded_optimizations
|
||||
sink = apply_opts(sink, ren, beam=ast.arg.beam)
|
||||
|
||||
# ** expander (expand_rewrite) **
|
||||
sink = graph_rewrite(sink, sym+pm_move_where_on_load, name="postopt symbolic")
|
||||
|
||||
# expand
|
||||
sink = graph_rewrite(sink, sym+pm_pre_expander+pm_group_for_reduce+expander, name="expander")
|
||||
|
||||
# add locals
|
||||
sink = graph_rewrite(sink, pm_add_buffers_local+rangeify_codegen, ctx=itertools.count(0), name="add local buffers")
|
||||
|
||||
# ** devectorizer (full_graph_rewrite) **
|
||||
# remove reduce
|
||||
sink = graph_rewrite(sink, pm_reduce+gep_pushing, ctx=ReduceContext(), name="remove_reduce")
|
||||
|
||||
# add gpu dims (late). this works after devectorize, but it's faster here
|
||||
sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
|
||||
|
||||
# **** optimizations are done, now we lower to actual code ****
|
||||
|
||||
# add loads
|
||||
sink = graph_rewrite(sink, pm_add_loads, name="** add loads (code)")
|
||||
|
||||
# create image buffers
|
||||
if IMAGE and ren.target.device in {"QCOM", "CL", "PYTHON", "NULL"}:
|
||||
sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True, ctx=ren.target.arch)
|
||||
|
||||
# devectorize
|
||||
sink = graph_rewrite(sink, sym+devectorize_alu+devectorize_buf_and_index+load_store_folding+correct_load_store+load_store_indexing,
|
||||
ctx=ren, name="devectorize")
|
||||
|
||||
# lower the index dtype to a concrete int
|
||||
sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing+gep_pushing, name="lower all index dtypes")
|
||||
sink = graph_rewrite(sink, symbolic, name="post index symbolic")
|
||||
|
||||
# optional pre matcher
|
||||
if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher")
|
||||
|
||||
# decompositions
|
||||
supported_ops = tuple(ren.code_for_op.keys())
|
||||
pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))
|
||||
pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
|
||||
sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="decompositions")
|
||||
sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
|
||||
sink = graph_rewrite(sink, pm_transcendental, name="transcendental")
|
||||
|
||||
# move gates from unrenderable INVALID where
|
||||
sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")
|
||||
|
||||
# final rules for the renderer (without sym)
|
||||
extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
|
||||
pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends
|
||||
sink = graph_rewrite(sink, pm_final_rewrite, ctx=ren, name="final rewrite")
|
||||
|
||||
# this was the linearizer
|
||||
sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
|
||||
|
||||
if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
|
||||
|
||||
# return the rewritten sink
|
||||
return sink
|
||||
|
||||
# inject IF/ENDIF. only needed if device doesn't support gated stores
|
||||
pm_linearize_cleanups = PatternMatcher([
|
||||
# if statements are not allowed in the graph
|
||||
(UPat((Ops.IF, Ops.ENDIF)), lambda: panic(RuntimeError, "if not allowed in graph")),
|
||||
# gated STORE becomes IF-STORE-ENDIF. this is the only use of IF-ENDIF
|
||||
(UPat(Ops.STORE, name="u", src=(UPat(Ops.INDEX).or_casted(), UPat(), UPat(name="gate", dtype=dtypes.bool))),
|
||||
lambda u, gate: (u, [mif:=UOp(Ops.IF, src=(gate, u.src[0])), u, UOp(Ops.ENDIF, src=(mif,))]))
|
||||
])
|
||||
|
||||
# requires lst be toposorted. like graph rewrite, but for lines
|
||||
def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]:
|
||||
newlst = []
|
||||
replaced: dict[UOp, UOp] = {}
|
||||
for u in lst:
|
||||
nu = u.replace(src=tuple([replaced.get(x, x) for x in u.src]))
|
||||
ret: tuple[UOp, list[UOp]] = cast(tuple[UOp, list[UOp]]|None, pm.rewrite(nu, ctx)) or (nu, [nu])
|
||||
replaced[u] = ret[0]
|
||||
newlst.extend(ret[1])
|
||||
return newlst
|
||||
|
||||
def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
|
||||
if DEBUG >= 3 and sink.arg.applied_opts: print(f"{sink.arg.function_name:<25} opts: {sink.arg.applied_opts}")
|
||||
lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
|
||||
if SPEC: type_verify(lst, spec_program)
|
||||
# isa renderers need to allocate registers
|
||||
if isinstance(ctx, ISARenderer):
|
||||
if ctx.pre_regalloc_matcher is not None: lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
|
||||
regalloc_ctx = LinearScanRegallocContext(lst, ctx)
|
||||
lst = line_rewrite(lst, pm_regalloc_rewrite, regalloc_ctx)
|
||||
lst = line_rewrite(lst, ctx.post_regalloc_matcher, regalloc_ctx)
|
||||
if DEBUG >= 4: print(ctx.asm_str(lst, sink.arg.function_name))
|
||||
return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
|
||||
|
||||
def do_estimates(prg:UOp, sink:UOp, lin:UOp) -> UOp|None:
|
||||
if sink.arg.estimates is not None: return None
|
||||
return prg.replace(src=(sink.replace(arg=replace(sink.arg, estimates=Estimates.from_uops(lin.src, ignore_indexing=True))),)+prg.src[1:])
|
||||
|
||||
def do_assemble(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
|
||||
src = "\n".join(str(u.arg) for u in lin.src)
|
||||
if DEBUG >= 4: print(src)
|
||||
binary = ctx.asm(prg, lin)
|
||||
return prg.replace(src=prg.src[:3]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
|
||||
|
||||
def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
|
||||
src = ctx.render(list(lin.src))
|
||||
new_arg = replace(prg.arg, aux=tuple(ctx.aux(list(lin.src)))) if ctx.has_aux else prg.arg
|
||||
return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=new_arg)
|
||||
|
||||
def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
|
||||
if DEBUG >= 4: print(source.arg)
|
||||
lib = ctx.compiler.compile_cached(source.arg)
|
||||
if DEBUG >= 7: ctx.compiler.disassemble(lib)
|
||||
return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
|
||||
|
||||
pm_to_program = PatternMatcher([
|
||||
(UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE)), name="prg"), do_linearize),
|
||||
(UPat(Ops.PROGRAM, src=(UPat(Ops.SINK, name="sink"), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_estimates),
|
||||
(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, src=UPat(Ops.INS), name="lin")), name="prg"), do_assemble),
|
||||
(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR, name="lin")), name="prg"), do_render),
|
||||
(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
|
||||
])
|
||||
|
||||
@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
|
||||
@Context(ALLOW_DEVICE_USAGE=0)
|
||||
def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
|
||||
"""
|
||||
Transform an AST into a compiled PROGRAM. May trigger BEAM search.
|
||||
|
||||
Args:
|
||||
ast: The Ops.SINK/Ops.PROGRAM rooted AST
|
||||
renderer: The renderer used to generate the code
|
||||
|
||||
Returns:
|
||||
The Ops.PROGRAM with SINK/DEVICE/LINEAR/SOURCE/BINARY.
|
||||
"""
|
||||
if ast.op is Ops.PROGRAM: prg = ast
|
||||
elif ast.op is Ops.SINK:
|
||||
assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
|
||||
full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
|
||||
prog_info = ProgramInfo.from_sink(full_sink)
|
||||
# instruction selection
|
||||
if isinstance(renderer, ISARenderer):
|
||||
full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True)
|
||||
full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True)
|
||||
prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=prog_info)
|
||||
else: raise RuntimeError(f"can't call to_program on {ast.op}")
|
||||
if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
|
||||
prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
|
||||
if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
|
||||
return prg
|
||||
|
||||
to_program_cache: dict[tuple, UOp] = {}
|
||||
def to_program(ast:UOp, renderer:Renderer) -> UOp:
|
||||
config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
|
||||
key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
|
||||
if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
|
||||
return prg
|
||||
@@ -0,0 +1,111 @@
|
||||
import math
|
||||
from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
|
||||
from tinygrad.helpers import dedup, get_contraction
|
||||
from tinygrad.dtype import dtypes, AddrSpace, Invalid
|
||||
from tinygrad.renderer import Renderer
|
||||
|
||||
def _dim_max(d:sint) -> int: return d if isinstance(d, int) else int(d.vmax)
|
||||
|
||||
def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
|
||||
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 i < (len(dims)-1) and _dim_max(dims[i]) * _dim_max(dims[i+1]) <= m:
|
||||
dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
|
||||
break
|
||||
else: return None
|
||||
return dims
|
||||
|
||||
def _split_dims(dims, max_sizes):
|
||||
if all(d <= m for d,m in zip(dims, max_sizes)): return dims
|
||||
_dims = list(dims) + [1]*(3-len(dims))
|
||||
for i in range(len(_dims)):
|
||||
while _dims[i] > max_sizes[i]:
|
||||
div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
|
||||
if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
|
||||
_dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
|
||||
return tuple(_dims[:2] if _dims[2] == 1 else _dims[0] if _dims[1:3] == [1,1] else _dims)
|
||||
|
||||
def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
|
||||
if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
|
||||
if max_sizes is None: limited = dims
|
||||
else:
|
||||
# try to group first: (a, b, c, d) -> (ab, c, d)
|
||||
limited = grouped if (grouped := _group_dims(dims, max_sizes)) else dims
|
||||
# check if grouping failed
|
||||
if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
|
||||
# try to split up dims: (a,) -> (b, c)
|
||||
if limited == dims: limited = _split_dims(dims, max_sizes)
|
||||
raw_idxs = [UOp.special(s, f"{prefix}{i}") for i,s in enumerate(limited)]
|
||||
if len(limited) < len(dims):
|
||||
ret = []
|
||||
if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(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
|
||||
elif (a:=len(limited)) > (b:=len(dims)):
|
||||
if a == 2 and b == 1: return [raw_idxs[0] * limited[1] + raw_idxs[1]]
|
||||
if a == 3 and b == 1: return [(raw_idxs[0] * limited[1] + raw_idxs[1]) * limited[2] + raw_idxs[2]]
|
||||
if limited != dims:
|
||||
# Convert to 1D
|
||||
flat = raw_idxs[0]*limited[1]+raw_idxs[1] if len(limited) == 2 else raw_idxs[0]*(limited[1]*limited[2])+raw_idxs[1]*limited[2]+raw_idxs[2]
|
||||
# Get back original indices from 1D
|
||||
return [flat//dims[1], flat%dims[1]] if len(dims) == 2 else [flat//(dims[2]*dims[1]), (flat//dims[2])%dims[1], flat%dims[2]]
|
||||
return raw_idxs
|
||||
|
||||
def add_gpudims(ctx:Renderer, s:UOp):
|
||||
if s.arg is None: return None
|
||||
s_topo = list(s.toposort())
|
||||
if any(x.op is Ops.SPECIAL for x in s_topo): return None
|
||||
|
||||
# get ranges
|
||||
all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
|
||||
|
||||
# extract global/local dims
|
||||
global_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)]))
|
||||
local_dims = sorted(dedup([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
|
||||
if not global_dims and not local_dims: return None
|
||||
|
||||
# get global and local shape
|
||||
ranges = [all_ranges[r] for r in global_dims+local_dims if r in all_ranges]
|
||||
global_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in global_dims])
|
||||
local_shape = tuple([ssimplify(r.src[0]) for r in ranges if r.arg[0:-1] in local_dims])
|
||||
|
||||
# get the idxs
|
||||
ki: KernelInfo = s.arg
|
||||
if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
|
||||
elif ki.dont_use_locals:
|
||||
assert not local_dims, "can't use locals if there's no local dims"
|
||||
idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
|
||||
else:
|
||||
# define indexes for GPU-like execution
|
||||
local_idxs = get_grouped_dims("lidx", local_shape, ctx.local_max)
|
||||
hw_local = [_dim_max(u.src[0]) for u in local_idxs if u.op is Ops.SPECIAL]
|
||||
global_max = ctx.global_max if ctx.global_prod_max is None else \
|
||||
tuple(min(gm, pm//l) for gm,pm,l in zip(ctx.global_max or ctx.global_prod_max, ctx.global_prod_max, hw_local+[1]*3))
|
||||
idxs = get_grouped_dims("gidx", global_shape, global_max, reverse=True) + local_idxs
|
||||
|
||||
# apply to multiple ranges
|
||||
subs = {}
|
||||
for r in s_topo:
|
||||
# look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
|
||||
if r.op is Ops.STORE and (idx := r.src[0]).src[0].ptrdtype.addrspace == AddrSpace.GLOBAL:
|
||||
missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
|
||||
if len(missing_locals):
|
||||
assert len(idx.src) == 2, "index has 2 sources"
|
||||
mask: UOp = UOp.uprod(*[x.eq(0) for x in missing_locals])
|
||||
subs[idx] = idx.replace(src=(idx.src[0], mask.broadcast(idx.src[1].dtype.count).where(idx.src[1], Invalid)))
|
||||
if r.op is not Ops.RANGE: continue
|
||||
try:
|
||||
ii = (global_dims+local_dims).index(r.arg[0:-1])
|
||||
if r.arg[1] == AxisType.REDUCE: continue
|
||||
subs[r] = idxs[ii]
|
||||
except ValueError: continue
|
||||
return s.substitute(subs)
|
||||
|
||||
pm_add_gpudims = PatternMatcher([
|
||||
# add gpudims must be last
|
||||
(UPat(Ops.SINK, name="s"), add_gpudims),
|
||||
])
|
||||
@@ -0,0 +1,388 @@
|
||||
from typing import Any, cast
|
||||
import functools, itertools
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
|
||||
from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
|
||||
from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
|
||||
from tinygrad.helpers import getenv, flatten, prod
|
||||
from tinygrad.renderer import Renderer
|
||||
|
||||
# ***** image load valid simplification *****
|
||||
|
||||
@functools.cache
|
||||
def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
|
||||
# can drop valid if idx is out of bound when valid is False
|
||||
drop_stmt = []
|
||||
for stmt in valid.split_uop(Ops.AND):
|
||||
if (res:=parse_valid(stmt)) is None: continue
|
||||
X, is_upper_bound, c = res
|
||||
|
||||
# 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 X.split_uop(Ops.ADD)):
|
||||
testidx = functools.reduce(lambda nowidx,u: nowidx.substitute({u:u.const_like(0)}), X.split_uop(Ops.ADD), idx)
|
||||
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, (width, height)):
|
||||
if i.is_increasing():
|
||||
rw = i.substitute({X:X.const_like(test_value)})
|
||||
if rw.vmin >= b or rw.vmax < 0:
|
||||
drop_stmt.append(stmt)
|
||||
break
|
||||
return drop_stmt
|
||||
|
||||
def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
|
||||
idx = uop_given_valid(valid, start_idx)
|
||||
return None if idx is start_idx else buf.index(idx.valid(valid), ptr=True)
|
||||
|
||||
def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|None:
|
||||
if not isinstance(buf.dtype, ImageDType): return None
|
||||
start_idx = UOp.vectorize(idx_x, idx_y)
|
||||
idx = uop_given_valid(valid, start_idx)
|
||||
drop_stmt = _drop_valid_stmts(valid, idx, buf.dtype.shape[0], buf.dtype.shape[1])
|
||||
|
||||
if not drop_stmt and idx is start_idx: return None
|
||||
new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
|
||||
idx_y, idx_x = idx.gep(1), idx.gep(0)
|
||||
return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), ptr=True) if new_valid is not None else buf.index(idx_y, idx_x, ptr=True)
|
||||
|
||||
load_store_indexing = PatternMatcher([
|
||||
# image load valid idx simplification
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), invalid_gate)), lambda buf,x,i,cond: simplify_valid_load(buf, x, cond)),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("valid").where(UPat.var("idx_y"), UPat(arg=Invalid)),
|
||||
UPat.var("valid").where(UPat.var("idx_x"), UPat(arg=Invalid)))), simplify_valid_image_load),
|
||||
])
|
||||
|
||||
# ***** load/store grouping *****
|
||||
|
||||
def expand_index(ctx, buf:UOp, vec:UOp):
|
||||
# determine optimal image shapes
|
||||
if isinstance(dt:=buf.dtype, ImageDType):
|
||||
x, valid = vec.get_idx().gep(0), vec.get_valid().gep(0)
|
||||
# search for dims that drop the most valid statements
|
||||
best_drop, cands = -1, []
|
||||
for ch, cw in ImageDType.valid_dims(dt, ctx.target.arch):
|
||||
if (dropped:=len(_drop_valid_stmts(valid, cidx:=uop_given_valid(valid, UOp.vectorize((x//4)%cw, x//(4*cw))), ch, cw))) > best_drop:
|
||||
best_drop, cands = dropped, [(ch, cw, cidx)]
|
||||
elif dropped == best_drop: cands.append((ch, cw, cidx))
|
||||
# and tiebreak with indexing complexity (ie. number of nodes)
|
||||
h, w, _ = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].gep(1).simplify().backward_slice))
|
||||
assert buf.op is Ops.RESHAPE
|
||||
buf = buf.src[0].replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4))).flatten()
|
||||
if getenv("UNSAFE_DISABLE_MASK", 0): vec = vec.get_idx()
|
||||
# generate the individual indexes
|
||||
return UOp(Ops.STACK, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count)))
|
||||
|
||||
def fold_expanded_index(midx:UOp):
|
||||
buf = midx.src[0].src[0]
|
||||
if not all(s.src[0] is buf for s in midx.src): return None
|
||||
if not all(isinstance(s.dtype, PtrDType) for s in midx.src): return None
|
||||
|
||||
# extract all the relevant offsets
|
||||
offsets_rootsrc: defaultdict[Any, dict[int, list[int]]] = defaultdict(dict)
|
||||
for i in range(len(midx.src)):
|
||||
idx: Any = midx.src[i].src[1].get_idx()
|
||||
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.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg
|
||||
elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0
|
||||
elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg
|
||||
else: root_src, arg = idx, 0
|
||||
root_src = (midx.src[i].src[1].get_valid(), root_src)
|
||||
offsets_rootsrc[root_src].setdefault(arg, []).append(i)
|
||||
|
||||
# then rewrite everything we can into groups
|
||||
ret = []
|
||||
idxs: list[int|None] = [None]*len(midx.src)
|
||||
global_offset = 0
|
||||
for offsets in offsets_rootsrc.values():
|
||||
grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
|
||||
for grp in grouped_offsets:
|
||||
# get the index offset for this element. using [0] is okay, because they are the same
|
||||
lidx = midx.src[offsets[grp[0]][0]]
|
||||
if len(grp) > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(len(grp)).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
|
||||
# set the idxs of the output
|
||||
for i,g in enumerate(grp):
|
||||
for oo in offsets[g]: idxs[oo] = global_offset+i
|
||||
# add this lidx to the CAT
|
||||
ret.append(lidx)
|
||||
global_offset += len(grp)
|
||||
assert None not in idxs, f"some idxs are missing {idxs}"
|
||||
# this base thing is for image, we want the CAT to be a normal pointer
|
||||
post_cat = UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace).vec(global_offset), tuple(ret))
|
||||
return post_cat.gep(tuple(cast(list[int], idxs)))
|
||||
|
||||
def cat_after_store(cat:UOp, data:UOp):
|
||||
# TODO: this is written in many places
|
||||
offset = 0
|
||||
ret: list[UOp] = []
|
||||
for s in cat.src:
|
||||
ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count)))))
|
||||
offset += s.dtype.count
|
||||
return UOp.group(*ret)
|
||||
|
||||
def gep_on_store(gep:UOp, st:UOp):
|
||||
# NOTE: we need to invert the gep here, but it may be an expanding gep
|
||||
# fake argsort. TODO: handle duplicates
|
||||
a = {}
|
||||
for i,x in enumerate(gep.arg): a[x] = i
|
||||
new_arg = tuple(x[1] for x in sorted(a.items()))
|
||||
return gep.src[0].store(st.gep(new_arg))
|
||||
|
||||
load_store_folding = PatternMatcher([
|
||||
(UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(name="buf")), UPat.var("vec"))), expand_index),
|
||||
(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="midx"), fold_expanded_index),
|
||||
# GEP after LOAD
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True),
|
||||
lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)),
|
||||
# GEP on data of STORE
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st"))), gep_on_store),
|
||||
# put PTRCAT after LOAD
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
|
||||
lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
|
||||
# put PTRCAT after STORE
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store),
|
||||
])
|
||||
|
||||
# *** correct load/store ***
|
||||
|
||||
def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
|
||||
# this splits loads and stores into multiple chunks
|
||||
|
||||
# if there's only one element to load/store, no splitting needed
|
||||
if (sz:=ls.src[0].dtype.count) == 1: return None
|
||||
buf = idx.src[0]
|
||||
|
||||
# determine fold lengths
|
||||
lengths = []
|
||||
must_divide = True
|
||||
if ctx is not None and ctx.target.device == "DSP":
|
||||
lengths = [128,64,32,16,8,4]
|
||||
must_divide = False
|
||||
elif buf.dtype.base not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
|
||||
pass
|
||||
elif buf.ptrdtype.addrspace == AddrSpace.REG:
|
||||
pass
|
||||
elif isinstance(buf.dtype, ImageDType):
|
||||
lengths = [4]
|
||||
elif ctx is not None and ctx.supports_float4:
|
||||
# TODO: a better way to get this than ctx
|
||||
lengths = [8,4,2] if buf.dtype.base == dtypes.half and getenv("ALLOW_HALF8") else ([16,8,4,2] if "AMX" in ctx.target.arch else [4,2])
|
||||
lengths.append(1) # worst case, it's not folded
|
||||
|
||||
# filter fold lengths that don't divide
|
||||
offset, mask = idx.src[1].get_idx(), idx.src[1].get_valid()
|
||||
if must_divide: lengths = [x for x in lengths if offset.divides(x) is not None]
|
||||
|
||||
# split based on the fold lengths
|
||||
global_offset = 0
|
||||
ret = []
|
||||
while global_offset < sz:
|
||||
# with 1 at the end of the lengths list, this will always hit
|
||||
for fold_length in lengths:
|
||||
if global_offset+fold_length > sz: continue
|
||||
lidx = buf.index((offset + global_offset).valid(mask), ptr=True)
|
||||
if fold_length > 1: lidx = lidx.cast(buf.ptrdtype.base.vec(fold_length).ptr(size=buf.ptrdtype.size, addrspace=buf.ptrdtype.addrspace))
|
||||
if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx,ls.src[1].gep(tuple(range(global_offset, global_offset+fold_length))))))
|
||||
else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar().vec(fold_length)))
|
||||
global_offset += fold_length
|
||||
break
|
||||
|
||||
# if it wasn't split, we return None. otherwise we CAT them
|
||||
if len(ret) <= 1: return None
|
||||
return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
|
||||
|
||||
def get_image_idx(idx:UOp, width:int):
|
||||
x, valid = idx.src[1].get_idx(), idx.src[1].get_valid()
|
||||
idx_x, idx_y = (x // 4) % width, x // (4*width)
|
||||
assert idx.src[0].op is Ops.RESHAPE, "image idx must be on reshape"
|
||||
return idx.replace(src=(idx.src[0].src[0], idx_y.valid(valid), idx_x.valid(valid)))
|
||||
|
||||
def image_fixup(ls:UOp):
|
||||
# normal image load or store, with the CAST from expand_index
|
||||
if isinstance(dt:=ls.src[0].src[0].dtype, ImageDType) and ls.src[0].op is Ops.CAST:
|
||||
assert ls.src[0].dtype.count == 4, "image must be casted to 4"
|
||||
return ls.replace(src=(get_image_idx(ls.src[0].src[0], dt.shape[1]),)+ls.src[1:])
|
||||
|
||||
# this is an unprocessed image without a cast, we should just make it a buffer
|
||||
if isinstance(dt, ImageDType) and len(ls.src[0].src) == 2:
|
||||
off = ls.src[0].src[1]
|
||||
assert ls.src[0].src[0].op is Ops.RESHAPE, "image idx must be on reshape"
|
||||
idx = ls.src[0].src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
|
||||
return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls.src[1].cast(new_dt)))
|
||||
|
||||
correct_load_store = PatternMatcher([
|
||||
# split LOAD/STORE
|
||||
(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, name="idx").cast(),), name="ls", allow_any_len=True), split_load_store),
|
||||
# image indexing, including unfoldable images
|
||||
(UPat((Ops.LOAD, Ops.STORE), name="ls"), image_fixup),
|
||||
])
|
||||
|
||||
# *** uop expander ***
|
||||
|
||||
# TODO: there's a lot shared with gep_through_wmma here
|
||||
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.STACK, wmma.dtype, tuple(wmma_ex))
|
||||
|
||||
def no_vectorized_alu(alu:UOp):
|
||||
if alu.dtype.vcount == 1: return None
|
||||
if alu.op is Ops.WHERE and alu.src[2].arg is Invalid: return None # image load/store has cond.where(idx.vec(2), Invalid) as the index
|
||||
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.STACK, alu.dtype, alus)
|
||||
|
||||
def no_vectorized_buf(buf:UOp):
|
||||
return buf.replace(dtype=buf.ptrdtype.base.scalar().ptr(buf.ptrdtype.size*buf.ptrdtype.count, buf.ptrdtype.addrspace)).cast(buf.dtype)
|
||||
|
||||
def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
|
||||
cnt = cast.dtype.count
|
||||
if bcast is not None and bcast.op is Ops.GEP:
|
||||
# GEP selects specific lanes; bcast.arg[k] is the offset for lane k, iterate groups × selected lanes
|
||||
pairs = [(k, g + bcast.arg[k]) for g, k in itertools.product(range(cast.dtype.vcount), range(len(bcast.arg)))]
|
||||
elif bcast is not None:
|
||||
# BROADCAST: cross product of components × lanes
|
||||
pairs = [(j, c) for c, j in itertools.product(range(cnt), range(bcast.dtype.vcount))]
|
||||
else:
|
||||
# simple scalar index: one lane, all components
|
||||
pairs = [(0, c) for c in range(cnt)]
|
||||
idx_lanes, offsets = (tuple(x) for x in zip(*pairs))
|
||||
return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.weakint.vec(len(pairs)), offsets), ptr=True)
|
||||
|
||||
devectorize_buf_and_index = PatternMatcher([
|
||||
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
|
||||
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").index(UPat.var("idx")), no_vectorized_index),
|
||||
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").broadcast(name="bcast").index(UPat.var("idx")),
|
||||
no_vectorized_index),
|
||||
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)).or_after(name="buf").cast(name="cast").gep(name="bcast").index(UPat.var("idx")),
|
||||
no_vectorized_index),
|
||||
])
|
||||
|
||||
devectorize_alu = PatternMatcher([
|
||||
# CAST after AFTER
|
||||
(UPat(Ops.CAST, name="c").f(Ops.AFTER, allow_any_len=True, name="a"), lambda c,a: c.src[0].after(*a.src[1:]).cast(c.dtype)),
|
||||
# no ALU on vectorized dtypes
|
||||
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
|
||||
(UPat(Ops.WMMA, name="wmma"), no_vectorized_wmma),
|
||||
])
|
||||
|
||||
pm_render = PatternMatcher([
|
||||
# for rendering, we use explicit VECTORIZE
|
||||
(UPat(Ops.CONST, name='c'),
|
||||
lambda c: UOp(Ops.STACK, c.dtype, (UOp.const(c.dtype.scalar(), c.arg),)*c.dtype.vcount) if c.dtype.vcount > 1 else None),
|
||||
(UPat(Ops.GEP, name='gep'), lambda gep: UOp(Ops.STACK, gep.dtype, tuple(gep.src[0].gep(x) for x in gep.arg)) if len(gep.arg) > 1 else None),
|
||||
(UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
|
||||
(UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
|
||||
])
|
||||
|
||||
# *** Ops.REDUCE -> Ops.DEFINE_ACC ***
|
||||
|
||||
@dataclass
|
||||
class ReduceContext:
|
||||
acc_num: int = 0
|
||||
|
||||
def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
|
||||
# if this has a horizontal reduction component, do that first
|
||||
if inp.dtype != out_dtype:
|
||||
# NOTE: [0 1 2 3 4 5 6 7] -> [0+4, 1+5, 2+6, 3+7]
|
||||
horizontal_amount = inp.dtype.count//out_dtype.count
|
||||
return [inp.gep(tuple(range(i, inp.dtype.count, horizontal_amount))) for i in range(0, horizontal_amount)]
|
||||
return [inp]
|
||||
|
||||
def reduce_to_acc(ctx:ReduceContext, red:UOp):
|
||||
inp, reduce_range = red.src[0], red.src[1:]
|
||||
lst = horizontal_reduce(inp, red.dtype)
|
||||
assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}"
|
||||
# if we have a range
|
||||
if len(reduce_range) != 0:
|
||||
topo = inp.toposort()
|
||||
ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
|
||||
input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in ended_ranges])
|
||||
identity = red.const(red.dtype, identity_element(red.arg[0], red.dtype.scalar()))
|
||||
acc = UOp.placeholder((1,), red.dtype, ctx.acc_num, AddrSpace.REG)
|
||||
acc_init = acc.after(*input_ranges).index(UOp.const(dtypes.weakint, 0)).store(identity)
|
||||
lst = [acc.after(acc_init, *reduce_range).index(UOp.const(dtypes.weakint, 0))] + lst # put acc as the first element
|
||||
ctx.acc_num += 1
|
||||
ret = functools.reduce(lambda x,y: x.alu(red.arg[0], y), lst)
|
||||
if len(reduce_range) == 0: return ret
|
||||
end = acc.index(UOp.const(dtypes.weakint, 0)).store(ret).end(*reduce_range).rtag("mergeable")
|
||||
return acc.after(end).index(UOp.const(dtypes.weakint, 0))
|
||||
|
||||
def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
|
||||
# merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
|
||||
# ENDs at different nesting depths get cloned RANGEs so each RANGE maps to one END
|
||||
range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
|
||||
for u in sink.backward_slice:
|
||||
if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
|
||||
subs: dict[UOp, UOp] = {}
|
||||
next_axis = max((u.arg[0] for u in sink.backward_slice if u.op is Ops.RANGE), default=-1) + 1
|
||||
for r, ends in range_to_ends.items():
|
||||
if len(ends) <= 1: continue
|
||||
by_ctx: dict[frozenset[UOp], list[UOp]] = {}
|
||||
for e in ends: by_ctx.setdefault(frozenset(e.ranges), []).append(e)
|
||||
for i, group in enumerate(by_ctx.values()):
|
||||
tr = r if i == 0 else tuple(rr.replace(arg=(next_axis + j, *rr.arg[1:])) for j, rr in enumerate(r))
|
||||
if i > 0: next_axis += len(r)
|
||||
mapped = [e.substitute(dict(zip(r, tr))) if i > 0 else e for e in group]
|
||||
merged = mapped[0] if len(mapped) == 1 else UOp.group(*(e.src[0] for e in mapped)).end(*tr)
|
||||
for e in group: subs[e] = merged
|
||||
return sink.substitute(subs) if subs else None
|
||||
|
||||
pm_reduce = PatternMatcher([
|
||||
# REDUCE -> DEFINE_ACC+ASSIGN, then merge ENDs with same range
|
||||
(UPat(Ops.REDUCE, name="red"), reduce_to_acc),
|
||||
(UPat(Ops.SINK, name="sink"), merge_reduce_ends),
|
||||
# tensor core built in accumulate
|
||||
(UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
|
||||
lambda add, wmma: UOp(wmma.op, wmma.dtype, (wmma.src[0], wmma.src[1], wmma.src[2]+add), wmma.arg)),
|
||||
])
|
||||
|
||||
# add loads
|
||||
|
||||
def add_load(idx:UOp):
|
||||
if isinstance(idx.dtype, PtrDType): return None
|
||||
assert isinstance(idx.src[0].dtype, PtrDType), f"param is not PtrDType {idx.src[0].dtype}"
|
||||
return idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)
|
||||
|
||||
pm_add_loads = PatternMatcher([
|
||||
# add loads to non ptr index
|
||||
(UPat(Ops.INDEX, name="idx"), add_load),
|
||||
# remove loads from stores
|
||||
(UPat(Ops.STORE, src=(UPat(Ops.LOAD),), allow_any_len=True, name="s"), lambda s: s.replace(src=(s.src[0].src[0],)+s.src[1:])),
|
||||
(UPat(Ops.LOAD, src=(UPat(Ops.LOAD),), allow_any_len=True, name="l"), lambda l: l.replace(src=(l.src[0].src[0],)+l.src[1:])),
|
||||
])
|
||||
|
||||
# make images
|
||||
|
||||
pm_imageh_store = PatternMatcher([
|
||||
# store<imageh>(idx, x) is actually store(idx, x.cast(half)) so we can pull the cast into the store
|
||||
(UPat.var("x", dtypes.float).cast(dtypes.half), lambda x: x),
|
||||
# store(imageh, a.where(b.half(), c).float()) -> store(imageh, a.where(b, c.float()))
|
||||
(UPat(Ops.WHERE, src=(UPat.var("a"), UPat.var("b", dtypes.float).cast(dtypes.half), UPat.var("c"))), lambda a,b,c: a.where(b,c.cast(dtypes.float))),
|
||||
# otherwise, we cast to float
|
||||
(UPat(GroupOp.All, name="x"), lambda x: x.cast(dtypes.float))
|
||||
])
|
||||
|
||||
def make_image(ctx, ls, buf, off):
|
||||
if (vcount:=buf.dtype.vcount) != 1: buf = buf.src[0]
|
||||
if buf.op == Ops.PARAM and not isinstance(dt:=buf.dtype, ImageDType) and (dims:=ImageDType.valid_dims(dt, ctx)):
|
||||
buf = buf.replace(dtype=(dtypes.imageh if dt.base == dtypes.half else dtypes.imagef)((*dims[0], 4))).flatten()
|
||||
if vcount != 1: buf = UOp.vectorize(*([buf] * vcount))
|
||||
if ls.op is Ops.LOAD: return ls.replace(src=(buf.index(off, ptr=True),), dtype=dtypes.float.vec(ls.dtype.vcount)).cast(dt.base)
|
||||
return buf.index(off, ptr=True).store(pm_imageh_store.rewrite(ls.src[1]) if dt.base == dtypes.half else ls.src[1])
|
||||
|
||||
pm_make_images = PatternMatcher([
|
||||
(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("off"))),), allow_any_len=True, name="ls"), make_image),
|
||||
# load<imageh> is actually load<half>.cast(float), so load<imageh>.half().float() -> load<half>.float().half().float() -> load<half>.float()
|
||||
(UPat(Ops.LOAD, name="li").cast(dtypes.half).cast(dtypes.float), lambda li: li if isinstance(li.src[0].dtype, ImageDType) else None),
|
||||
])
|
||||
@@ -0,0 +1,160 @@
|
||||
# this converts a lowerer program into a vectorized program
|
||||
import functools, itertools
|
||||
from tinygrad.dtype import dtypes, PtrDType, AddrSpace
|
||||
from tinygrad.helpers import dedup, flatten, all_same, prod, partition
|
||||
from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, AxisType, range_start
|
||||
from tinygrad.schedule.rangeify import BufferizeOpts
|
||||
|
||||
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.cache
|
||||
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 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 in range_start and i >= range_start[root.op]:
|
||||
# for any range args of REDUCE/WMMA/END/etc., pass them through
|
||||
new_srcs.append(src)
|
||||
elif root.op is Ops.INDEX and i >= 1 and not isinstance(root.dtype, PtrDType):
|
||||
new_srcs.append(src)
|
||||
elif src.dtype.count > 1:
|
||||
# put any input dtype > 1 grouped together
|
||||
new_srcs.append(UOp(Ops.VCAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
|
||||
else:
|
||||
# repeat the arg
|
||||
new_srcs.append(src.broadcast(expand_sz))
|
||||
|
||||
# for non-PtrDType INDEX on REG buffers, expand into individual scalar INDEXes instead of one vectorized INDEX
|
||||
# this avoids creating a VECTORIZE of REG pointers which the devectorizer can't resolve
|
||||
if root.op is Ops.INDEX and not isinstance(root.dtype, PtrDType) and \
|
||||
isinstance(root.src[0].dtype, PtrDType) and root.src[0].dtype.addrspace == AddrSpace.REG:
|
||||
idxs = []
|
||||
for j in range(expand_sz):
|
||||
idx_srcs = tuple(s.gep(j) if isinstance(s.dtype, PtrDType) or s.dtype.count > 1 else s for s in new_srcs)
|
||||
idxs.append(UOp(Ops.INDEX, root.dtype, idx_srcs, root.arg))
|
||||
return UOp(Ops.UNROLL, root.dtype, (UOp(Ops.STACK, root.dtype.vec(expand_sz), tuple(idxs)),), expand_args)
|
||||
|
||||
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.STACK, con.dtype, con.src*con.dtype.count)
|
||||
# CONTRACT may remove several axes from UNROLL
|
||||
assert con.dtype == dtypes.void or 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 end_unrolls(u:UOp):
|
||||
unrolls, src = partition(u.src[1:], lambda x: x.op is Ops.UNROLL)
|
||||
if not len(unrolls): return None
|
||||
ret = UOp(Ops.CONTRACT, dtypes.void, (u.src[0],), sum([x.arg for x in unrolls], start=()))
|
||||
return u.replace(src=(ret,)+tuple(src))
|
||||
|
||||
expander = PatternMatcher([
|
||||
# push broadcast through AFTER/END
|
||||
(UPat.var("x").broadcast(name="b").after(name="a", allow_any_len=True), lambda x,b,a: x.after(*a.src[1:]).broadcast(len(b.src))),
|
||||
(UPat.var("x").broadcast(name="b").end(name="a", allow_any_len=True), lambda x,b,a: x.end(*a.src[1:]).broadcast(len(b.src))),
|
||||
# END on UNROLL ends the UNROLL
|
||||
(UPat(Ops.END, name="u"), end_unrolls),
|
||||
# BUFFERIZE puts UNROLLs for ranges as contract
|
||||
(UPat(Ops.STAGE, src=(UPat(Ops.UNROLL), UPat(Ops.UNROLL)), name="x"),
|
||||
lambda x: x.replace(src=tuple(UOp(Ops.CONTRACT, dtype=s.dtype.vec(x.src[1].src[0].dtype.count), src=(s,), arg=x.src[1].arg) for s in x.src))),
|
||||
# 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.STAGE,
|
||||
Ops.STACK, Ops.REDUCE, Ops.END, Ops.AFTER), name="root", custom_early_reject=set([Ops.UNROLL])), do_expand),
|
||||
(UPat(Ops.CONTRACT, name="con"), do_contract),
|
||||
# empty UNROLL is NOOP
|
||||
(UPat(Ops.UNROLL, src=(UPat.var('x'),), arg=()), lambda x: x),
|
||||
])
|
||||
|
||||
# ****
|
||||
|
||||
def fix_reduce_unroll(x:UOp):
|
||||
reduce_range, reduce_expand = partition(x.src[1:], lambda y: y.op is Ops.RANGE)
|
||||
if len(reduce_expand) == 0: return None
|
||||
reduce_expand = [x for x in reduce_expand if x.op is not Ops.CONST]
|
||||
assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand}"
|
||||
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), tag=1)
|
||||
# REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
|
||||
return x.replace(src=(ret,)+tuple(reduce_range))
|
||||
|
||||
def fix_store_unroll(x:UOp):
|
||||
store_expand, store_range = partition(x.src[2:], lambda y: y.op is Ops.UNROLL)
|
||||
if len(store_expand) == 0: return None
|
||||
return UOp(Ops.CONTRACT, dtypes.void, (x.replace(src=x.src[:2]+tuple(store_range)),), tuple(flatten(x.arg for x in store_expand)), tag=1)
|
||||
|
||||
def fix_group_for_reduce(x:UOp):
|
||||
reduce_gfr, reduce_r = partition(x.src[1:], lambda u: u.op is Ops.RANGE and u.arg[1] == AxisType.GROUP_REDUCE)
|
||||
if len(reduce_gfr) == 0: return None
|
||||
|
||||
# NOTE: if there's other locals here, we need them in the buffer too
|
||||
upstream_locals = [u for u in x.toposort() if u.op is Ops.RANGE and u.arg[1] == AxisType.LOCAL]
|
||||
|
||||
# do only the non grouped reduces early
|
||||
ret = x.replace(src=(x.src[0],)+tuple(reduce_r))
|
||||
reduce_loop = [x.replace(arg=(x.arg[0]+100, AxisType.REDUCE)) for x in reduce_gfr]
|
||||
buf = ret.bufferize(*upstream_locals, *reduce_gfr, arg=BufferizeOpts(reduce_gfr[0].arg[0], AddrSpace.LOCAL)).index(*upstream_locals, *reduce_loop)
|
||||
|
||||
# do the final reduce (if/barrier are added in gpudims step)
|
||||
return buf.reduce(*reduce_loop, arg=x.arg)
|
||||
|
||||
pm_pre_expander = PatternMatcher([
|
||||
# rewrite UPCAST/UNROLL range to something to be expanded
|
||||
(UPat(Ops.RANGE, name="r"),
|
||||
lambda r: UOp(Ops.UNROLL, r.dtype, (UOp.const(r.dtype.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
|
||||
if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
|
||||
# fix REDUCEs with UNROLLs
|
||||
(UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
|
||||
(UPat(Ops.STORE, name="x"), fix_store_unroll),
|
||||
])
|
||||
|
||||
pm_group_for_reduce = PatternMatcher([
|
||||
# fix group for reduce
|
||||
(UPat(Ops.REDUCE, name="x"), fix_group_for_reduce),
|
||||
])
|
||||
@@ -0,0 +1,25 @@
|
||||
# this is a temporary intermediate step while we remove this index style
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops
|
||||
from tinygrad.dtype import Invalid, dtypes
|
||||
|
||||
pm_move_gates_from_index = PatternMatcher([
|
||||
# here we create the alt value for load to be 0s and remove the where Invalid
|
||||
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").load(name="l"),
|
||||
lambda buf,gate,idx,cast,l: buf.index(idx, ptr=True).cast(cast.dtype).load(l.const_like(0), gate, dtype=l.dtype)),
|
||||
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").store(UPat.var("data")),
|
||||
lambda buf,gate,idx,cast,data: buf.index(idx, ptr=True).cast(cast.dtype).store(data, gate)),
|
||||
|
||||
# for image idx
|
||||
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
|
||||
UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).or_casted(name="cast").load(name="l"),
|
||||
lambda buf,gate,idx_y,idx_x,cast,l: buf.index(idx_y, idx_x, ptr=True).cast(cast.dtype).load(l.const_like(0), gate, dtype=l.dtype)),
|
||||
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
|
||||
UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).or_casted(name="cast").store(UPat.var("data")),
|
||||
lambda buf,gate,idx_y,idx_x,cast,data: buf.index(idx_y, idx_x, ptr=True).cast(cast.dtype).store(data, gate)),
|
||||
|
||||
# Where after gated load becomes alt value
|
||||
(UPat.var("gate").where(UPat().load(UPat(), UPat.var("gate", dtype=dtypes.bool), name="l").or_casted(), UPat.var("a")), lambda gate,l,a:
|
||||
l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
|
||||
(UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()), lambda gate,l,a:
|
||||
l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
|
||||
])
|
||||
@@ -0,0 +1,96 @@
|
||||
import heapq
|
||||
from typing import Any
|
||||
from collections import defaultdict
|
||||
from tinygrad.uop.ops import PatternMatcher, UOp, Ops, UPat, multirange_str
|
||||
from tinygrad.helpers import prod, getenv, TUPLE_ORDER
|
||||
|
||||
def linearize(sink:UOp) -> list[UOp]:
|
||||
# this is a toposort with priority
|
||||
lst = list(sink.toposort())
|
||||
out_degree:defaultdict[UOp, int] = defaultdict(int)
|
||||
priorities:dict[UOp, tuple[int, int, Any]] = {}
|
||||
|
||||
# get consumers and assign priorities
|
||||
# NOTE: this requires the lst be locally toposorted
|
||||
for u in reversed(lst):
|
||||
for s in u.src: out_degree[s] += 1
|
||||
|
||||
# we place UOps with higher run_counts later
|
||||
run_count = prod([int(r.vmax)+1 for r in u.ranges])
|
||||
|
||||
# simple priority override. this is all bottom up now, smaller numbers will be closer to the top
|
||||
extra = None
|
||||
match u.op:
|
||||
# the order and placement of these defines is important
|
||||
case Ops.PARAM: priority, extra = -20, u.arg
|
||||
case Ops.DEFINE_VAR: priority, extra = -19, u.arg
|
||||
case Ops.DEFINE_REG: priority = -18
|
||||
case Ops.DEFINE_LOCAL: priority = -17
|
||||
case Ops.LOAD: priority = -1 # place loads early
|
||||
case Ops.STORE: priority = 1 # place stores late
|
||||
case Ops.RANGE: priority = 5 # placing RANGE is good
|
||||
case Ops.END: priority = -5 # placing END is bad
|
||||
case _: priority = 0 # everything else has priority 0
|
||||
priorities[u] = (run_count, priority, extra)
|
||||
|
||||
# number the uops in "ideal" order
|
||||
nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}
|
||||
|
||||
# then force them to be toposorted in as close to the ideal order as possible
|
||||
heap = [(-nkey[sink], sink)]
|
||||
newlst = []
|
||||
while heap:
|
||||
newlst.append(u:=heapq.heappop(heap)[1])
|
||||
for v in u.src:
|
||||
out_degree[v] -= 1
|
||||
if out_degree[v] == 0: heapq.heappush(heap, (-nkey[v],v))
|
||||
newlst = newlst[::-1]
|
||||
|
||||
if getenv("DEBUG_LINEARIZE"):
|
||||
for i,u in enumerate(newlst):
|
||||
print(f"{i:4d} {str(u.op):20s} {multirange_str(u.ranges, color=True, pad=10)} {priorities[u]}")
|
||||
return newlst
|
||||
|
||||
class CFGContext:
|
||||
def __init__(self, sink:UOp):
|
||||
# there are 3 relationships between ranges:
|
||||
# nested, meaning endrange y is a dependency of endrange x and range x is a dependency of endrange y
|
||||
# dependent, meaning endrange y is a dependency of endrange x and range x is not a dependency of endrange y
|
||||
# independent, endrange y is not a dependency of endrange x
|
||||
# everything is nested inside the sink
|
||||
deps: dict[UOp, dict[UOp, None]] = {}
|
||||
nesting: dict[UOp, UOp] = {}
|
||||
for u in sink.toposort():
|
||||
# get the deps from the src
|
||||
deps[u] = {}
|
||||
for s in u.src: deps[u] |= deps[s]
|
||||
|
||||
if u.op in (Ops.END, Ops.SINK):
|
||||
nesting |= {x:u for x in deps[u] if x.op is Ops.END and (u.op is Ops.SINK or u.src[1] in deps[x]) and x not in nesting}
|
||||
if u.op in (Ops.RANGE, Ops.END): deps[u][u] = None
|
||||
|
||||
self.edges: dict[UOp, UOp] = {}
|
||||
siblings: dict[UOp, list[UOp]] = {}
|
||||
for k,vv in nesting.items(): siblings.setdefault(vv, []).append(k)
|
||||
for k,v in siblings.items():
|
||||
# ranges that have dependencies on other siblings need to be scheduled after them
|
||||
order = sorted(v, key=lambda x: len([u for u in v if u in deps[x]]))
|
||||
zipped = zip(order, order[1:]) if k.op is Ops.SINK else zip([k.src[1]] + order, order)
|
||||
for x,y in zipped:
|
||||
# TODO: this can happen! it causes infinite loop in shufflenet
|
||||
assert y.src[1] not in x.backward_slice_with_self
|
||||
self.edges[y.src[1]] = x
|
||||
|
||||
pm_add_control_flow = PatternMatcher([
|
||||
(UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(src=x.src+(y,)) if (y:=ctx.edges.get(x)) is not None else None),
|
||||
])
|
||||
|
||||
def do_split_ends(e:UOp):
|
||||
ret = e.src[0]
|
||||
for r in sorted(UOp.sink(*e.src[1:]).ranges, key=lambda x: x.arg, reverse=True): ret = ret.end(r)
|
||||
return ret
|
||||
|
||||
pm_split_ends = PatternMatcher([
|
||||
# split the ends
|
||||
(UPat(Ops.END, name="e"), do_split_ends),
|
||||
])
|
||||
@@ -0,0 +1,137 @@
|
||||
import itertools
|
||||
from tinygrad.helpers import dedup
|
||||
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
|
||||
from tinygrad.renderer.isa import ISARenderer, Register
|
||||
from tinygrad.dtype import dtypes, PtrDType
|
||||
|
||||
PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP}
|
||||
|
||||
class LinearScanRegallocContext:
|
||||
# returns the uop that defines the virtual register
|
||||
def vdef(self, v:Register) -> UOp: return self.uops[self.live_range[v][0]]
|
||||
def __init__(self, uops:list[UOp], ren:ISARenderer):
|
||||
self.uops = uops
|
||||
self.ren = ren
|
||||
self.idx = itertools.count()
|
||||
# the label associated with each loop NOTE: this is only used post regalloc and should be removed
|
||||
self.loop_label: dict[UOp, str] = {}
|
||||
|
||||
# compute live ranges
|
||||
self.live_range: dict[Register, list[int]] = {}
|
||||
lr = self.live_range
|
||||
ranges: list[Register] = []
|
||||
for i,u in enumerate(reversed(uops)):
|
||||
if u.op in PSEUDO_OPS: continue
|
||||
defs = u.tag if isinstance(u.tag, tuple) else ()
|
||||
for v in defs + tuple(s.reg for s in dedup(u.src)):
|
||||
if isinstance(v, Register): lr.setdefault(v, []).insert(0, len(uops) - 1 - i)
|
||||
for v in defs:
|
||||
if v in lr and (n:=max((lr[rng][-1] for rng in ranges if lr[rng][0] <= lr[v][-1] < lr[rng][-1]), default=None)): lr[v].append(n)
|
||||
if u.op is Ops.RANGE: ranges.append(u.reg)
|
||||
|
||||
# allocate registers
|
||||
self.stack_size: int = 0
|
||||
self.locals: dict[UOp, UOp] = {}
|
||||
self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot
|
||||
self.reals: dict[int, dict[Register, Register]] = {} # mapping from virtual to real at each program point
|
||||
self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # fills to be inserted at each program point
|
||||
live: dict[Register, Register] = {} # mapping from virtual to real that's currently assigned to it
|
||||
live_ins: list[dict[Register, Register]] = [] # mapping from virtual to real at loop entry
|
||||
|
||||
def alloc(cons:tuple[Register, ...], i:int) -> Register:
|
||||
live_inv = {v:k for k,v in live.items()}
|
||||
# allocate the best register. Registers not in live or not used again are free and have priority,
|
||||
# otherwise pick the one with the furthest next use. Regs that appear first in cons have priority in case of a tie
|
||||
reg,vreg = max(((r,live_inv.get(r)) for r in cons),
|
||||
key=lambda rv: next((j-i for j in ([] if rv[1] is None else lr[rv[1]]) if j >= i), len(uops)))
|
||||
return live.pop(vreg) if vreg is not None else reg
|
||||
|
||||
# assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot
|
||||
def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register:
|
||||
if v not in self.spills:
|
||||
dt = self.vdef(v).dtype
|
||||
sz = dt.scalar().itemsize * dt.count if not isinstance(dt, PtrDType) else 8
|
||||
offset = self.stack_size + (sz - self.stack_size % sz) % sz
|
||||
self.spills[v] = UOp.const(dtypes.int32, offset)
|
||||
self.stack_size = offset + sz
|
||||
r = alloc(cons if cons is not None else v.cons, i)
|
||||
self.insert_before.setdefault(i, []).append((v, r))
|
||||
return r
|
||||
|
||||
for i,u in enumerate(uops):
|
||||
if u.op in PSEUDO_OPS: continue
|
||||
# allocate uses
|
||||
for s in u.src:
|
||||
# HACK: cause of later hacks to lower range
|
||||
if u.op is Ops.END: continue
|
||||
if not isinstance(v:=s.reg, Register): continue
|
||||
if v not in live: live[v] = fill(v, i)
|
||||
self.reals.setdefault(i, {})[v] = live[v]
|
||||
|
||||
# allocate defs
|
||||
if isinstance(u.tag, tuple):
|
||||
for j,v in enumerate(u.tag):
|
||||
# register should only be defined once
|
||||
assert isinstance(v, Register) and lr[v][0] == i
|
||||
cons = v.cons
|
||||
# two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak
|
||||
if ren.is_two_address(u) and j == 0:
|
||||
uses = tuple(live.get(s.reg) for s in u.src)
|
||||
cons = ((uses[0],) if uses[0] in cons else ()) + tuple(r for r in cons if r not in uses)
|
||||
# HACK: cause the range is missing the comparison
|
||||
live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i)
|
||||
self.reals.setdefault(i, {})[v] = live[v]
|
||||
|
||||
# allocate stack array
|
||||
if u.op is Ops.DEFINE_LOCAL:
|
||||
self.locals[u] = UOp.const(dtypes.int32, self.stack_size)
|
||||
self.stack_size += u.dtype.nbytes()
|
||||
|
||||
# loop prologue, avoid loading inside the loop
|
||||
if u.op is Ops.RANGE:
|
||||
# we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
|
||||
used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < lr[u.reg][-1] for l in lr[v])]
|
||||
sorted_uses = sorted(used_in_loop, key=lambda k: (next(l-i for l in lr[k] if l >= i), lr[k][0], k.name, k.index))
|
||||
live_in: dict[Register, Register] = {}
|
||||
for v in sorted_uses:
|
||||
# if all the possible registers are already in live_in there's no space for this var
|
||||
if set(v.cons).issubset(live_in.values()): continue
|
||||
if v not in live: live[v] = fill(v, i)
|
||||
live_in[v] = live[v]
|
||||
live_ins.append(live_in)
|
||||
|
||||
# loop epilogue, reload registers that were live at loop entry
|
||||
if u.op is Ops.END:
|
||||
# TODO: if a uop is in a different reg in live out vs live in move between registers instead of loading
|
||||
# TODO: don't reload if first use in loop is a load
|
||||
for v,r in live_ins.pop().items():
|
||||
if v not in live or live[v] != r: live[v] = fill(v, i, (r,))
|
||||
|
||||
def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
|
||||
i = next(ctx.idx)
|
||||
if x.op in PSEUDO_OPS: return None
|
||||
nsrc = []
|
||||
for j,s in enumerate(x.src):
|
||||
# v here is the virtual defined by the original s as s is the rewritten version
|
||||
if i in ctx.reals and (v:=ctx.uops[i].src[j].reg) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
|
||||
else: nsrc.append(s)
|
||||
ndefs = tuple(ctx.reals[i][v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag
|
||||
if x.op is Ops.DEFINE_LOCAL: nx = ctx.ren.isel_matcher.rewrite(ctx.ren.stack_pointer().index(ctx.locals[x], dtype=x.dtype, tag=ndefs))
|
||||
else: nx = x.replace(src=tuple(nsrc), tag=ndefs)
|
||||
|
||||
before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), r) for v,r in ctx.insert_before.get(i, [])]
|
||||
after = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else []
|
||||
|
||||
# alloc/dealloc stack
|
||||
if ctx.stack_size > 0:
|
||||
sp = ctx.ren.stack_pointer()
|
||||
offset = UOp(Ops.CONST, sp.dtype, arg=ctx.stack_size)
|
||||
if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, sp.dtype, (sp, offset), tag=sp.tag))] + before
|
||||
elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, sp.dtype, (sp, offset), tag=sp.tag))]
|
||||
|
||||
return nx, before + [nx] + after
|
||||
|
||||
pm_regalloc_rewrite = PatternMatcher([
|
||||
(UPat({Ops.INS, Ops.RANGE, Ops.END, Ops.DEFINE_REG, Ops.DEFINE_LOCAL, Ops.PARAM, Ops.DEFINE_VAR, Ops.SPECIAL} | PSEUDO_OPS, name="x"),
|
||||
regalloc_rewrite),
|
||||
])
|
||||
@@ -0,0 +1,20 @@
|
||||
# opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
|
||||
from __future__ import annotations
|
||||
from enum import Enum, auto
|
||||
from dataclasses import dataclass
|
||||
|
||||
class OptOps(Enum):
|
||||
TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto(); THREAD = 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
|
||||
|
||||
@dataclass(frozen=True, order=True)
|
||||
class Opt:
|
||||
op: OptOps
|
||||
axis: int|None = None
|
||||
arg: int|tuple|None = None
|
||||
def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})"
|
||||
|
||||
class KernelOptError(Exception): pass
|
||||
def check(cond:bool, msg:str=""):
|
||||
if not cond: raise KernelOptError(msg)
|
||||
@@ -0,0 +1,191 @@
|
||||
import itertools
|
||||
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
|
||||
from tinygrad.helpers import getenv, DEBUG, prod, NOLOCALS, TC_OPT, TC_SELECT, USE_TC, IMAGE
|
||||
from tinygrad.dtype import PtrDType, ImageDType
|
||||
from tinygrad.uop.ops import Ops, resolve, AxisType
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
|
||||
def hand_coded_optimizations(k:Scheduler) -> Scheduler:
|
||||
# first try the tensor cores
|
||||
""" 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_select -- specifies which tensor core(s) to use for optimization (default -1)
|
||||
-1: iterates through all available tensor cores in order and uses the first one that matches the requirements (dims and dtypes)
|
||||
[0-N]: uses only the n'th tensor core available; useful for search
|
||||
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 Ops.LOAD into Ops.MUL
|
||||
1: allows kernels with multiple reduce axes and also multiplication of Ops.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
|
||||
"""
|
||||
# NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
|
||||
if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
|
||||
good_tc_opt = False
|
||||
tk = k.copy()
|
||||
try: # check TC first and apply hand-coded opts if successful
|
||||
rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value)))
|
||||
good_tc_opt = True
|
||||
except KernelOptError:
|
||||
pass
|
||||
# skip hand-coded TC opts if AMX, upcasting will make kernel slower
|
||||
if good_tc_opt and "AMX" not in k.ren.target.arch:
|
||||
if rngs is not None:
|
||||
for tc_dim in [1,0]: # attempt to upcast M and N
|
||||
szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
|
||||
if szs:
|
||||
# set it to the replaced range
|
||||
rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
|
||||
if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
|
||||
tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
|
||||
return tk
|
||||
|
||||
# make a copy so it does not mutate the input
|
||||
k = k.copy()
|
||||
|
||||
# upcast float4 images, this must be early so we don't accidentally add locals before the upcast
|
||||
if IMAGE:
|
||||
for buf_index,buf in enumerate(k.bufs):
|
||||
if isinstance(buf.src[0].dtype, PtrDType) and ImageDType.valid_dims(buf.src[0].dtype, k.ren.target.arch):
|
||||
# part of is_expanded
|
||||
unit_stride_axes_mul_4 = [k.rngs.index(c) for c in k.bufs[buf_index].src[1].get_idx().split_uop(Ops.ADD) if
|
||||
c.op is Ops.RANGE and (c.vmax+1)%4 == 0]
|
||||
if len(unit_stride_axes_mul_4):
|
||||
if (axis:=unit_stride_axes_mul_4[0]) in k.upcastable_dims:
|
||||
k.apply_opt(Opt(OptOps.UPCAST, axis, 4))
|
||||
elif axis in k.unrollable_dims:
|
||||
k.apply_opt(Opt(OptOps.UNROLL, k.unrollable_dims.index(axis), 4))
|
||||
|
||||
# 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 k.ren.has_local and getenv("MV",1) != 0 and (MV_BLOCKSIZE > 1 or MV_THREADS_PER_ROW > 1 or MV_ROWS_PER_THREAD > 1) and \
|
||||
k.reduceop is not None and k.reduceop.arg[0] is Ops.ADD and len(k.full_shape) >= 2 and k.ren.has_shared and \
|
||||
(mulop:=k.reduceop.src[0]).op is Ops.MUL and mulop.src[0].op is Ops.INDEX and mulop.src[1].op is Ops.INDEX:
|
||||
idx0, idx1 = mulop.src[0].src[1].get_idx(), mulop.src[1].src[1].get_idx()
|
||||
if k.ranges_of(AxisType.REDUCE):
|
||||
first_reduce_rng = k.ranges_of(AxisType.REDUCE)[0]
|
||||
if any(u is first_reduce_rng for u in idx0.split_uop(Ops.ADD)) and all(r in idx1.ranges for r in idx0.ranges):
|
||||
for global_idx in k.axes_of(AxisType.GLOBAL):
|
||||
if first_reduce_rng.src[0].divides(MV_THREADS_PER_ROW) is not None and k.full_shape[global_idx]%(MV_BLOCKSIZE*MV_ROWS_PER_THREAD) == 0:
|
||||
if DEBUG >= 3:
|
||||
print(f"MATVEC: {k.full_shape=} {first_reduce_rng.render()} {MV_BLOCKSIZE=} {MV_THREADS_PER_ROW=} {MV_ROWS_PER_THREAD=}")
|
||||
try:
|
||||
if MV_THREADS_PER_ROW > 1: k.apply_opt(Opt(OptOps.GROUP, 0, MV_THREADS_PER_ROW))
|
||||
except KernelOptError: pass
|
||||
if MV_BLOCKSIZE > 1: k.apply_opt(Opt(OptOps.LOCAL, global_idx, MV_BLOCKSIZE))
|
||||
if MV_ROWS_PER_THREAD > 1: k.apply_opt(Opt(OptOps.UPCAST, global_idx, MV_ROWS_PER_THREAD))
|
||||
return k
|
||||
|
||||
# are we grouping? (requires local shape support)
|
||||
if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
|
||||
for axis, sz in itertools.product((0, 1, 2), (16,)):
|
||||
try:
|
||||
k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
|
||||
break
|
||||
except KernelOptError: pass
|
||||
|
||||
# no more opt if we are grouping
|
||||
if k.group_for_reduces: return k
|
||||
|
||||
# **** below this line need to be optional and benchmarked ****
|
||||
|
||||
# if there are small dims with lots of valid masks, upcast them (they might be from Tensor.stack)
|
||||
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 k.upcastable_dims:
|
||||
# for Schedule, we check if the range is used in INDEX gates or WHERE gates
|
||||
is_masked = any(any(o is k.rngs[axis] for o in u.src[0].backward_slice) for u in k.ast.backward_slice if u.op is Ops.WHERE)
|
||||
if k.full_shape[axis] <= 7 and is_masked and prod(k.full_shape[j] for j in to_upcast) * k.full_shape[axis] <= 7 * 7:
|
||||
if DEBUG >= 4: print(f"upcasting masked axis : {axis}")
|
||||
to_upcast.append(axis)
|
||||
for axis in to_upcast[::-1]: k.apply_opt(Opt(OptOps.UPCAST, axis, 0))
|
||||
|
||||
# potentially do more upcasts of non reduce axes based on a heuristic
|
||||
is_dsp = k.ren is not None and k.ren.target.device == "DSP"
|
||||
upcasted_axis: set[int] = set()
|
||||
while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
|
||||
xb_choices = []
|
||||
# consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
|
||||
for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
|
||||
# if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
|
||||
if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
|
||||
rng = k.rngs[axis]
|
||||
if any(rng not in b.src[1].get_idx().backward_slice and all(r2 in b.src[1].get_idx().backward_slice
|
||||
for r2 in k.ranges_of(AxisType.UPCAST, AxisType.UNROLL)) for b in k.bufs):
|
||||
num_strides, sum_strides = 0, 0
|
||||
for b in k.bufs:
|
||||
idx = b.src[1].get_idx()
|
||||
if rng in idx.backward_slice: num_strides += 1
|
||||
for c in idx.split_uop(Ops.ADD):
|
||||
if c is rng: sum_strides += 1
|
||||
if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg
|
||||
if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg
|
||||
xb_choices.append((num_strides, sum_strides, axis, upcast_amount))
|
||||
if xb_choices:
|
||||
xb_choices = sorted(xb_choices)
|
||||
if DEBUG >= 4: print(f"more upcast axis : {xb_choices}")
|
||||
k.apply_opt(Opt(OptOps.UPCAST, xb_choices[0][2], xb_choices[0][3]))
|
||||
upcasted_axis.add(xb_choices[0][2])
|
||||
else: break
|
||||
|
||||
# if last reduce dim is small(ish), loop unroll the reduce
|
||||
# NOTE: this can fail on multireduce with mismatching dimensions, this is okay
|
||||
try:
|
||||
if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
|
||||
if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
|
||||
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
|
||||
# if it's small, upcast a second reduce dimension too
|
||||
if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
|
||||
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
|
||||
else:
|
||||
for splits in [4]:
|
||||
if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
|
||||
k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
|
||||
break
|
||||
except KernelOptError: pass
|
||||
|
||||
# if nothing at all is upcasted and it's easy to, do an upcast
|
||||
for splits in [4]:
|
||||
if not k.upcasted and k.upcastable_dims and k.full_shape[k.upcastable_dims[-1]] % splits == 0:
|
||||
k.apply_opt(Opt(OptOps.UPCAST, k.upcastable_dims[-1], splits))
|
||||
|
||||
# **** local groups ****
|
||||
|
||||
if k.ren.has_local:
|
||||
if NOLOCALS:
|
||||
k.apply_opt(Opt(OptOps.NOLOCALS))
|
||||
else:
|
||||
# prioritize making expand axes local
|
||||
local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
|
||||
for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
|
||||
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: int|None = next((x for x in ([32] * (axis == 0) + [16,8,4,3,2]) if k.full_shape[axis] % x == 0 and local_size * x <= 128), None)
|
||||
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 == k.full_shape[axis]
|
||||
k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
|
||||
if will_delete_shape: deleted_shape += 1
|
||||
|
||||
# **** threading ****
|
||||
|
||||
if k.ren.has_threads and k.ren.global_max is not None:
|
||||
for threads in [32,16,12,8,6,5,4,3,2]:
|
||||
# Skip if too many threads. Heuristic: use about 128K ops per thread
|
||||
if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
|
||||
for axis in k.axes_of(AxisType.LOOP):
|
||||
if k.full_shape[axis] % threads == 0:
|
||||
try: k.apply_opt(Opt(OptOps.THREAD, axis, threads))
|
||||
except KernelOptError: pass
|
||||
break
|
||||
if k.applied_opts and k.applied_opts[-1].op is OptOps.THREAD: break
|
||||
|
||||
return k
|
||||
@@ -0,0 +1,352 @@
|
||||
from __future__ import annotations
|
||||
import math, itertools
|
||||
from collections import defaultdict
|
||||
from typing import cast, Final
|
||||
from tinygrad.uop.ops import Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, GroupOp, remove_all_tags
|
||||
from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
|
||||
from tinygrad.device import Buffer
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.helpers import colored, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
|
||||
from tinygrad.helpers import ALLOW_TF32, count, Context
|
||||
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
|
||||
from tinygrad.codegen.simplify import pm_flatten_range
|
||||
from tinygrad.renderer import Renderer
|
||||
|
||||
class Scheduler:
|
||||
def __init__(self, ast:UOp, ren:Renderer):
|
||||
self.ast, self.ren = ast, ren
|
||||
self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
|
||||
self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
|
||||
self.opt_range = count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
|
||||
|
||||
@property
|
||||
def rngs(self):
|
||||
# always in order by axistype
|
||||
return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.vmax > 0], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
|
||||
@property
|
||||
def shape_len(self) -> int: return len(self.rngs)
|
||||
@property
|
||||
def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
|
||||
@property
|
||||
def axis_types(self) -> list[AxisType]: return [x.arg[-1] for x in self.rngs]
|
||||
|
||||
# strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
|
||||
def shape_str(self) -> list[str]:
|
||||
ret: list[str] = []
|
||||
cnt: dict[AxisType, int] = {}
|
||||
for x in self.axis_types:
|
||||
cnt[x] = (cnt[x] + 1) if x in cnt else 0
|
||||
ret.append(f"{axis_letters[x]}{cnt[x]}")
|
||||
return ret
|
||||
def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms])
|
||||
|
||||
def copy(self) -> Scheduler:
|
||||
ret = Scheduler(self.ast, self.ren)
|
||||
ret.dont_use_locals = self.dont_use_locals
|
||||
ret.applied_opts = self.applied_opts[:]
|
||||
if hasattr(self, 'tensor_core'): ret.tensor_core = self.tensor_core
|
||||
return ret
|
||||
|
||||
kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
|
||||
def get_optimized_ast(self, name_override:str|None=None) -> UOp:
|
||||
if name_override is not None: name = name_override
|
||||
else:
|
||||
k_type = "r" if self.reduceop is not None else "E"
|
||||
special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg)
|
||||
special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops]
|
||||
name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
|
||||
Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
|
||||
num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
|
||||
name += colored(num, 'BLACK')
|
||||
self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
|
||||
return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)
|
||||
|
||||
def _output_rngs(self) -> list[UOp]:
|
||||
return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
|
||||
def _globalizable_rngs(self) -> list[UOp]:
|
||||
ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
|
||||
# exclude any output ranges from global that don't appear in all BUFFERIZE
|
||||
for x in self.ast.toposort():
|
||||
if x.op is Ops.STAGE:
|
||||
ret = [r for r in ret if r in x.ranges]
|
||||
return ret
|
||||
|
||||
def convert_loop_to_global(self) -> None:
|
||||
if not self.ren.has_local: return
|
||||
|
||||
globalizible_rngs = self._globalizable_rngs()
|
||||
rng = [x.replace(arg=x.arg[0:-1]+(AxisType.GLOBAL,)) if x in globalizible_rngs else x for x in self.rngs]
|
||||
|
||||
self.ast = self.ast.substitute(dict(zip(self.rngs, rng)))
|
||||
|
||||
def colors(self) -> list[str]:
|
||||
output_rngs = self._output_rngs()
|
||||
globalizible_rngs = self._globalizable_rngs()
|
||||
ret = []
|
||||
for x,r in zip(self.axis_types, self.rngs):
|
||||
if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
|
||||
elif r not in output_rngs and x == AxisType.LOOP: ret.append("BLACK")
|
||||
elif r not in globalizible_rngs and x == AxisType.LOOP: ret.append("white")
|
||||
else: ret.append(axis_colors[x])
|
||||
return ret
|
||||
def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
|
||||
|
||||
def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng:UOp|None=None):
|
||||
if (old_sz:=rng.src[0].divides(amount)) is None:
|
||||
raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
|
||||
new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng
|
||||
replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
|
||||
sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
|
||||
self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
|
||||
return replaced_rng, new_rng
|
||||
|
||||
def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
|
||||
def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
|
||||
|
||||
def upcast_size(self): return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
|
||||
|
||||
# copied from kernel.py
|
||||
@property
|
||||
def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
|
||||
if isinstance(s:=self.full_shape[i], int) and s > 1]
|
||||
@property
|
||||
def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
|
||||
if isinstance(s:=self.full_shape[i], int) and s > 1]
|
||||
|
||||
def real_axis(self, op:OptOps, axis:int|None) -> int:
|
||||
try:
|
||||
if axis is None or op is OptOps.TC: return -1
|
||||
if op is OptOps.UNROLL: return self.unrollable_dims[axis]
|
||||
if op in {OptOps.GROUP, OptOps.GROUPTOP}: return self.axes_of(AxisType.REDUCE)[axis]
|
||||
check(axis < self.shape_len, f"invalid axis on {axis=} {op=} {self.shape_len=}")
|
||||
return axis
|
||||
except IndexError as e: raise KernelOptError from e
|
||||
|
||||
def apply_opt(self, opt:Opt, append_opt:bool=True):
|
||||
if opt.op is OptOps.NOLOCALS:
|
||||
check(all(x not in {AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE} for x in self.axis_types), "no locals can't have locals")
|
||||
if append_opt: self.applied_opts.append(opt)
|
||||
self.dont_use_locals = True
|
||||
return
|
||||
|
||||
if opt.op in {OptOps.LOCAL, OptOps.GROUP, OptOps.GROUPTOP}:
|
||||
check(self.ren.has_local, "locals needed for opt")
|
||||
|
||||
rng = self.rngs[real_axis] if (real_axis:=self.real_axis(opt.op, opt.axis)) >= 0 else UOp(Ops.NOOP)
|
||||
|
||||
opt_to_at = {
|
||||
OptOps.LOCAL: AxisType.LOCAL, OptOps.UPCAST: AxisType.UPCAST,
|
||||
OptOps.UNROLL: AxisType.UNROLL, OptOps.GROUP: AxisType.GROUP_REDUCE,
|
||||
OptOps.GROUPTOP: AxisType.GROUP_REDUCE, OptOps.THREAD: AxisType.THREAD}
|
||||
|
||||
ret = None
|
||||
if opt.op in opt_to_at:
|
||||
amt:int = int(rng.vmax+1) if opt.arg == 0 else cast(int, opt.arg)
|
||||
|
||||
# copied from kernel.py. prevents METAL compiler hangs
|
||||
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})):
|
||||
upcast_local_sz = prod([self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
|
||||
smem_sz = amt*upcast_local_sz*self.reduceop.dtype.itemsize
|
||||
check(smem_sz <= self.ren.shared_max, f"exceeds maximum shared memory size: needs {smem_sz}, max {self.ren.shared_max}")
|
||||
if self.reduceop is not None and (opt.op in {OptOps.GROUP, OptOps.GROUPTOP}):
|
||||
# We currently dont support a group within another rudece, TODO: fix if-contexts
|
||||
reduce = [u for u in self.ast.backward_slice if u.op is Ops.REDUCE and rng in merge_dicts([r.ranges for r in u.src[1:]])][0]
|
||||
check(not any(u.arg[-1] in (AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE) for u in reduce.ranges),
|
||||
"cannot have a GROUP_REDUCE inside another reduce")
|
||||
|
||||
if opt.op is OptOps.UNROLL:
|
||||
check(amt <= 32, "don't unroll more than 32")
|
||||
check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
|
||||
if opt.op is OptOps.UPCAST:
|
||||
check((self.ren is not None and self.ren.target.device == "DSP") or amt <= 16, "don't upcast more than 16")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
|
||||
if opt.op is OptOps.LOCAL:
|
||||
check(not self.dont_use_locals, "can't use locals")
|
||||
check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
|
||||
if opt.op is OptOps.THREAD:
|
||||
check(self.ren is not None and self.ren.has_threads, "target does not support threads")
|
||||
check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads")
|
||||
check(all(x is not AxisType.THREAD for x in self.axis_types), "already threaded")
|
||||
check(rng in self._globalizable_rngs(), "can't apply range to this dim")
|
||||
if opt.op in {OptOps.GROUP, OptOps.GROUPTOP}:
|
||||
check(all(x.op is not OptOps.TC for x in self.applied_opts), "no grouping with tensor cores") # TODO: why is this wrong?
|
||||
check(not self.dont_use_locals, "can't use locals")
|
||||
check(rng.arg[-1] == AxisType.REDUCE, "group is for reduce")
|
||||
ret = self.shift_to(rng, amt, opt_to_at[opt.op], top=opt.op in {OptOps.GROUPTOP, OptOps.THREAD})
|
||||
elif opt.op is OptOps.TC:
|
||||
check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps
|
||||
check(opt.axis is not None, "tensor core opts must have an axis")
|
||||
check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg")
|
||||
check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.ren.tensor_cores), "tensor core opts must have valid tc_select")
|
||||
check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt")
|
||||
check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid")
|
||||
try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt)
|
||||
except ValueError as e: raise KernelOptError(str(e))
|
||||
check(ret is not None, "no tensor core available")
|
||||
elif opt.op is OptOps.PADTO:
|
||||
check(rng.src[0].op is Ops.CONST, "only pad const axes")
|
||||
check(rng.arg[-1] not in {AxisType.UPCAST, AxisType.UNROLL}, "cannot pad upcasted") # TODO: why is this wrong?
|
||||
check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread")
|
||||
# ok to pad SUM if all parent ALU ops have f(0) = 0
|
||||
if (r:=self.reduceop) is not None and rng.arg[-1] in (AxisType.GROUP_REDUCE, AxisType.REDUCE):
|
||||
check(r.arg[0] is Ops.ADD and not r.op_in_backward_slice_with_self(*GroupOp.UnsafePad), f"cannot pad {r}")
|
||||
new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg))
|
||||
check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work")
|
||||
replaced_rng = UOp.range(new_sz, *rng.arg)
|
||||
replaces = {rng:replaced_rng}
|
||||
valid = replaced_rng < rng.vmax+1
|
||||
for b in self.bufs:
|
||||
if rng in (i:=b.src[1].get_idx()).backward_slice_with_self:
|
||||
replaces[b] = b.replace(src=(b.src[0],(valid&b.src[1].get_valid()).where(i, UOp.invalid())))
|
||||
self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
|
||||
elif opt.op is OptOps.SWAP:
|
||||
try:
|
||||
altrng:UOp = self.rngs[opt.arg]
|
||||
except IndexError:
|
||||
raise KernelOptError
|
||||
check(rng.arg[-1] == AxisType.GLOBAL and altrng.arg[-1] == AxisType.GLOBAL, "swap only for globals")
|
||||
self.ast = self.ast.substitute({rng:rng.replace(arg=(*altrng.arg[0:-1], rng.arg[-1]), tag=1),
|
||||
altrng:altrng.replace(arg=(*rng.arg[0:-1], altrng.arg[-1]), tag=1)},
|
||||
name=f"swap {rng.arg[:-1]} {altrng.arg[:-1]}")
|
||||
self.ast = graph_rewrite(self.ast, remove_all_tags, name="swap remove tags")
|
||||
else:
|
||||
raise KernelOptError(f"unsupported opt {opt.op}")
|
||||
|
||||
if append_opt: self.applied_opts.append(opt)
|
||||
return ret
|
||||
|
||||
def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]:
|
||||
if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore")
|
||||
reduceop = reduceops[0]
|
||||
if use_tensor_cores and reduceop.arg[0] is Ops.ADD:
|
||||
mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0]
|
||||
if mul.op is not Ops.MUL: return None
|
||||
in0, in1 = mul.src
|
||||
try:
|
||||
tensor_cores = self.ren.tensor_cores if tc_select == -1 else [self.ren.tensor_cores[tc_select]]
|
||||
except IndexError:
|
||||
raise KernelOptError(f"invalid tensor core choice {tc_select}")
|
||||
for tc in tensor_cores:
|
||||
if self.ren.target.device in ("CUDA", "NV") and tc.dtype_in == dtypes.float and not ALLOW_TF32: continue
|
||||
if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
|
||||
# tensor cores have three ranges. X, Y, and REDUCE
|
||||
in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)
|
||||
in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0], reverse=True)
|
||||
red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0], reverse=True)
|
||||
if DEBUG >= 3:
|
||||
print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
|
||||
f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")
|
||||
if not len(in0_ranges) or not len(in1_ranges) or not len(red_ranges): continue
|
||||
|
||||
# pick ranges
|
||||
# NOTE: why are in1 and in0 switched?
|
||||
axis_choices = list(itertools.product(in1_ranges, in0_ranges, red_ranges))
|
||||
if not (axis < len(axis_choices)): continue
|
||||
axes = list(axis_choices[axis])
|
||||
|
||||
# tag the reduceop
|
||||
self.ast = self.ast.substitute({reduceop: reduceop.replace(tag="TC")})
|
||||
|
||||
# do optimizations and save the ranges
|
||||
try:
|
||||
for i,a in enumerate(axes):
|
||||
idx = self.rngs.index(a)
|
||||
if (a.vmax+1) % tc.dims[i] != 0:
|
||||
if opt_level < 2: raise KernelOptError("tc padding requires opt_level >= 2")
|
||||
# apply_opt should return the updated range?
|
||||
self.apply_opt(Opt(OptOps.PADTO, idx, tc.dims[i]), append_opt=False) # PADTO might fail
|
||||
axes[i] = self.rngs[idx]
|
||||
except KernelOptError: continue
|
||||
|
||||
# we create the warp as a whole thing, in case some of these ranges are moved/removed later
|
||||
warp = UOp.range(tc.threads, -1, AxisType.WARP)
|
||||
ne: list[UOp] = []
|
||||
for opt in tc.opts:
|
||||
if opt[0] == "l":
|
||||
axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.LOCAL, input_new_rng=warp%2)
|
||||
warp //= 2
|
||||
elif opt[0] == "u":
|
||||
axes[int(opt[1])], new_range = self.shift_to(axes[int(opt[1])], 2, AxisType.UPCAST)
|
||||
else: raise RuntimeError(f"unsupported opt {opt[0]} in tensor cores")
|
||||
ne.append(new_range)
|
||||
|
||||
for _, amt in tc.get_reduce_axes():
|
||||
axes[2], new_range = self.shift_to(axes[2], amt, AxisType.UNROLL)
|
||||
ne.append(new_range)
|
||||
|
||||
if use_tensor_cores != 2:
|
||||
# fix the srcs
|
||||
reduceop = get_single_element([x for x in self.ast.toposort() if x.op is Ops.REDUCE and x.tag == "TC"])
|
||||
tne = [x.replace(tag=1) for x in ne]
|
||||
ret = reduceop.substitute(dict(zip(ne, tne)))
|
||||
srcs = list((ret.src[0] if ret.src[0].op is not Ops.CAST else ret.src[0].src[0]).src)
|
||||
srcs = [x.substitute(dict(zip(tne, [ne[i] for i in argsort(p)]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
|
||||
|
||||
# get reduce/upcast axes for the tensor cores
|
||||
tc_reduce_axes = self.shape_str_to_axis([f"r{i}" for i in range(len(tc.get_reduce_axes()))])
|
||||
base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())])
|
||||
tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
|
||||
|
||||
# axes to range number (was done in lowerer)
|
||||
tc_upcast_axes = tuple([tuple([(self.rngs[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes])
|
||||
tc_reduce_axes = tuple([self.rngs[a].arg[0] for a in tc_reduce_axes])
|
||||
|
||||
# construct the op
|
||||
# TODO: remove tc_upcast_axes from the arg
|
||||
# do the reduce_axes always disappear? i think they don't
|
||||
# they need to be moved into the WMMA srcs
|
||||
wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.target.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
|
||||
wmma = UOp(Ops.WMMA, dtype=tc.dtype_out.vec(tc.elements_per_thread[2]), src=(
|
||||
UOp(Ops.CONTRACT, dtype=srcs[0].dtype.vec(tc.elements_per_thread[0]), src=(srcs[0],), arg=tc_upcast_axes[0], tag=1),
|
||||
UOp(Ops.CONTRACT, dtype=srcs[1].dtype.vec(tc.elements_per_thread[1]), src=(srcs[1],), arg=tc_upcast_axes[1], tag=1),
|
||||
UOp.const(tc.dtype_out.vec(tc.elements_per_thread[2]), 0.0)), arg=wmma_arg, tag=1)
|
||||
tc_uop = UOp(Ops.UNROLL, tc.dtype_out, (wmma,), arg=tc_upcast_axes[2], tag=1)
|
||||
|
||||
# preserve extra reduces
|
||||
reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
|
||||
if len(reduce_ranges): tc_uop = UOp(Ops.REDUCE, tc_uop.dtype, (tc_uop,)+tuple(reduce_ranges), (Ops.ADD, ()))
|
||||
self.ast = self.ast.substitute({reduceop: tc_uop})
|
||||
self.tensor_core = tc
|
||||
return axes
|
||||
return None
|
||||
|
||||
# helpers for hand_coded_optimizations
|
||||
@property
|
||||
def reduceops(self) -> list[UOp]: return [x for x in self.ast.backward_slice if x.op is Ops.REDUCE]
|
||||
@property
|
||||
def reduceop(self) -> UOp|None:
|
||||
if not (red := self.reduceops): return None
|
||||
return UOp(Ops.REDUCE, red[0].dtype, red[0].src, red[0].arg)
|
||||
@property
|
||||
def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]
|
||||
@property
|
||||
def output_shape(self):
|
||||
return [s if at not in {AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE} else 1 for s,at in zip(self.full_shape, self.axis_types)]
|
||||
@property
|
||||
def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
|
||||
@property
|
||||
def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))
|
||||
|
||||
def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
|
||||
glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg)
|
||||
return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]
|
||||
|
||||
def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
|
||||
if ast.tag is not None: return ast
|
||||
k = Scheduler(ast, ren)
|
||||
k.convert_loop_to_global()
|
||||
if ast.arg is not None and ast.arg.opts_to_apply is not None:
|
||||
for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
|
||||
elif beam >= 1:
|
||||
from tinygrad.codegen.opt.search import beam_search
|
||||
rawbufs = bufs_from_ast(ast, ren.target.device)
|
||||
# beam search may open devices
|
||||
with Context(ALLOW_DEVICE_USAGE=1):
|
||||
k = beam_search(k, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
|
||||
elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
|
||||
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
|
||||
# NOTE: hand_coded_optimizations doesn't support multiblock opts yet
|
||||
if not any(u.op is Ops.STAGE for u in ast.backward_slice):
|
||||
k = hand_coded_optimizations(k)
|
||||
return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)
|
||||
@@ -0,0 +1,187 @@
|
||||
import math, time, multiprocessing, traceback, signal, atexit
|
||||
from dataclasses import replace
|
||||
from tinygrad.uop.ops import sym_infer, AxisType, UOp
|
||||
from tinygrad.uop.render import pyrender
|
||||
from tinygrad.device import Device, Buffer
|
||||
from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
|
||||
from tinygrad.helpers import IGNORE_BEAM_CACHE
|
||||
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
|
||||
from tinygrad.engine.realize import time_call
|
||||
from tinygrad.codegen import to_program
|
||||
from tinygrad.codegen.opt.postrange import Scheduler
|
||||
|
||||
actions = [Opt(op=OptOps.UPCAST, axis=axis, arg=amt) for amt in [0,2,3,4,5,7] for axis in range(8)]
|
||||
actions += [Opt(op=OptOps.UNROLL, axis=axis, arg=amt) for amt in [0,4,7] for axis in range(5)]
|
||||
actions += [Opt(op=OptOps.LOCAL, axis=axis, arg=amt) for amt in [2,3,4,8,13,16,29] for axis in range(6)]
|
||||
actions += [Opt(op=OptOps.GROUPTOP, axis=axis, arg=amt) for amt in [13,16,28,29,32,49,64,256] for axis in range(3)]
|
||||
actions += [Opt(op=OptOps.GROUP, axis=axis, arg=amt) for amt in [0,4,8,16] for axis in range(3)]
|
||||
if getenv("BEAM_PADTO", 0): actions += [Opt(op=OptOps.PADTO, axis=axis, arg=amt) for amt in [32] for axis in range(7)]
|
||||
actions += [Opt(op=OptOps.LOCAL, axis=0, arg=32), Opt(op=OptOps.LOCAL, axis=6, arg=2)]
|
||||
actions += [Opt(op=OptOps.TC, axis=0, arg=(-1, 0, getenv("TC", 1)))]
|
||||
# covers resnet kernels (3 global * 3 reduce)
|
||||
actions += [Opt(op=OptOps.TC, axis=axis, arg=(-1, getenv("TC_OPT", 2), getenv("TC", 1))) for axis in range(9)]
|
||||
actions += [Opt(op=OptOps.SWAP, axis=axis_0, arg=axis_1) for axis_0 in range(5) for axis_1 in range(axis_0+1, 5)]
|
||||
actions += [Opt(op=OptOps.THREAD, axis=axis, arg=amt) for amt in [2,3,4,5,8,12,16,24,32,64] for axis in range(3)]
|
||||
if getenv("NOLOCALS"): actions += [Opt(op=OptOps.NOLOCALS)]
|
||||
|
||||
def get_test_global_size(global_size, max_global_size, var_vals):
|
||||
test_global_size = [sym_infer(sz, var_vals) for sz in global_size]
|
||||
input_size = prod(test_global_size)
|
||||
while prod(test_global_size) > max_global_size:
|
||||
for j in range(len(global_size)-1,-1,-1):
|
||||
if test_global_size[j] > 16:
|
||||
test_global_size[j] //= 2
|
||||
break
|
||||
return test_global_size, input_size / prod(test_global_size)
|
||||
|
||||
def _time_program(prg:UOp, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
|
||||
allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test", dev_timeout=False) -> list[float]:
|
||||
timeout = int(early_stop * 1e3) if dev_timeout and early_stop is not None and early_stop < math.inf else None
|
||||
factor = 1
|
||||
if allow_test_size and max_global_size is not None:
|
||||
global_size, factor = get_test_global_size(prg.arg.global_size, max_global_size, var_vals)
|
||||
prg = prg.replace(arg=replace(prg.arg, global_size=tuple(global_size)))
|
||||
call = prg.call(*[UOp.from_buffer(b) for b in rawbufs])
|
||||
tms = []
|
||||
for _ in range(cnt):
|
||||
try: tms.append(time_call(call, var_vals, timeout=timeout, clear_l2=clear_l2) * factor)
|
||||
except AssertionError: return [math.inf] * cnt
|
||||
if early_stop is not None and early_stop < min(tms): break
|
||||
return tms
|
||||
|
||||
class TimeoutException(Exception): pass
|
||||
def timeout_handler(signum, frame):
|
||||
if DEBUG >= 2: print("*** BEAM COMPILE TIMEOUT")
|
||||
raise TimeoutException()
|
||||
|
||||
def _try_compile(x:tuple[int,Scheduler]) -> tuple[int, tuple[UOp, float]|None]:
|
||||
if hasattr(signal, "alarm"):
|
||||
signal.signal(getattr(signal, 'SIGALRM'), timeout_handler)
|
||||
# set timeout
|
||||
signal.alarm(getenv("BEAM_TIMEOUT_SEC", 10))
|
||||
ret = None
|
||||
try:
|
||||
st = time.perf_counter()
|
||||
prg = to_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren)
|
||||
et = time.perf_counter() - st
|
||||
uops = prg.src[2].src
|
||||
if len(uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
|
||||
if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many uops. {len(uops)=}, {uops_max=}")
|
||||
raise RuntimeError("too many uops")
|
||||
ret = (prg, et)
|
||||
except RuntimeError:
|
||||
if DEBUG >= 4: traceback.print_exc()
|
||||
except Exception as e:
|
||||
if getenv("BEAM_STRICT_MODE"): raise e
|
||||
finally:
|
||||
if hasattr(signal, "alarm"): signal.alarm(0)
|
||||
return x[0], ret
|
||||
|
||||
# workers should not open devices and should ignore ctrl c and should not launch VIZ
|
||||
def _init_worker():
|
||||
Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
|
||||
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
||||
|
||||
def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_allocated() if buf is not None else buf for buf in bufs]
|
||||
|
||||
# *** external API ***
|
||||
|
||||
# get dictionary of all possible actions
|
||||
def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dict[int, Scheduler]:
|
||||
acted, max_up, max_lcl = {0:s} if include_0 else {}, getenv("BEAM_UPCAST_MAX", 256) if max_up is None else max_up, getenv("BEAM_LOCAL_MAX", 1024)
|
||||
kernel_actions = actions.copy()
|
||||
|
||||
for i,a in enumerate(kernel_actions):
|
||||
if a.axis is not None and a.op is not OptOps.TC:
|
||||
try: ax = s.real_axis(a.op, a.axis)
|
||||
except KernelOptError: continue
|
||||
if (ax >= s.shape_len) or (s.full_shape[ax] == a.arg and Opt(a.op, a.axis, 0) in kernel_actions): continue
|
||||
s2 = s.copy()
|
||||
try:
|
||||
s2.apply_opt(a)
|
||||
up, lcl, tc_up = 1, 1, prod(tc.dims)//tc.threads if hasattr(s2, 'tensor_core') and (tc:=s2.tensor_core) else 1
|
||||
for x,t in zip(s2.full_shape, s2.axis_types):
|
||||
if t in (AxisType.UPCAST, AxisType.UNROLL): up *= x
|
||||
elif t in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE): lcl *= x
|
||||
if up//tc_up > max_up or lcl > max_lcl:
|
||||
if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too many upcast/local. {up//tc_up=}, {max_up=}, {lcl=}, {max_lcl=}")
|
||||
continue
|
||||
acted[i+1] = s2
|
||||
except KernelOptError: pass
|
||||
return acted
|
||||
|
||||
beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
|
||||
def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
|
||||
global beam_pool
|
||||
key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix}
|
||||
if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
|
||||
ret = s.copy()
|
||||
for o in val[len(s.applied_opts):]: ret.apply_opt(o)
|
||||
return ret
|
||||
|
||||
beam: list[tuple[Scheduler, float]] = [(s, float("inf"))]
|
||||
seen_libs = set()
|
||||
|
||||
default_parallel = multiprocessing.cpu_count() if s.ren.target.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
|
||||
if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)):
|
||||
beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
|
||||
@atexit.register
|
||||
def close_pool(): beam_pool.close()
|
||||
|
||||
min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
|
||||
if BEAM_DEBUG:
|
||||
print("BEAM_SEARCH:")
|
||||
print(pyrender(s.ast.replace(arg=None)))
|
||||
if DEBUG >= 2: print(f" 0.00s: from 1 -> 1 actions {s.colored_shape()}")
|
||||
|
||||
try:
|
||||
rawbufs = _ensure_buffer_alloc(rawbufs)
|
||||
var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in s.ast.variables()}
|
||||
exiting, st = False, time.perf_counter()
|
||||
dev = Device[s.ren.target.device]
|
||||
while not exiting:
|
||||
candidates: list[Scheduler] = flatten([get_kernel_actions(si, include_0=False).values() for si,_ in beam])
|
||||
timed: list[tuple[Scheduler, float]] = []
|
||||
least_compute_ops = math.inf
|
||||
for i, proc in ((map if beam_pool is None else beam_pool.imap_unordered)(_try_compile, enumerate(candidates))):
|
||||
if proc is None: continue
|
||||
prg, compile_et = proc
|
||||
if (lib:=prg.src[4].arg) in seen_libs: continue
|
||||
# filter out kernels that use 1000x more compute than the smallest
|
||||
estimates = prg.src[0].arg.estimates
|
||||
least_compute_ops = min(this_compute_ops:=sym_infer(estimates.ops if estimates is not None else 0, var_vals), least_compute_ops)
|
||||
if least_compute_ops*1000 < this_compute_ops:
|
||||
if getenv("BEAM_LOG_SURPASS_MAX"): print(f"too much compute. {this_compute_ops} when least is {least_compute_ops}")
|
||||
continue
|
||||
seen_libs.add(lib)
|
||||
try: tms = _time_program(prg, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0,
|
||||
allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'),
|
||||
dev_timeout=getenv("BEAM_DEV_TIMEOUT", 1))
|
||||
except Exception as e:
|
||||
if BEAM_DEBUG: print(f"BEAM failed for opts: {candidates[i].applied_opts}\n{e}")
|
||||
if isinstance(e, RuntimeError): continue
|
||||
raise
|
||||
timed.append((candidates[i], min(tms)))
|
||||
if BEAM_DEBUG > 1:
|
||||
print(f"{time.perf_counter() - st:7.2f}s: {i:5d} {len(prg.src[2].src):5d} uops",
|
||||
f"{time_to_str(compile_et, w=12)} compile/{time_to_str(timed[-1][1], w=12)} run",
|
||||
f" {len(timed):4d}/{len(candidates):4d} {timed[-1][0].colored_shape()}")
|
||||
elif DEBUG >= 2:
|
||||
print(f"\r{time.perf_counter() - st:7.2f}s: {time_to_str(timed[-1][1], w=12)}",
|
||||
f" {len(timed):4d}/{len(candidates):4d} {timed[-1][0].colored_shape()}\033[K", end="")
|
||||
|
||||
# done
|
||||
opts = sorted(timed, key=lambda x: x[1])
|
||||
exiting = len(opts) == 0 or (opts[0][1] < min_progress) or (len(beam) > 0 and ((beam[0][1]-opts[0][1]) < min_progress))
|
||||
if not exiting: beam = opts[:amt]
|
||||
elif len(opts) > 0 and opts[0][1] < beam[0][1]: beam = opts[:1]
|
||||
if DEBUG >= 2:
|
||||
print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None),
|
||||
f"from {len(candidates):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())
|
||||
except KeyboardInterrupt as e:
|
||||
if beam_pool is not None: beam_pool.terminate()
|
||||
raise e
|
||||
|
||||
if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)
|
||||
if BEAM_DEBUG: print(f"BEAM_SEARCH: final tm={time_to_str(beam[0][1], w=0)}, applied_opts={beam[0][0].applied_opts}")
|
||||
return beam[0][0]
|
||||
@@ -0,0 +1,159 @@
|
||||
import math, functools
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.dtype import DType, dtypes
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
|
||||
dims: tuple[int,int,int] # N, M, K
|
||||
threads: int # number of threads that construct the warp
|
||||
elements_per_thread: tuple[int, int, int] # elements per-thread to load/store from A/B/C
|
||||
dtype_in: DType # dtype for A and B
|
||||
dtype_out: DType # dtype for C and D
|
||||
opts: tuple[str, ...] # ordered tuple of "ux" or "lx" specifying kernel opts to perform. "ux" upcasts dim x and "lx" localizes dim x
|
||||
# (local_swizzle, upcast_swizzle, reduce_swizzle)
|
||||
# l<num> is the num axis of the locals, similar for u<num> and upcasts, r<num> and reduces
|
||||
swizzle: tuple[tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]], tuple[tuple[str, ...], tuple[str, ...], tuple[str, ...]]]
|
||||
@functools.cache # pylint: disable=method-cache-max-size-none
|
||||
def _remaps(self) -> list[dict[str, str]]:
|
||||
local_axes, upcast_axes, reduce_axes = len(self.get_local_axes()), len(self.get_upcast_axes()), len(self.get_reduce_axes())
|
||||
fwd_st = [f"l{i}" for i in range(local_axes)] + [f"u{i}" for i in range(upcast_axes)] + [f"r{i}" for i in range(reduce_axes)]
|
||||
return [dict(zip(fwd_st, sum(s, ()))) for s in self.swizzle]
|
||||
def permutes_for_shape_str(self, shape_str:list[str]) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
ret = [[shape_str.index(remap[ss]) if ss in remap else i for i,ss in enumerate(shape_str)] for remap in self._remaps()]
|
||||
return tuple(ret[0]), tuple(ret[1])
|
||||
@functools.cache # pylint: disable=method-cache-max-size-none
|
||||
def base_shape_str(self) -> list[str]:
|
||||
ret = []
|
||||
cnt = {'u': 0, 'l': 0}
|
||||
for opt in self.opts:
|
||||
ret.append(f"{opt[0]}{cnt[opt[0]]}")
|
||||
cnt[opt[0]] += 1
|
||||
# assumes you do the UNROLL after the opts
|
||||
return ret + [f"r{i}" for i in range(len(self.get_reduce_axes()))]
|
||||
def get_reduce_axes(self): return [(i, 2) for i in range(int(math.log2(self.dims[2])))]
|
||||
def get_upcast_axes(self): return [opt for opt in self.opts if opt[0] == "u"]
|
||||
def get_local_axes(self): return [opt for opt in self.opts if opt[0] == "l"]
|
||||
def base_upcast_axes(self):
|
||||
# this is defined in the swizzle. first we use the upcast axes, then the reduce
|
||||
return ([f"r{i}" for i in range(len(self.get_reduce_axes()))] + [f"u{i}" for i in range(len(self.get_upcast_axes()))])[::-1]
|
||||
def __str__(self): return "_".join(["WMMA"] + list(map(str, self.dims)) + [self.dtype_in.name, self.dtype_out.name])
|
||||
def __post_init__(self):
|
||||
# all axes have size 2, <local> <reduce> <upcast> is the order
|
||||
local_axes, upcast_axes, reduce_axes = len(self.get_local_axes()), len(self.get_upcast_axes()), len(self.get_reduce_axes())
|
||||
assert self.dims[0] * self.dims[1] == 2**(local_axes + upcast_axes), \
|
||||
f"N({self.dims[0]}) x M({self.dims[1]}) != local({2**local_axes}) x upcast({2**upcast_axes}) with opts({self.opts})"
|
||||
assert 2**local_axes == self.threads, f"{self.threads} threads construct the warp but found {2**local_axes} in {self.opts}"
|
||||
assert 2**upcast_axes == self.elements_per_thread[2], \
|
||||
f"{self.elements_per_thread[2]} elements from C are processed per thread but found {2**upcast_axes} in {self.opts}"
|
||||
# check dims match opts
|
||||
assert self.dims[0] == 2**len(gd:=[x for x in self.opts if x[1] == '0']), f"opts wrong on dims[0], {self.dims[0]} vs {gd}"
|
||||
assert self.dims[1] == 2**len(gd:=[x for x in self.opts if x[1] == '1']), f"opts wrong on dims[1], {self.dims[1]} vs {gd}"
|
||||
# NOTE: the K opts is implictly set by the dim
|
||||
# check swizzle
|
||||
assert len(self.swizzle[0]) == 3 and len(self.swizzle[1]) == 3, "swizzle has wrong part count"
|
||||
assert len(self.swizzle[0][0]) == len(self.swizzle[1][0]) == local_axes, "local swizzle size is wrong"
|
||||
assert len(self.swizzle[0][1]) == len(self.swizzle[1][1]) == upcast_axes, "upcast swizzle size is wrong"
|
||||
assert len(self.swizzle[0][2]) == len(self.swizzle[1][2]) == reduce_axes, "reduce swizzle size is wrong"
|
||||
assert all(len(s) == local_axes+upcast_axes+reduce_axes for s in self._remaps()), "remaps are the wrong size"
|
||||
# check elements_per_thread
|
||||
un, ln = 0, 0
|
||||
zero_stride_0 = []
|
||||
zero_stride_1 = []
|
||||
for o in self.opts:
|
||||
if o[1] == '0': zero_stride_0.append(o[0] + str(un if o[0] == 'u' else ln))
|
||||
if o[1] == '1': zero_stride_1.append(o[0] + str(un if o[0] == 'u' else ln))
|
||||
if o[0] == 'u': un += 1
|
||||
if o[0] == 'l': ln += 1
|
||||
# NOTE: all the zero_stride dims can be placed in any order in the swizzle
|
||||
upcasted_0 = [x for x in (self.swizzle[0][1] + self.swizzle[0][2]) if x not in zero_stride_0 and x[0] != 'l']
|
||||
upcasted_1 = [x for x in (self.swizzle[1][1] + self.swizzle[1][2]) if x not in zero_stride_1 and x[0] != 'l']
|
||||
assert 2**len(upcasted_0) == self.elements_per_thread[0], f"mismatch in elements_per_thread[0], {upcasted_0} vs {self.elements_per_thread[0]}"
|
||||
assert 2**len(upcasted_1) == self.elements_per_thread[1], f"mismatch in elements_per_thread[1], {upcasted_1} vs {self.elements_per_thread[1]}"
|
||||
|
||||
# ***** NVIDIA *****
|
||||
|
||||
cuda_tc_opts = ("u0","l0","l0","l1","l1","l1","u1") # shared by all shapes with M=16 N=8
|
||||
|
||||
# https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-multiply-accumulate-instructions
|
||||
cuda_81616 = [TensorCore(dims=(8,16,16), threads=32, elements_per_thread=(8,4,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
|
||||
swizzle=((('r1', 'r2', 'l2', 'l3', 'l4'), ('u1', 'r3'), ('l0', 'l1', 'u0', 'r0')),
|
||||
(('r1', 'r2', 'u0', 'l0', 'l1'), ('r0', 'r3'), ('l2', 'l3', 'l4', 'u1'))))
|
||||
for di,do in [(dtypes.half,dtypes.float), (dtypes.bfloat16,dtypes.float), (dtypes.half,dtypes.half)]]
|
||||
cuda_81632_f8 = [TensorCore(dims=(8,16,32), threads=32, elements_per_thread=(16,8,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
|
||||
swizzle=((('r2', 'r3', 'l2', 'l3', 'l4'), ('u1', 'r4'), ('l0', 'l1', 'u0', 'r0', 'r1')),
|
||||
(('r2', 'r3', 'u0', 'l0', 'l1'), ('r1', 'r4'), ('l2', 'l3', 'l4', 'u1', 'r0'))))
|
||||
for di,do in [(dtypes.fp8e4m3,dtypes.float),(dtypes.fp8e5m2,dtypes.float)]]
|
||||
cuda_8168_f16 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,4), dtype_in=di, dtype_out=do, opts=cuda_tc_opts,
|
||||
swizzle=((('r1', 'r2', 'l2', 'l3', 'l4'), ('r0', 'u1'), ('l0', 'l1', 'u0')),
|
||||
(('r1', 'r2', 'u0', 'l0', 'l1'), ('u1', 'r0'), ('l2', 'l3', 'l4'))))
|
||||
for di,do in [(dtypes.half,dtypes.float), (dtypes.half,dtypes.half)]]
|
||||
cuda_8168_tf32 = [TensorCore(dims=(8,16,8), threads=32, elements_per_thread=(4,2,4), dtype_in=dtypes.float, dtype_out=dtypes.float, opts=cuda_tc_opts,
|
||||
swizzle=((('r0', 'r1', 'l2', 'l3', 'l4'), ('u1', 'r2'), ('l0', 'l1', 'u0')),
|
||||
(('r0', 'r1', 'u0', 'l0', 'l1'), ('u1', 'r2'), ('l2', 'l3', 'l4'))))]
|
||||
cuda_sm75: list[TensorCore] = cuda_8168_f16
|
||||
cuda_sm80: list[TensorCore] = cuda_81616 + cuda_8168_f16 + cuda_8168_tf32
|
||||
cuda_sm89: list[TensorCore] = cuda_sm80 + cuda_81632_f8
|
||||
|
||||
def get_cuda(arch): return cuda_sm89 if (ver:=int(arch[3:])) >= 89 else cuda_sm80 if ver >= 80 else cuda_sm75 if ver >= 75 else []
|
||||
|
||||
# ***** AMD *****
|
||||
|
||||
# https://gpuopen.com/learn/wmma_on_rdna3/
|
||||
amd_rdna3 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(16,16,8), dtype_in=di, dtype_out=do,
|
||||
opts=("l0","l0","l0","l0","l1","u1","u1","u1"),
|
||||
swizzle=((('l4', 'u0', 'u1', 'u2', 'l0'), ('r1', 'r2', 'r3'), ('l1', 'l2', 'l3', 'r0')),
|
||||
(('l0', 'l1', 'l2', 'l3', 'l4'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))
|
||||
for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float)]]
|
||||
amd_rdna4 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(8,8,8), dtype_in=di, dtype_out=do,
|
||||
opts=("l0","l0","l0","l0","u1","u1","u1","l1"),
|
||||
swizzle=((('u0', 'u1', 'u2', 'l4', 'r2'), ('r0', 'r1', 'r3'), ('l0', 'l1', 'l2', 'l3')),
|
||||
(('l0', 'l1', 'l2', 'l3', 'r2'), ('r0', 'r1', 'r3'), ('l4', 'u0', 'u1', 'u2'))))
|
||||
for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.bfloat16,dtypes.bfloat16)]]
|
||||
|
||||
# https://gpuopen.com/learn/amd-lab-notes/amd-lab-notes-matrix-cores-readme
|
||||
amd_cdna_161616 = [TensorCore(dims=(16,16,16), threads=64, elements_per_thread=(4,4,4), dtype_in=di, dtype_out=do,
|
||||
opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
|
||||
swizzle=((('u0', 'u1', 'l4', 'l5', 'r2', 'r3'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3')),
|
||||
(('l0', 'l1', 'l2', 'l3', 'r2', 'r3'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1'))))
|
||||
for di,do in [(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
|
||||
|
||||
amd_cdna_161632 = [TensorCore(dims=(16,16,32), threads=64, elements_per_thread=(8,8,4), dtype_in=di, dtype_out=do,
|
||||
opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
|
||||
swizzle=((('u0', 'u1', 'l4', 'l5', 'r3', 'r4'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3', 'r2')),
|
||||
(('l0', 'l1', 'l2', 'l3', 'r3', 'r4'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1', 'r2'))))
|
||||
for di,do in [(dtypes.fp8e5m2,dtypes.float),(dtypes.fp8e4m3,dtypes.float),(dtypes.half,dtypes.float),(dtypes.bfloat16,dtypes.float)]]
|
||||
|
||||
amd_cdna_1616128 = [TensorCore(dims=(16,16,128), threads=64, elements_per_thread=(32,32,4), dtype_in=di, dtype_out=do,
|
||||
opts=("l0","l0","l0","l0","u1","u1","l1","l1"),
|
||||
swizzle=((('u0', 'u1', 'l4', 'l5', 'r5', 'r6'), ('r0', 'r1'), ('l0', 'l1', 'l2', 'l3', 'r2', 'r3', 'r4')),
|
||||
(('l0', 'l1', 'l2', 'l3', 'r5', 'r6'), ('r0', 'r1'), ('l4', 'l5', 'u0', 'u1', 'r2', 'r3', 'r4'))))
|
||||
for di,do in [(dtypes.fp8e5m2,dtypes.float),(dtypes.fp8e4m3,dtypes.float)]]
|
||||
|
||||
amd_cdna3 = amd_cdna_161632[:2] + amd_cdna_161616
|
||||
|
||||
amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616
|
||||
|
||||
def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)
|
||||
|
||||
# ***** Apple Metal *****
|
||||
|
||||
metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,
|
||||
opts=("u0","l0","l1","l1","l0","l1"),
|
||||
swizzle=((('r1', 'l1', 'l2', 'r2', 'l4'), ('r0',), ('u0', 'l0', 'l3')),
|
||||
(('l0', 'r0', 'r1', 'l3', 'r2'), ('u0',), ('l1', 'l2', 'l4'))))
|
||||
for di,do in [(dtypes.float,dtypes.float),(dtypes.half,dtypes.float),
|
||||
(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.bfloat16,dtypes.bfloat16)]]
|
||||
|
||||
# ***** Apple AMX *****
|
||||
|
||||
amx = [TensorCore(dims=(sz,sz,1), threads=1, elements_per_thread=(sz,sz,sz*sz), dtype_in=dt, dtype_out=dt,
|
||||
swizzle=(((), ('u0', 'u1', 'u2', 'u3', 'u4', 'u5', 'u6', 'u7'), ()),
|
||||
((), ('u4', 'u5', 'u6', 'u7', 'u0', 'u1', 'u2', 'u3'), ())),
|
||||
opts=("u0","u0","u0","u0","u1","u1","u1","u1")) for dt,sz in [(dt, 64 // dt.itemsize) for dt in [dtypes.float]]]
|
||||
|
||||
# ***** Intel ****
|
||||
|
||||
intel = [TensorCore(dims=(8,8,16), threads=8, elements_per_thread=(16,16,8), dtype_in=dtypes.half, dtype_out=dtypes.float,
|
||||
opts=("l0","l0","l0","u1","u1","u1"),
|
||||
swizzle=((('r1', 'r2', 'r3'), ('u0', 'u1', 'u2'), ('l0', 'l1', 'l2', 'r0')),
|
||||
(('l0', 'l1', 'l2'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))]
|
||||
@@ -0,0 +1,157 @@
|
||||
import itertools
|
||||
from typing import Callable
|
||||
from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
|
||||
from tinygrad.uop.symbolic import symbolic
|
||||
from tinygrad.helpers import partition
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
def flatten_range(r:UOp) -> UOp|None:
|
||||
off = range_start[r.op]
|
||||
rngs = r.src[off:]
|
||||
if not len(rngs): return None
|
||||
new_rngs = [x for x in UOp.sink(*rngs).toposort() if x.op is Ops.RANGE]
|
||||
return r.replace(src=r.src[:off]+tuple(new_rngs))
|
||||
|
||||
pm_flatten_range = PatternMatcher([
|
||||
# real ranges only
|
||||
(UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range),
|
||||
])
|
||||
|
||||
# index/range arithmetic uses FLOORDIV/FLOORMOD prior to late rewrite
|
||||
def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.FLOORDIV, Ops.FLOORMOD} for u in x.backward_slice)
|
||||
def simplify_merge_adjacent(u:UOp) -> UOp|None:
|
||||
reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
|
||||
# on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
|
||||
for r0, r1 in (zip(u.ended_ranges, u.ended_ranges[1:]) if u.op is Ops.END else itertools.permutations(u.ended_ranges, 2)):
|
||||
# check same type
|
||||
if r0.arg[-1] == r1.arg[-1]:
|
||||
# check if the ranges to merge are in the same reduces
|
||||
if all((r0 in rngs) == (r1 in rngs) for rngs in reduce_ranges):
|
||||
s0, s1 = r0.src[0], r1.src[0]
|
||||
# do the merge
|
||||
new_range = r0.replace(src=(s0*s1,))
|
||||
nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
|
||||
name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
|
||||
|
||||
# check if it simplifies
|
||||
if count_divmod(nidx) <= count_divmod(u):
|
||||
u = nidx
|
||||
return u
|
||||
|
||||
def mark_gated(ctx, idx):
|
||||
if idx.src[1].op is Ops.WHERE:
|
||||
x, cond = idx.src[1].get_idx(), idx.src[1].get_valid()
|
||||
# get all ranges r with guards "r < c" for some const c
|
||||
guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST}
|
||||
else: x, guards = idx, {}
|
||||
# ensure that we choose max(c_i) for all i where r < c_i
|
||||
ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].arg < c.arg)}
|
||||
# but if a range is ever ungated, we cannot shrink it
|
||||
ctx |= {r:r.src[0] for r in x.ranges if r not in guards}
|
||||
|
||||
pm_simplify_ranges = PatternMatcher([
|
||||
(UPat((Ops.END, Ops.REDUCE), name="u"), simplify_merge_adjacent),
|
||||
(UPat(Ops.INDEX, name="idx"), mark_gated),
|
||||
# reduce ranges can't be shrunk
|
||||
(UPat(Ops.REDUCE, name="red"), lambda ctx, red: ctx.update({r:r.src[0] for r in red.src[1:]})),
|
||||
(UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x, lambda r,c: r.replace(src=(c,)))),
|
||||
])
|
||||
|
||||
def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
|
||||
if r not in ctx and r.arg[-1] is not AxisType.WARP and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
|
||||
|
||||
def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
|
||||
ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
|
||||
ctx.clear()
|
||||
return None if ret is x else ret.simplify()
|
||||
|
||||
pm_split_ranges = PatternMatcher([
|
||||
(UPat(Ops.RANGE, name="r")%UPat.cvar("c"), mark_range_mod),
|
||||
(UPat(Ops.SINK, name="x"), lambda ctx, x: do_substitute(ctx, x,
|
||||
lambda k,v: k.replace(src=(k.src[0]//v,), arg=k.arg[0:-1]+(0,k.arg[-1]))*v + k.replace(src=(v,), arg=k.arg[0:-1]+(1,k.arg[-1])))),
|
||||
])
|
||||
|
||||
# **** reduce simplification ****
|
||||
|
||||
def no_range(u:UOp) -> bool: return not any(x.op is Ops.RANGE for x in u.backward_slice_with_self)
|
||||
|
||||
def reduce_unparented(red:UOp) -> UOp|None:
|
||||
if red.arg[0] not in {Ops.ADD, Ops.MAX, Ops.MUL}: return None
|
||||
assert all(x.op is Ops.RANGE for x in red.src[1:]), "some reduce srcs aren't ranges"
|
||||
reduce_parented, reduce_unparented = partition(red.src[1:], lambda x: x in red.src[0].ranges)
|
||||
if len(reduce_unparented) == 0: return None
|
||||
ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
|
||||
if red.arg[0] is Ops.ADD:
|
||||
for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
|
||||
if red.arg[0] is Ops.MUL:
|
||||
for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype.scalar()).broadcast(ret.dtype.count)
|
||||
return ret
|
||||
|
||||
pm_reduce_unparented = PatternMatcher([
|
||||
# remove any ranges from a REDUCE that aren't referenced in the reduce source
|
||||
(UPat(Ops.REDUCE, name="red"), reduce_unparented),
|
||||
])
|
||||
|
||||
pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
|
||||
# lift x+y out of reduce on lt
|
||||
((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
|
||||
# lift x*y out of reduce
|
||||
((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
|
||||
lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
|
||||
# fold the range
|
||||
# bound from below
|
||||
((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(0, UPat.var("val")).reduce(UPat.var("r"), arg=Ops.ADD),
|
||||
lambda r,cut,val: (r.src[0]-cut).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
|
||||
# bound from two sides
|
||||
(((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0).reduce(UPat.var("r"),
|
||||
arg=Ops.ADD), lambda r,lower,upper,val:
|
||||
(upper.minimum(r.src[0])-lower.maximum(0)).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
|
||||
# bound from above
|
||||
((UPat(Ops.RANGE, name="r") < UPat.var("cut")).where(UPat.var("val"), 0).reduce(UPat.var("r"), arg=Ops.ADD),
|
||||
lambda r,cut,val: cut.maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
|
||||
# REDUCE on ADD
|
||||
((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
|
||||
lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
|
||||
# AND on WHERE
|
||||
((UPat(Ops.DEFINE_VAR, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
|
||||
lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
|
||||
# MUL casted bool
|
||||
((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
|
||||
])+symbolic
|
||||
|
||||
pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
|
||||
# lift x+y out of reduce on ne
|
||||
((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
|
||||
# reduce on gated load becomes can substitute the range and remove the reduce
|
||||
((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
|
||||
lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),
|
||||
])
|
||||
|
||||
def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp|None:
|
||||
for r in red.src[1:]:
|
||||
included = u.toposort(gate=lambda x: r in x.ranges)
|
||||
if any(x.op in {Ops.STORE, Ops.REDUCE} for x in included): return None
|
||||
replaces: dict[UOp, UOp] = {}
|
||||
for u in included:
|
||||
for s in u.src:
|
||||
if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR}: continue
|
||||
replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
|
||||
collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
|
||||
sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
|
||||
if not no_range(sink): return None
|
||||
u = sink.substitute({v:k for k,v in replaces.items()})
|
||||
return u
|
||||
|
||||
def reduce_load_collapse(red:UOp, u:UOp) -> UOp|None: return reduce_collapse(red, u, pm=pm_reduce_load_collapse)
|
||||
|
||||
# remove REDUCE without loads (generic arange opt / indexing).
|
||||
pm_reduce_simplify = pm_reduce_unparented + PatternMatcher([
|
||||
(UPat(Ops.REDUCE, src=(UPat.var("u"),), allow_any_len=True, arg=(Ops.ADD, ()), name="red"), reduce_collapse),
|
||||
])
|
||||
# remove REDUCE on load, comes from indexing a tensor with another tensor
|
||||
def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward_slice_with_self)
|
||||
pm_load_collapse = PatternMatcher([
|
||||
(UPat(Ops.REDUCE, arg=(Ops.ADD, ()), src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
|
||||
# we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
|
||||
((UPat.var("x", dtypes.weakint)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
|
||||
])
|
||||
Reference in New Issue
Block a user