From 7969b205dd89f91109ed22e09fb5fbfddcd3c0e4 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 30 Apr 2026 07:22:06 -0700 Subject: [PATCH] fix --- tinygrad/codegen/late/devectorizer.py | 8 ++++++-- tinygrad/uop/ops.py | 4 ++-- 2 files changed, 8 insertions(+), 4 deletions(-) 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