github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

simplify idx before counting backward_slice (#15117)

Browse Source
This commit is contained in:
Christopher Milan
2026-03-03 20:53:50 -08:00
committed by GitHub
parent 2d72a4a90c
commit e7e70a3c95
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tinygrad/codegen/late/devectorizer.py
View File

@@ -189,7 +189,10 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
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, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)), -len(idx.backward_slice)))
# maximize number of valids removed
(len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)),
# and minimize idx complexity (number of nodes)
-len(idx.simplify().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