diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 975a5ef06b..0e756a31b0 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -74,6 +74,7 @@ class DType(metaclass=DTypeMetaClass): def vcount(self): return self.count @functools.cache # pylint: disable=method-cache-max-size-none def vec(self, sz:int) -> DType: + raise RuntimeError("vec is no longer supported") assert self.count == 1, f"can't vectorize {self} with size {sz}" if sz == 1 or self == dtypes.void: return self # void doesn't vectorize, and sz=1 is scalar return DType(self.priority, self.bitsize*sz, f"{INVERSE_DTYPES_DICT[self.name]}{sz}", None, sz, self) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 81902a4934..26495ad290 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -65,9 +65,9 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str: return ret def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp: - if len(arg) == 0: return UOp(Ops.STACK, dtypes.weakint.vec(0)) - elif all_int(arg): return UOp.const(dtypes.weakint.vec(len(arg)), arg) - else: return UOp(Ops.STACK, dtypes.weakint.vec(len(arg)), tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg)) + if len(arg) == 0: return UOp(Ops.STACK, dtypes.weakint) + elif all_int(arg): return UOp.const(dtypes.weakint, arg) + else: return UOp(Ops.STACK, dtypes.weakint, tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg)) def consumer_map_from_toposort(lst:Iterable[UOp]): ret: dict[UOp, dict[UOp, None]] = {} @@ -212,7 +212,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): match self.op: # late ops don't have shape case Ops.UNIQUE | Ops.LUNIQUE | Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \ - Ops.STACK | Ops.GEP | Ops.UNROLL | Ops.CONTRACT | Ops.SINK | \ + Ops.STACK | Ops.GEP | Ops.CONTRACT | Ops.SINK | Ops.END | \ Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.BINARY | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION: return None @@ -233,17 +233,22 @@ class UOp(OpMixin, metaclass=UOpMetaClass): return None case Ops.INDEX: + shp = [] + for s in self.src[1:]: shp.extend(list(s.shape)) + return tuple(shp) + """ # non pointer index doesn't have a shape if not isinstance(self.dtype, PtrDType): return None # fully indexed doesn't have a shape. TODO: remove this if self.src[0]._shape is None or len(self.src[1:]) == len(self.src[0].shape): return None # pointer index return self.src[0].shape[len(self.src[1:]):] + """ # some ops init the shape - case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL: return () - # TODO: VCONST should have the shape of the arg - case Ops.VCONST: return () + case Ops.CONST | Ops.DEFINE_VAR | Ops.BIND | Ops.RANGE | Ops.SPECIAL | Ops.UNROLL: return () + case Ops.VCONST: return (len(self.arg),) + case Ops.STACK: return (len(self.src),) case Ops.BUFFER: return (self.arg,) case Ops.BUFFER_VIEW: return (self.arg[0],) case Ops.CUSTOM_FUNCTION: return None @@ -259,7 +264,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): case Ops.SHAPED_WMMA: return self.src[2]._shape # passthrough ops - case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END | Ops.LOAD: + case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD: return self.src[0]._shape # TODO: disallow shape changing bitcast @@ -422,7 +427,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass): if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0] return UOp(Ops.GROUP, dtypes.void, tuple([x for x in srcs if x is not None])) def vectorize(self, *srcs, **kwargs): - return UOp(Ops.STACK, self.dtype.vec(len(srcs)+1), (self,)+srcs, **kwargs) + return UOp(Ops.STACK, self.dtype, (self,)+srcs, **kwargs) def index(self, *srcs:UOp|None, ptr=False, **kwargs): return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype if ptr else self.dtype.base), (self,)+tuple([x for x in srcs if x is not None]), **kwargs) def __getitem__(self, idx): @@ -486,13 +491,13 @@ class UOp(OpMixin, metaclass=UOpMetaClass): @staticmethod def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None): if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b - if isinstance(b, tuple) and all_same(b): - assert len(b) > 0, "can't create const from empty tuple" - b = b[0] # doesn't have to be a VCONST if they are all the same + #if isinstance(b, tuple) and all_same(b): + # assert len(b) > 0, "can't create const from empty tuple" + # b = b[0] # doesn't have to be a VCONST if they are all the same ret = UOp(Ops.VCONST if isinstance(b, tuple) else Ops.CONST, dtype, arg=dtype.const(b), src=(UOp(Ops.DEVICE, arg=device),) if device is not None else ()) - return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None else ret + return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and shape is not () else ret @staticmethod def unique_const(fill_value:ConstType, dtype:DTypeLike|None=None, device:str|tuple[str, ...]|None=None, # type: ignore[override] shape:tuple[sint, ...]|None=None, unique=True): diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 86440a15e5..eeb55bfa9a 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -294,8 +294,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # after with 1 src is just src[0] (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s), # VECTORIZE/CONST - (UPat(Ops.STACK, src=UPat(Ops.CONST), name="vec"), - lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src)) if len(vec.src) > 0 else None), + #(UPat(Ops.STACK, src=UPat(Ops.CONST), name="vec"), + # lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src)) if len(vec.src) > 0 else None), ])+div_and_mod_symbolic+gep_pushing # ******** we take a small aside to "simplify_valid" to rewrite valids ******** diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 7d70f7379f..964d7afd89 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -116,7 +116,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: if u.op is Ops.VCONST and u.dtype.scalar() == dtypes.weakint and u is not x: excluded.add(u) if u.op is Ops.STACK and len(u.src) == 0: excluded.add(u) # exclude RESHAPE/EXPAND that only serve to broadcast a CONST - if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u) + #if u.op in {Ops.RESHAPE, Ops.EXPAND} and len(u.src) >= 1 and u.src[0] in excluded and u is not x: excluded.add(u) for u in toposort: if u in excluded: continue argst = codecs.decode(str(u.arg), "unicode_escape")