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

consider indexing math cost for IMAGE=1 (#14973)

Browse Source
This commit is contained in:
Christopher Milan
2026-02-23 15:57:45 -08:00
committed by GitHub
parent 0bda5585c7
commit 57ade7608a
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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