From d51e55aa1783cab9b9b7ed3c9a094f4a59a07fc8 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 5 Aug 2026 21:28:38 -0400 Subject: [PATCH] remove some pm_fold_cast_const [pr] (#17426) --- tinygrad/codegen/__init__.py | 10 +++++----- tinygrad/codegen/simplify.py | 2 +- tinygrad/schedule/__init__.py | 4 ++-- tinygrad/schedule/indexing.py | 8 ++++---- tinygrad/uop/ops.py | 4 ++-- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 93ee01837e..a7ebfbe24b 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -310,7 +310,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = apply_opts(sink, ren, beam=ast.arg.beam) # ** expander (expand_rewrite) ** - sink = graph_rewrite(sink, sym+pm_fold_cast_const+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic") + sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic") # expand sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander") @@ -326,7 +326,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # **** optimizations are done, now we lower to actual code **** - sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads") + sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads") # devectorize sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+devectorizer2+indexing_simplify, ctx=ren, name="devectorize2") @@ -336,18 +336,18 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # do memory coalescing (late) sink = memory_coalescing(sink, ren) - sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+ew_devectorizer+pm_simplify_add_image, + sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True) # extra symbolic before decomp. crashes without this? - sink = graph_rewrite(sink, sym+pm_fold_cast_const, name="extra symbolic") + sink = graph_rewrite(sink, sym, name="extra symbolic") # lower index dtype # NOTE: we need indexing_simplify to remove the cast to long using the Invalid sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") # final symbolic before decomp - sink = graph_rewrite(sink, symbolic+pm_fold_cast_const, name="final symbolic") + sink = graph_rewrite(sink, symbolic, name="final symbolic") sink = graph_rewrite(sink, pm_cast_float_alu, name="cast float alu operands") diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index dbaa6c44f3..879dd0e6c9 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -119,7 +119,7 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([ lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x), # MUL casted bool ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)), -])+symbolic+pm_fold_cast_const +])+symbolic pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([ # lift x+y out of reduce on ne diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 3941c281fd..a98e97b7c4 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -151,8 +151,8 @@ def simplify_copy_kernel(call:UOp, ast:UOp, dst:UOp, src:UOp): if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None from tinygrad.codegen.simplify import pm_flatten_range, pm_simplify_ranges from tinygrad.schedule.rangeify import pm_mops - from tinygrad.uop.symbolic import sym, pm_fold_cast_const - sink = graph_rewrite(ast, sym+pm_fold_cast_const+pm_mops+pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges in copy") + from tinygrad.uop.symbolic import sym + sink = graph_rewrite(ast, sym+pm_mops+pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges in copy") return call.replace(src=(sink,) + call.src[1:]) pm_copy_from_store = PatternMatcher([ diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 10317a37b7..22f89362c6 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field, replace from tinygrad.dtype import dtypes, AddrSpace from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes from tinygrad.uop.ops import gate_kernel_sink -from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, pm_simplify_valid, pm_drop_and_clauses +from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE, @@ -168,7 +168,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U axes_out.append(combined_axes % s) combined_axes //= s # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code - return graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_fold_cast_const+pm_simplify_valid+pm_drop_and_clauses, name="reshape") + return graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic+pm_simplify_valid+pm_drop_and_clauses, name="reshape") # this is the definition of the movement ops @functools.cache @@ -182,7 +182,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local` # wraps the pad with only the newly added valid rngs = tuple(r if (sz == sh and off == 0) else (r-off).valid(graph_rewrite((r >= off) & (r < (sh+off)), - symbolic+pm_fold_cast_const+pm_simplify_valid, name="pad")) for r,sh,(off,sz) in zip(rngs, in_shape, arg)) + symbolic+pm_simplify_valid, name="pad")) for r,sh,(off,sz) in zip(rngs, in_shape, arg)) case Ops.RESHAPE: sink = UOp.sink(*rngs).simplify() # NOTE: this applies any commutative flips to the rngs early sub_array = {r:r.replace(src=r.src[:1], arg=(i, AxisType.PLACEHOLDER)) for i,r in enumerate(sink.ranges)} @@ -263,7 +263,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: if all_all_same or (PCONTIG and all_same(local_rngs)): # the new valid is the OR of all the children valids minimum_valid = UOp.const(False).usum(valids) - _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic+pm_fold_cast_const, name="minimum_valid")) + _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic, name="minimum_valid")) else: _out_rngs.append(rctx.new_range(x.shape[i])) _realize_axis.append(i) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index e68b879e7c..5f0706d964 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -895,7 +895,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def contiguous_view_offset(self) -> int|None: """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None.""" from tinygrad.schedule.rangeify import pm_mops - from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const + from tinygrad.uop.symbolic import symbolic # WEBGPU and CL do not support views. # WEBGPU requires that minUniformBufferOffsetAlignment be at least 32 bytes: https://gpuweb.github.io/gpuweb/#adapter-capability-guarantees @@ -905,7 +905,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if (dev:=self.device) is not None and any(d.startswith(("WEBGPU", "CL")) for d in ((dev,) if isinstance(dev, str) else dev)): return None idx = self.flatten().index(UOp.range(self.numel(), 0)) - out = graph_rewrite(idx, pm_mops+symbolic+pm_fold_cast_const+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") + out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None def has_buffer_identity(self, after_ok=False):