diff --git a/tinygrad/codegen/rangeify.py b/tinygrad/codegen/rangeify.py index b9b61fa44a..e63e20e11d 100644 --- a/tinygrad/codegen/rangeify.py +++ b/tinygrad/codegen/rangeify.py @@ -1,7 +1,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, KernelInfo, GroupOp, AxisType, TRACK_MATCH_STATS from tinygrad.opt.kernel import axis_colors, Opt, OptOps from dataclasses import dataclass -from tinygrad.dtype import dtypes +from tinygrad.dtype import dtypes, AddrSpace from tinygrad.helpers import argsort, colored, prod, all_same @dataclass @@ -84,20 +84,44 @@ def capture_sink(ctx:RangeifyContext, x: UOp): if k.op is Ops.CHILDREN and all([vi.op is Ops.INDEX for vi in v]): idxs = list(zip(*[vi.src[1:] for vi in v])) new_idxs = [] + only_new_idxs = [] + save_shape = [] + full_shape = [] for idx in idxs: if all_same(idx): - new_idxs.append(idx) + new_idxs.append(idx[0]) + save_shape.append(1) + full_shape.append(idx[0].src[0].arg) else: ll = [z.vmax+1 for z in idx] assert all_same(ll) + save_shape.append(ll[0]) + full_shape.append(ll[0]) new_idxs.append(UOp.range(dtypes.int, ll[0], (ctx.idx, AxisType.LOOP))) - return None + only_new_idxs.append(new_idxs[-1]) + ctx.idx += 1 + new_idxs = tuple(new_idxs) + inp = k.src[0] + print(save_shape, full_shape) + if len(save_shape): + buf = UOp(Ops.DEFINE_REG, inp.dtype.ptr(size=prod(save_shape), addrspace=AddrSpace.REG), arg=(ctx.idx,)) + buf = buf.reshape(tuple(save_shape)).expand(tuple(full_shape)) + store = UOp(Ops.INDEX, buf.dtype, (buf,)+new_idxs).store(UOp(Ops.INDEX, inp.dtype, (inp,)+new_idxs), *only_new_idxs, tag=1) + for vi in v: + late_subs[vi] = UOp(Ops.INDEX, buf.dtype, (buf,)+vi.src[1:]).load(store) + else: + print("no replace") + for vi in v: + assert new_idxs == vi.src[1:] + late_subs[vi] = UOp(Ops.INDEX, inp.dtype, (inp,)+new_idxs) + if not len(late_subs): return None + return x.substitute(late_subs) if x.arg is not None and x.arg.opts_to_apply is not None: ctx.opts = x.arg.opts_to_apply replace_children = {} for k,v in x.get_children_map().items(): if k.op not in {Ops.CHILDREN, Ops.DEVICE} and len(v) > 1: replace_children[k] = UOp(Ops.CHILDREN, dtype=k.dtype, src=(k.replace(tag=len(v)),)) - #if TRACK_MATCH_STATS > 0: x = x.substitute(replace_children) + if TRACK_MATCH_STATS > 0: x = x.substitute(replace_children) return x.replace(arg=None, tag=1) pm_rangeify = PatternMatcher([ diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 867b1b61b9..483df693c7 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -472,5 +472,5 @@ sym = symbolic_flat+PatternMatcher([ # move const multiply after REDUCE (NOTE: the mul chain can do this, but only if it's a same dtype reduce) ((UPat.var("x")*UPat.cvar("c", vec=False)).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg), # reduce mul chain, move muls after the reduce - (UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain), + #(UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain), ])