From ab71e51ae131940d52c388904caf2fcf36f533fd Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 26 May 2026 13:36:31 -0700 Subject: [PATCH] cast also --- tinygrad/codegen/late/devectorizer.py | 5 +++-- tinygrad/renderer/cstyle.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 21abad2d32..38b2021b23 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -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 *** diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index bcb13568e1..2637b6f80b 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -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]})"),