forked from tinygrad/tinygrad
remove dtype vec, try 2
This commit is contained in:
@@ -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)
|
||||
|
||||
+18
-13
@@ -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):
|
||||
|
||||
@@ -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 ********
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user