diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index dcac5eabd4..7f866cd9c8 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -188,8 +188,8 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]: x, valid = idx.src[1].get_idx(), idx.src[1].get_valid() h, w = dt.shape[0], dt.shape[1] if IMAGE == 1 and valid is not None: - h, w = max(ImageDType.valid_dims(dt), - key=lambda hw: len(_drop_valid_stmts(valid, uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw))) + h, w = max(ImageDType.valid_dims(dt), key=lambda hw: + (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)), -len(idx.backward_slice))) buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize)) oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % w, (x // (4*w)))) return x, idx.replace(src=(buf, oidx.valid(valid))), w, h