From 6ea665ed66a8fc02543f6893e73b76fe68619e0f Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 14 Aug 2026 16:24:51 -0400 Subject: [PATCH] remove pm_fold_cast_const from dtype decomp [pr] (#17536) * remove pm_fold_cast_const from dtype decomp [pr] * fix --- tinygrad/codegen/__init__.py | 2 +- tinygrad/codegen/late/gater.py | 2 +- tinygrad/renderer/llvmir.py | 4 ++-- tinygrad/renderer/nir.py | 3 ++- tinygrad/uop/weak.py | 1 + 5 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 0f0877d9f2..02150e4320 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -357,7 +357,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # floordiv+mod / dtype decomp (early) supported_ops = tuple(ren.code_for_op.keys()) - pm_decomp = symbolic_simple+pm_fold_cast_const+get_simplifying_rewrite_patterns(supported_ops) + pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) sink = graph_rewrite(sink, pm_decomp, name="early decompositions") # late decomps + move gates from unrenderable INVALID where diff --git a/tinygrad/codegen/late/gater.py b/tinygrad/codegen/late/gater.py index b9b3a4efb7..e9b46ecfd5 100644 --- a/tinygrad/codegen/late/gater.py +++ b/tinygrad/codegen/late/gater.py @@ -3,7 +3,7 @@ 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 + return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else l.const_like(a.val) if a.op is Ops.CONST 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([ diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 838e1c95f0..3ce6fc8919 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -238,8 +238,8 @@ class AMDLLVMRenderer(LLVMRenderer): (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x: f" {ctx[x]} = call i8 @f32_to_fp8({ldt(x.src[0].dtype)} {ctx[x.src[0]]}, i1 {'1' if x.dtype == dtypes.fp8e5m2 else '0'})"), (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), lambda ctx,x,y: - f" {ctx[x.src[0]]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n" - f" {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x.src[0]]}_i32, i32 0)"), + f" {ctx[x]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n" + f" {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x]}_i32, i32 0)"), ]) + base_rewrite extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([ # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index b0d3a38f2c..a2153394f6 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -137,7 +137,8 @@ class NIRRenderer(Renderer): (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)), # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace( - src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None), + src=(buf,UOp.const(off.val, dtypes.long) if off.op is Ops.CONST else off.cast(dtypes.long))+x.src[2:]) + if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None), # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype) (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"), lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))), diff --git a/tinygrad/uop/weak.py b/tinygrad/uop/weak.py index 0460845664..398c68fb3c 100644 --- a/tinygrad/uop/weak.py +++ b/tinygrad/uop/weak.py @@ -65,6 +65,7 @@ def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None: pm_cast_weak = PatternMatcher([ (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs), + (UPat(Ops.CAST, name="c", src=(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"),)), lambda c,u: commit_weak(u, c.dtype)), ]) pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([