diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 2cdc894d34..cbac28946b 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -285,7 +285,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp): topo = inp.toposort() stored_ranges = flatten([x.src[2:] for x in topo if x.op is Ops.STORE]) input_ranges = tuple([x for x in topo if x.op is Ops.RANGE and x not in reduce_range and x not in stored_ranges]) - identity = red.const_like(identity_element(red.arg, red.dtype.scalar())) + identity = red.const(red.dtype, identity_element(red.arg, red.dtype.scalar())) acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,)).index(UOp.const(dtypes.int, 0)) do_store = acc.store(identity, UOp(Ops.NOOP, src=input_ranges)) if len(input_ranges) else acc.store(identity) lst = [acc.load(do_store, *reduce_range)] + lst # put acc as the first element diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 2f925d1d3c..88c68b0123 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -146,7 +146,7 @@ class CStyleLanguage(Renderer): if u.arg is not None: name = u.arg.function_name continue if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR): - r[u] = f"data{u.arg}" if u.op is Ops.DEFINE_GLOBAL else u.arg[0] + r[u] = (f"data{u.arg}_{sz}" if (sz:=cast(PtrDType, u.dtype).size) > 0 else f"data{u.arg}") if u.op is Ops.DEFINE_GLOBAL else u.arg[0] bufs[u] = (r[u], (u.dtype, False)) continue diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index d99a0c4067..d5de866118 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -2933,11 +2933,11 @@ class Tensor(MathTrait): """ return self*-1 if self.dtype != dtypes.bool else self.logical_not() - def contiguous(self) -> Tensor: + def contiguous(self, **kwargs) -> Tensor: """ Returns a contiguous tensor. """ - return self._apply_uop(UOp.contiguous) + return self._apply_uop(UOp.contiguous, **kwargs) def fuse(self) -> Tensor: """ diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 3ae9a6e398..dfea89bfe0 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -86,6 +86,9 @@ class GroupOp: Ternary = {Ops.WHERE, Ops.MULACC} ALU = set.union(Unary, Binary, Ternary) + # TODO: is BITCAST always Elementwise if it's shape changing? + Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST}) + Defines = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG} Irreducible = {Ops.CONST, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.RANGE} diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8114b1a43e..f03c1c0a8f 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -182,6 +182,19 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @property def size(self) -> int: return self.arg[0] if self.op is Ops.BUFFER_VIEW else self.arg if self.op is Ops.BUFFER else unwrap(self.st).size + # determine what ranges this is in + @functools.cached_property + def ranges(self) -> dict[UOp, None]: + if self.op is Ops.RANGE: return {self:None} + if self.op in {Ops.CONTIGUOUS, Ops.REDUCE, Ops.STORE}: + ret = self.src[0].ranges.copy() + for s in self.src[1:]: + if s in ret: del ret[s] + else: + ret = {} + for s in self.src: ret.update(s.ranges) + return ret + # *** uop evaluation *** def simplify(self): @@ -219,7 +232,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return ret def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs) def detach(self): return UOp(Ops.DETACH, self.dtype, (self,)) - def index(self, idx:UOp, valid:UOp|None=None): return UOp(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx)) + def index(self, *srcs:UOp|None): return UOp(Ops.INDEX, self.dtype, (self,)+tuple([x for x in srcs if x is not None])) def __getitem__(self, idx): return self.index(idx) def const_like(self, b:ConstLike): # constants can optionally have a DEVICE source @@ -275,7 +288,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, new_axis)) return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)])) def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs) - def contiguous(self): return self.alu(Ops.CONTIGUOUS) + def contiguous(self, *args, **kwargs): return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs) def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD) def fuse(self): return self.alu(Ops.FUSE) def allreduce(self, op, device:str|tuple[str, ...]|UOp): diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 7126371d47..0f272a904a 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -75,11 +75,13 @@ def uop_to_json(x:UOp) -> dict[int, dict]: if x in excluded: if x.op is Ops.CONST and dtypes.is_float(u.dtype): label += f"\nCONST{idx} {x.arg:g}" else: label += f"\n{x.op.name}{idx} {x.arg}" - try: - if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: + if u.op not in {Ops.VIEW, Ops.BUFFER, Ops.KERNEL, Ops.ASSIGN, Ops.COPY, Ops.SINK, *GroupOp.Buffer} and u.st is not None: + try: label += f"\n{shape_to_str(u.shape)}" - except Exception: - label += "\n" + except Exception: + label += "\n" + elif len(rngs:=u.ranges): + label += f"\n{str(sorted([x.arg for x in rngs]))}" if (ref:=ref_map.get(u.arg.ast) if u.op is Ops.KERNEL else None) is not None: label += f"\ncodegen@{ctxs[ref]['name']}" # NOTE: kernel already has metadata in arg if TRACEMETA >= 2 and u.metadata is not None and u.op is not Ops.KERNEL: label += "\n"+repr(u.metadata)