From 984a0edbc9be619fd22b64268b07bdbfab243141 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 24 Jul 2025 16:13:43 -0700 Subject: [PATCH] identity store for DEFINE_REG --- tinygrad/codegen/devectorizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 8b7b8c5c1d..588310d2ea 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -284,7 +284,7 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp): assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}" # if we have a range if len(reduce_range) != 0: - input_ranges = tuple([x for x in inp.toposort(gate=lambda x: x.op is Ops.STORE) if x.op is Ops.RANGE and x not in reduce_range]) + input_ranges = tuple([x for x in inp.toposort(gate=lambda x: x.op is not Ops.STORE) if x.op is Ops.RANGE and x not in reduce_range]) identity = red.const_like(identity_element(red.arg, red.dtype.scalar())) acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), (identity,), (ctx.acc_num,)).index(UOp.const(dtypes.int, 0)) lst = [acc.store(identity, UOp(Ops.NOOP, src=input_ranges)).load(*reduce_range)] + lst # put acc as the first element