diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index b2b23fc02d..680bd4608d 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -247,7 +247,7 @@ class CStyleLanguage(Renderer): (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))): r[u] = l else: - if u.op not in {Ops.RANGE, Ops.STORE, Ops.BUFFER} and u.dtype != dtypes.void: + if u.op not in {Ops.RANGE, Ops.BUFFER} and u.dtype != dtypes.void: l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "") kernel.append("\n".join(" "*depth + line for line in l.split("\n"))) if prefix: c[prefix] += 1 # if it was used, increment diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 9a7ea339af..a5d532ead2 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -80,7 +80,6 @@ class PythonProgram(Program['PythonDevice']): # in the python emulator, the warp is always in sync i += 1 continue - assert u.dtype is not None, f"{u.op} is missing a dtype" if u.op is Ops.STORE: assert len(src_values) == 2, f"STORE must be lowered to 2 srcs, got {len(src_values)}" store_gate = exec_masks[-1] diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index e77fdd5d4a..65305dacad 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -68,8 +68,8 @@ spec_shared = PatternMatcher([ (UPat(GroupOp.Comparison, dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: matches_dtype(x, y.dtype) or matches_dtype(y, x.dtype) or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks), (UPat((Ops.AND, Ops.OR, Ops.XOR, Ops.SHL, Ops.SHR), name="x"), lambda x: False if any(dtypes.is_float(s.dtype) for s in x.src) else None), - (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("c")), name="a"), lambda a,x,c: (matches_dtype(x, a.dtype) or x.dtype is dtypes.weakint) - and (matches_dtype(c, a.dtype) or c.dtype in (dtypes.uint, dtypes.weakint) or x.base.is_invalid)), + (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("c")), name="a"), lambda a,x,c: + matches_dtype(c, a.dtype) or c.dtype in (dtypes.uint, dtypes.weakint) or x.base.is_invalid), (UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) or any(s.base.is_invalid for s in x.src) else False), (UPat(GroupOp.ALU, name="x"), lambda x: all(matches_dtype(y, x.dtype) or y.dtype in dtypes.weaks for y in x.src)), @@ -78,8 +78,7 @@ spec_shared = PatternMatcher([ (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: isinstance(x.arg, DType)), # RANGE can be in the big graph now. a void RANGE is a bound-less loop header, the arg is an axis id like RANGE - (UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x: - matches_dtype(x, rng.dtype) and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \ + (UPat(Ops.RANGE, src=(UPat(),), allow_any_len=True, name="rng"), lambda rng: 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)), (UPat(Ops.INDEX, name="x"), lambda x: len(x.src)>0 and all(dtypes.is_int(y.dtype) or y.base.is_invalid for y in x.src[1:]) or None), # END closes RANGEs @@ -97,7 +96,7 @@ spec_shared = PatternMatcher([ # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.INDEX, Ops.AFTER, Ops.UNSHARD, Ops.BITCAST, Ops.INS})),), - allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)), + allow_any_len=True), lambda: True), # CUSTOM (inline and non inline): the arg is the source string and the dtype it produces, void for a bare statement (UPat((Ops.CUSTOMI, Ops.CUSTOM), name="x"), @@ -132,7 +131,7 @@ spec_shared = PatternMatcher([ def is_device(d): return isinstance(d, str) or (isinstance(d, tuple) and all(isinstance(s, str) for s in d)) -def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg < len(t.src) and matches_dtype(t.src[g.arg], g.dtype) +def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg < len(t.src) # these ops can exist in tensor but not programs. example: movement spec_tensor = PatternMatcher([ @@ -160,7 +159,7 @@ spec_tensor = PatternMatcher([ (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), valid_gettuple), # SPECIAL is index before index lowering. custom_kernel currently has this - (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.weakint),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)), + (UPat(Ops.SPECIAL, src=(UPat(dtype=dtypes.weakint),), name="s"), lambda s: isinstance(s.arg, str)), # movement ops (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat())), lambda: True), @@ -173,20 +172,19 @@ spec_tensor = PatternMatcher([ and isinstance(x.arg[1], int) and all(y.dtype in (dtypes.weakint, dtypes.int) for y in x.src[1:])), # COPY - (UPat(Ops.COPY, name="copy", src=(UPat.var("x"),)), lambda copy,x: matches_dtype(x, copy.dtype) and is_device(copy.arg)), - (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"),)), lambda red,x: matches_dtype(x, red.dtype) and isinstance(red.arg, tuple) and - len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])), + (UPat(Ops.COPY, name="copy", src=(UPat(),)), lambda copy: is_device(copy.arg)), + (UPat(Ops.ALLREDUCE, name="red", src=(UPat(),)), + lambda red: isinstance(red.arg, tuple) and len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])), # UNSHARD/MSELECT/MSTACK # an UNSHARD carries the value and one sharding range per sharded axis (usually a DEVICE RANGE, but can be a derived expression) - (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 1+len(multi.arg) and matches_dtype(multi.src[0], multi.dtype) + (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 1+len(multi.arg) and all(isinstance(a, int) for a in multi.arg) and all(r.dtype in dtypes.weaks for r in multi.src[1:])), (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)), (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)), # CONTIGUOUS ensures the source UOp realizes - (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), name="root", src=(UPat.var("x"),), arg=None), - lambda root,x: matches_dtype(x, root.dtype)), + (UPat((Ops.DETACH, Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD), src=(UPat(),), arg=None), lambda: True), # TODO: this should not be here. STAGE is transformed to BUFFER later (UPat(Ops.STAGE, src=(UPat(),), allow_any_len=True), lambda: True), @@ -224,7 +222,7 @@ spec_program = PatternMatcher([ (UPat(Ops.ENDIF, dtype=dtypes.void, src=(UPat(Ops.IF),)), lambda: True), # SPECIAL is int32 after index lowering - (UPat(Ops.SPECIAL, src=(UPat.var("x", dtypes.int32),), name="s"), lambda s,x: matches_dtype(x, s.dtype) and isinstance(s.arg, str)), + (UPat(Ops.SPECIAL, src=(UPat(dtype=dtypes.int32),), name="s"), lambda s: isinstance(s.arg, str)), ])+spec_shared spec_hcq = PatternMatcher([ @@ -271,7 +269,7 @@ spec_kernel_graph = PatternMatcher([ (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True), # after on PARAM or AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),), - allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)), + allow_any_len=True), lambda: True), ]) # **** pyrender (move this) ****