IQ.Pilot Prebuilt Release @ 27f668a
This commit is contained in:
29
artifacts/package_runtime/tinygrad/codegen/late/gater.py
Normal file
29
artifacts/package_runtime/tinygrad/codegen/late/gater.py
Normal file
@@ -0,0 +1,29 @@
|
||||
# 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
|
||||
|
||||
def move_where_load(gate, l, a, w):
|
||||
return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
|
||||
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(w.dtype)
|
||||
|
||||
pm_move_gates_from_index = PatternMatcher([
|
||||
# for image idx (must be first)
|
||||
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
|
||||
UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).load(name="l"),
|
||||
lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x, dtype=dtypes.float).load(l.vconst_like(0), gate)),
|
||||
(UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
|
||||
UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).store(UPat.var("data")),
|
||||
lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x, dtype=dtypes.float).store(data, gate)),
|
||||
|
||||
# here we create the alt value for load to be 0s and remove the where Invalid
|
||||
(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \
|
||||
.load(name="l"), lambda mop,gate,idx,l: mop.replace(src=(mop.src[0],idx)+mop.src[2:]).load(l.vconst_like(0), gate)),
|
||||
(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \
|
||||
.store(UPat.var("data")), lambda mop,gate,idx,data: mop.replace(src=(mop.src[0],idx)+mop.src[2:]).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")).named("w"),
|
||||
move_where_load),
|
||||
(UPat.var("gate").where(UPat.var("a"), UPat().load(UPat(), ~UPat.var("gate", dtype=dtypes.bool), name="l").or_casted()).named("w"),
|
||||
move_where_load),
|
||||
])
|
||||
Reference in New Issue
Block a user