sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

update rules for INDEX mops (#17105)

Browse Source
This commit is contained in:
George Hotz
2026-07-20 18:43:00 -07:00
committed by GitHub
parent 13ca9bd8a6
commit 2a81616492
2 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/__init__.py
-3
View File
@@ -175,9 +175,6 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
# EXPAND on scalar -> STACK
(UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
# INDEX on INDEX is INDEX
(UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
])
def fix_group_for_reduce(x:UOp):
tinygrad/uop/movement.py
+7 -1
View File
@@ -16,5 +16,11 @@ mop_cleanup = PatternMatcher([
lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].arg for x in stk.src] else None),
# const INDEX into STACK is src
(UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i")), name="idx", allow_any_len=True),
lambda a,i,idx: a.src[i.arg] if len(idx.src) <= 2 else a.src[i.arg].index(*idx.src[2:])),
lambda a,i,idx: a.src[i.arg] if len(idx.src) <= 2 else a.src[i.arg].index(*idx.src[2:])),
# INDEX on INDEX is INDEX
(UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
lambda idx1,idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:]) if all(x.shape == () for x in idx1.src[1:]+idx2.src[1:]) else None),
# INDEX on shaped INDEX (TODO: this can be more generic)
(UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx1_arg"))),), allow_any_len=True, name="idx2"),
lambda buf,idx1_arg,idx2: buf.index(idx1_arg.index(*idx2.src[1:])) if len(idx1_arg.shape) == len(idx2.src[1:]) else None),
])