diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index a3eb42bf0d..65625a4ccb 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -134,7 +134,7 @@ def fix_group_for_reduce(x:UOp): # 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] + reduce_loop = [x.replace(arg=(x.arg[0]+"_gfr", 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) diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index c9dad5c85d..0eec0bd874 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -30,7 +30,7 @@ class Scheduler: @property def axis_types(self): return [x.arg[-1] for x in self.rngs] @property - def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0) + def maxarg(self): return max([int(x.arg[0]) for x in self.rngs], default=0) # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2'] def shape_str(self) -> list[str]: @@ -231,9 +231,9 @@ class Scheduler: for tc in tensor_cores: 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]) - in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: -x.arg[0]) - red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0]) + 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]}") diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 3e3514e7cb..12137438a3 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -47,7 +47,7 @@ def do_substitute(ctx, x: UOp): subs = {} for k,v in ctx.items(): if v is not None: - subs[k] = 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])) + subs[k] = k.replace(src=(k.src[0]//v,), arg=(k.arg[0]+"_0", k.arg[-1]))*v + k.replace(src=(v,), arg=(k.arg[0]+"_1", k.arg[-1])) if not len(subs): return None ret = x.substitute(subs).simplify() ctx.clear() @@ -152,7 +152,7 @@ def cut_store_range(ctx, store:UOp, r:UOp): if r.src[0].op is not Ops.CONST or ctx!="CPU": return None if not (cuts:=[c.src[1].arg for c in store.get_consumer_map()[r] if c.op is Ops.CMPLT and r is c.src[0] and c.src[1].op is Ops.CONST]): return None cuts = sorted(dedup([0] + cuts + [r.src[0].arg])) - ranges = [UOp.range((end-start), *(r.arg[0:-1]+(i,r.arg[-1]))) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))] + ranges = [UOp.range((end-start), r.arg[0]+f"_{i}", r.arg[-1]) for i,(start,end) in enumerate(zip(cuts[:-1], cuts[1:]))] return UOp.group(*[store.substitute({r: new_r+start}).end(new_r) for new_r, start in zip(ranges, cuts[:-1])]) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index c989da46b5..567be2fcdf 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -412,7 +412,7 @@ def renumber_range(ctx:LocalAddBufferContext, r:UOp): if r.arg[-1] == AxisType.OUTER: # for outer range, we replace with a bound variable return UOp.variable("range_"+range_str(r), r.vmin, r.vmax).bind(r.replace(tag=None)) - ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None) + ret = r.replace(arg=(str(ctx.range),)+r.arg[1:], tag=None) ctx.range += 1 return ret diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 2f5877e772..243b2d4059 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -49,8 +49,8 @@ def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop def range_str(u:UOp, color=False) -> str: - ret = '_'.join([str(x) if x >= 0 else "m"+str(-x) for x in u.arg[0:-1]]) - return colored(ret, axis_colors[u.arg[-1]]) if color else ret + assert len(u.arg) == 2 + return colored(u.arg[0], axis_colors[u.arg[-1]]) if color else u.arg[0] def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str: ret = ','.join([range_str(x, color=color) for x in sorted(rngs, key=lambda x: x.arg)]) @@ -86,6 +86,7 @@ class UOpMetaClass(type): if _buffer is not None: assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}" buffers[created] = _buffer + if op is Ops.RANGE: assert isinstance(arg[0], str) if SPEC > 1: from tinygrad.uop.spec import full_spec, test_pyrender if SPEC > 2: test_pyrender(created) @@ -425,8 +426,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass): if shape is not None: ret = ret.reshape((1,)*len(shape)).expand(shape) return ret @staticmethod - def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.index, src=(), **kwargs): - return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs) + def range(end:sint, axis_id:str|int, axis_type=AxisType.LOOP, dtype=dtypes.index, src=(), **kwargs): + assert isinstance(axis_type, AxisType), f"{axis_type} must be an AxisType" + return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(str(axis_id), axis_type), **kwargs) @staticmethod def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, dtype=dtype, src=(sint_to_uop(end, dtype),), arg=name) def r(self, op:Ops, axis:tuple[int, ...]): @@ -1350,7 +1352,7 @@ pm_pyrender_extra = PatternMatcher([ (UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}.r({r.arg[0]}, {r.arg[1]})"), # NOTE: range has srcs sometimes after control flow (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c: - "UOp.range("+', '.join([str(c.arg)] + [str(y) for y in x.arg])+ + "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+ (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+")"), # TODO: index shouldn't mismatch dtype (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 8fcd09a041..a6fc4cffb5 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -37,8 +37,7 @@ shared_spec = PatternMatcher([ # RANGE can be in the big graph now (UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x: - rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \ - all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)), + rng.dtype == x.dtype and isinstance(rng.arg, tuple) and len(rng.arg) == 2 and isinstance(rng.arg[0], str) and isinstance(rng.arg[-1], AxisType)), (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None), # RANGE/SPECIAL define loops, END closes them