sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

some hand coded opts for postrange

Browse Source
This commit is contained in:
geohot
2025-08-27 16:08:45 -07:00
parent cb5295168d
commit e863b2ea6f
5 changed files with 43 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/__init__.py
+7 -6
View File
@@ -56,19 +56,20 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
# view pushing
ret.extend(rewrites_for_views)
# this is kernel.py
if not _RANGEIFY: ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
if not _POSTOPT and not _RANGEIFY: ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
if not _POSTOPT and not _RANGEIFY:
# this is kernel.py
ret.append(RewriteStep(pm_get_optimization, ctx=lambda _: opts, name="get optimization"))
ret.append(RewriteStep(pm_do_optimize, ctx=lambda _: opts, name="optimize ast"))
if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
# symbolic before post opt
ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
if _POSTOPT or _RANGEIFY: ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
# ** expander (expand_rewrite) **
ret.append(RewriteStep(sym+migrate_indexing, name="initial symbolic"))
# expand
ret.append(RewriteStep(sym+pm_pre_expander+expander, name="expander"))
# add locals
tinygrad/codegen/gpudims.py
+2 -2
View File
@@ -59,8 +59,8 @@ def add_gpudims(ctx:Renderer, s:UOp):
all_ranges = {x.arg[0]%1000:x for x in s_topo if x.op is Ops.RANGE}
# extract global/local dims
global_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[1] is AxisType.GLOBAL]))
local_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
global_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[-1] is AxisType.GLOBAL]))
local_dims = sorted(dedup([x.arg[0]%1000 for x in all_ranges.values() if x.arg[-1] in (AxisType.LOCAL, AxisType.GROUP_REDUCE)]))
if not global_dims and not local_dims: return None
# get global and local shape
tinygrad/codegen/late/expander.py
+3 -3
View File
@@ -134,11 +134,11 @@ def fix_store_unroll(x:UOp):
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)
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]
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))
@@ -153,7 +153,7 @@ pm_pre_expander = PatternMatcher([
# rewrite UPCAST/UNROLL range to something to be expanded
(UPat(Ops.RANGE, name="r"),
lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0],s),)) \
if r.arg[1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
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),
tinygrad/codegen/opt/postrange.py
+30 -2
View File
@@ -1,7 +1,9 @@
from typing import cast
from dataclasses import replace
from tinygrad.dtype import dtypes, AddrSpace, PtrDType
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, KernelInfo
from tinygrad.helpers import colored
from tinygrad.codegen.opt.kernel import axis_colors
from tinygrad.codegen.opt.kernel import axis_colors, AxisType
def rename_sink(s:UOp):
if s.arg is not None and s.arg.name != "test": return None
@@ -10,9 +12,35 @@ def rename_sink(s:UOp):
rngs = sorted([u for u in s.parents if u.op is Ops.RANGE], key=lambda x: x.arg[0:-1])
# add name to kernel
name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[1]]) for x in rngs])
name = "k" + colored('_', 'BLACK').join(['']+[colored(x.src[0].render(), axis_colors[x.arg[-1]]) for x in rngs])
return s.replace(arg=KernelInfo(name=name) if s.arg is None else replace(s.arg, name=name))
def global_stores_are_global(s:UOp):
if cast(PtrDType, s.src[0].dtype).addrspace != AddrSpace.GLOBAL: return None
return s.substitute({u:u.replace(arg=u.arg[0:-1]+(AxisType.GLOBAL,)) for u in s.src[2:] if u.op is Ops.RANGE and u.arg[-1] is AxisType.LOOP})
def split_range(r:UOp):
if len(r.arg) > 2: return None
N = 4
rd = r.src[0].divides(N)
if rd is None: return None
sr = r.replace(src=(rd,), arg=r.arg[0:-1]+(0,r.arg[-1]))
er = UOp(Ops.RANGE, dtypes.int, src=(UOp.const(dtypes.int, N),),
arg=r.arg[0:-1]+(1,AxisType.UNROLL if r.arg[-1] is AxisType.REDUCE else AxisType.UPCAST))
return sr*N+er
def flatten_range_in_terminators(r:UOp):
off = 2 if r.op is Ops.STORE else 1
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_postrange_opt = PatternMatcher([
# flatten ranges
(UPat((Ops.REDUCE, Ops.STORE), name="r"), flatten_range_in_terminators),
(UPat(Ops.RANGE, name="r"), split_range),
(UPat(Ops.STORE, name="s"), global_stores_are_global),
(UPat(Ops.SINK, name="s"), rename_sink),
])
tinygrad/viz/serve.py
+1 -1
View File
@@ -80,7 +80,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None:
label += f"\n{shape_to_str(u.shape)}"
elif len(rngs:=u.ranges):
label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[1]]) for x in sorted(rngs, key=lambda x: x.arg[0])])})"
label += f"\n({','.join([colored(str(x.arg[0]), axis_colors[x.arg[-1]]) for x in sorted(rngs, key=lambda x: x.arg[0:-1])])})"
except Exception:
label += "\n<ISSUE GETTING LABEL>"
if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}"