diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 45260e4e3a..e19d3a98de 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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) ** diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index ac8efc60e9..c6d557bbbd 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -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), diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index f38c3dc174..657993a7c8 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -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), ]) diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 99fe219230..cf85bd3036 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -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