github-actions[bot]
de77c89cf6
version: sunnypilot v2026.002.000 (staging) date: 2026-07-21T07:44:02 master commit: 2c334ede443d7391d27575af8c854a095ba702a8
22 lines
1.6 KiB
Python
22 lines
1.6 KiB
Python
# this is a temporary intermediate step while we remove this index style
|
|
from tinygrad.uop.ops import PatternMatcher, UPat, Ops
|
|
from tinygrad.dtype import Invalid, dtypes
|
|
|
|
pm_move_gates_from_index = PatternMatcher([
|
|
# here we create the alt value for load to be 0s and remove the where Invalid
|
|
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").load(name="l"),
|
|
lambda buf,gate,idx,cast,l: buf.index(idx, ptr=True).cast(cast.dtype).load(l.const_like(0), gate, dtype=l.dtype)),
|
|
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid))).or_casted(name="cast").store(UPat.var("data")),
|
|
lambda buf,gate,idx,cast,data: buf.index(idx, ptr=True).cast(cast.dtype).store(data, gate)),
|
|
|
|
# Where after gated load becomes alt value
|
|
(UPat.var("gate").where(UPat().load(UPat(), UPat.var("gate", dtype=dtypes.bool), name="l").or_casted(), UPat.var("a")), lambda gate,l,a:
|
|
l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
|
|
(UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()), lambda gate,l,a:
|
|
l.replace(src=(l.src[0], a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(a.dtype)),
|
|
|
|
# images use 2D INDEX now (y,x)
|
|
(UPat(Ops.INDEX, src=(UPat(), UPat((Ops.CONST, Ops.VCONST, Ops.STACK), name="vec")), name="idx"),
|
|
lambda idx,vec: idx.replace(src=(idx.src[0], vec.gep(1).cast(dtypes.int), vec.gep(0).cast(dtypes.int))) if vec.dtype.count == 2 else None),
|
|
])
|