From 433ebecee7fd8bf4479b291b16376ed7dd2f4060 Mon Sep 17 00:00:00 2001 From: Roelof van Dijk <3604013+roelofvandijk@users.noreply.github.com> Date: Wed, 13 Nov 2024 17:06:59 +0100 Subject: [PATCH] refactor: double if statement [pr] (#7674) --- tinygrad/codegen/uopgraph.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py index 1ee8585320..98e29348dd 100644 --- a/tinygrad/codegen/uopgraph.py +++ b/tinygrad/codegen/uopgraph.py @@ -50,14 +50,12 @@ def fold_expanded(ex, buf): else: # for non image, we upcast the index pointer new_src[0] = new_src[0].cast(new_src[0].dtype.base.vec(fold_length).ptr(new_src[0].dtype.local)) - # vectorize the store - if not is_load: - new_src[1] = UOp(Ops.VECTORIZE, new_src[1].dtype.vec(fold_length), tuple(new_srcs[offsets[o+i]].src[1] for i in range(fold_length))) # generate the folded new_srcs if is_load: new_load = UOp(Ops.LOAD, load_1.dtype.vec(fold_length), tuple(new_src)) for i in range(fold_length): new_srcs[offsets[o+i]] = new_load.gep(i) - else: + else: # vectorize the store + new_src[1] = UOp(Ops.VECTORIZE, new_src[1].dtype.vec(fold_length), tuple(new_srcs[offsets[o+i]].src[1] for i in range(fold_length))) for i in range(fold_length): new_srcs[offsets[o+i]] = UOp(Ops.STORE, dtypes.void, tuple(new_src)) if i == 0 else None for i in range(fold_length): used.add((rootsrc,o+i))