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

cast also

Browse Source
This commit is contained in:
geohot
2026-05-27 01:25:14 +00:00
parent 0cb3ce4811
commit ab71e51ae1
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/late/devectorizer.py
+3 -2
View File
@@ -284,8 +284,9 @@ pm_render = PatternMatcher([
(UPat(Ops.GEP, name='gep'), lambda gep: gep.src[0] if gep.src[0].dtype.vcount == 1 and gep.arg == (0,) else None),
(UPat(Ops.STACK, src=(UPat(name='x'),)), lambda x: x),
# rewrite INDEX to SLICE
(UPat(Ops.INDEX, name="x"), lambda x:
UOp(Ops.SLICE, dtype=x.dtype, src=(x.src[0], x.src[1]), arg=0 if x.dtype.vcount == 1 else x.dtype.vcount)),
(UPat(Ops.INDEX, name="x"), lambda x: UOp(Ops.SLICE, dtype=x.dtype, src=x.src, arg=0 if x.dtype.count == 1 else x.dtype.count)),
# rewrite CAST on SLICE to SLICE
(UPat(Ops.SLICE, name="bv").cast(name="x"), lambda bv,x: bv.replace(dtype=x.dtype, arg=x.dtype.count))
])
# *** Ops.REDUCE -> Ops.DEFINE_ACC ***
tinygrad/renderer/cstyle.py
+2 -2
View File
@@ -44,8 +44,8 @@ base_rewrite = PatternMatcher([
# default const render
(UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
# slice is the new index
(UPat(Ops.SLICE, src=(UPat.var("buf"), UPat.var('idx'))),
lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
(UPat(Ops.SLICE, src=(UPat.var("buf"), UPat.var('idx')), name="x"),
lambda ctx,buf,idx,x: ctx.render_cast(x.dtype, f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})")),
# new load/store
(UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"(*{ctx[bidx]})"),
(UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),