tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 11:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

test plus passes

Browse Source
This commit is contained in:
geohot
2025-07-31 22:39:42 -07:00
parent 5706e2d845
commit 149c3f8fe9
2 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/rangeify.py
+4 -1
View File
@@ -1,4 +1,4 @@
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, KernelInfo
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, KernelInfo, GroupOp
from tinygrad.dtype import dtypes
from tinygrad.helpers import argsort
@@ -61,4 +61,7 @@ pm_rangeify = PatternMatcher([
(UPat(Ops.INDEX, src=(UPat(Ops.EXPAND),), allow_any_len=True, name="x"), map_expand),
# TODO: CONST shouldn't have src
(UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),)), lambda c: c.replace(src=())),
# move MAP through elementwise ALU
(UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.LOAD})),), allow_any_len=True, name="x"),
lambda x: x.src[0].replace(src=tuple([UOp(Ops.INDEX, dtype=s.dtype, src=(s,)+x.src[1:]) for s in x.src[0].src]))),
])
tinygrad/uop/__init__.py
+1
View File
@@ -82,6 +82,7 @@ class GroupOp:
Ops.XOR, Ops.SHL, Ops.SHR, Ops.OR, Ops.AND, Ops.THREEFRY, Ops.SUB, Ops.FDIV, Ops.POW}
Ternary = {Ops.WHERE, Ops.MULACC}
ALU = set.union(Unary, Binary, Ternary)
Elementwise = set.union(ALU, {Ops.CAST, Ops.BITCAST})
Defines = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}