diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 54b09f1c6c..dd4c4556e8 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -358,10 +358,14 @@ pm_reduce = PatternMatcher([ # add loads +def add_load(idx:UOp): + if isinstance(idx.dtype, PtrDType): return None + assert isinstance(idx.src[0].dtype, PtrDType), f"param is not PtrDType {idx.src[0].dtype}" + return idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base) + pm_add_loads = PatternMatcher([ # add loads to non ptr index - (UPat(Ops.INDEX, name="idx"), lambda idx: None if isinstance(idx.dtype, PtrDType) else - idx.replace(dtype=idx.src[0].dtype).load(dtype=idx.dtype.base)), + (UPat(Ops.INDEX, name="idx"), add_load), # remove loads from stores (UPat(Ops.STORE, src=(UPat(Ops.LOAD), UPat(name="val")), name="s"), lambda s,val: s.replace(src=(s.src[0].src[0], val))), ]) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 509d58433d..c0603fa0ae 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -270,9 +270,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass): # REDUCE with empty axis is passthrough (lowered form) case Ops.REDUCE if len(self.arg[1]) == 0: # these can mismatch if there's a horizonal reduce - if self.dtype.count > 1: + if self.src[0].dtype.count > 1: assert len(self.src[0]._shape) == 1 - return (self.dtype.count,) + return () if self.dtype.count == 1 else (self.dtype.count,) return self.src[0]._shape # TODO: disallow shape changing bitcast