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

tuple ish

Browse Source
This commit is contained in:
geohot
2025-08-27 20:39:34 -07:00
parent df660ce904
commit 82504bc5ea
4 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/__init__.py
+1
View File
@@ -69,6 +69,7 @@ def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVEC
if _POSTOPT or _RANGEIFY:
ret.append(RewriteStep(pm_postrange_opt_early, ctx=lambda _: ({}, opts), name="early post opt ast"))
ret.append(RewriteStep(sym, name="mid symbolic"))
ret.append(RewriteStep(pm_postrange_opt, ctx=lambda _: opts, name="post optimize ast"))
# ** expander (expand_rewrite) **
tinygrad/codegen/late/expander.py
+1 -1
View File
@@ -152,7 +152,7 @@ def fix_group_for_reduce(x:UOp):
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),)) \
lambda r: UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s:=r.vmax+1), tuple(range(s))),), ((r.arg[0:-1],s),)) \
if r.arg[-1] in {AxisType.UNROLL, AxisType.UPCAST} else None),
# fix REDUCEs with UNROLLs
(UPat(Ops.REDUCE, name="x"), fix_reduce_unroll),
tinygrad/codegen/lowerer.py
+1 -1
View File
@@ -83,5 +83,5 @@ pm_lowerer = PatternMatcher([
# axis fixups for WMMA
(UPat((Ops.CONTRACT, Ops.UNROLL), name="x"),
lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0], sz) for a,sz in x.arg])) if x.tag is None else None),
lambda ctx,x: x.replace(tag=1, arg=tuple([(ctx.idxs[a].arg[0:-1], sz) for a,sz in x.arg])) if x.tag is None else None),
])
tinygrad/codegen/opt/postrange.py
+4 -4
View File
@@ -61,8 +61,8 @@ def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp,
srcs = [x.substitute(dict(zip(tne, [ne[i] for i in p]))) for x,p in zip(srcs, tc.permutes_for_shape_str(tc.base_shape_str()))]
ned = dict(zip(tc.base_shape_str(), ne))
tc_reduce_axes = tuple([ned[f"r{i}"].arg[0] for i in range(len(tc.get_reduce_axes()))])
base_upcast_axes = tuple([(ned[s].arg[0], 2) for s in tc.base_upcast_axes()])
tc_reduce_axes = tuple([ned[f"r{i}"].arg[0:-1] for i in range(len(tc.get_reduce_axes()))])
base_upcast_axes = tuple([(ned[s].arg[0:-1], 2) for s in tc.base_upcast_axes()])
tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)])
# construct the op
@@ -100,12 +100,12 @@ axis_typemap = { # (is_reduce, is_local)
def split_range(r:UOp):
if r.arg[-1] not in {AxisType.LOOP, AxisType.GLOBAL, AxisType.REDUCE}: return None
if len(r.arg) > 2: return None
if r.tag is not None: return None
# any divisor is an option
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]))
sr = r.replace(src=(rd,), arg=r.arg[0:-1]+(0, r.arg[-1]), tag=1)
er = UOp(Ops.RANGE, dtypes.int, src=(UOp.const(dtypes.int, N),), arg=r.arg[0:-1]+(1, axis_typemap[(r.arg[-1] is AxisType.REDUCE, False)]))
return sr*N+er