From 6b50df940a7e938c3bbf74769d8267c6269a4ce1 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 11 Mar 2026 03:54:42 -0400 Subject: [PATCH] remove mop_cleanup [pr] (#15217) no kernel diff, i think this was needed due to force_reshape? test/external/external_benchmark_schedule.py is about the same speed --- tinygrad/engine/jit.py | 3 +-- tinygrad/schedule/rangeify.py | 7 +------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index c7d307290e..e1e6356bc2 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -8,7 +8,6 @@ from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, EncDec, CompiledRunner, Runner, Estimates from tinygrad.engine.memory import _internal_memory_planner from tinygrad.nn.state import get_parameters -from tinygrad.schedule.rangeify import mop_cleanup from dataclasses import dataclass, replace from weakref import WeakKeyDictionary @@ -265,7 +264,7 @@ def _prepare_jit_inputs(args, kwargs): raise JitError("JIT inputs cannot be const, create a buffer with .contiguous()") input_buffers: list[Buffer] = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for u in input_uops if (b:=u.base.realized) is not None]) if len(set(input_buffers)) != len(input_buffers): raise JitError("duplicate inputs to JIT") - inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops] + inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}).unbind_all()), u.dtype, u.device) for u in input_uops] _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))]) var_vals = {k.expr:v for k,v in _var_vals.items()} expected_input_info = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in inputs] diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 9be14b4573..fb482d8df4 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -87,11 +87,6 @@ def split_reduceop(reduce:UOp, x:UOp): # reduce original axes, then split return splitted.r(*reduce.arg).contiguous().r(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape) -mop_cleanup = PatternMatcher([ - # merge adjacent RESHAPES - (UPat(Ops.RESHAPE, src=(UPat(Ops.RESHAPE, name="x2"), UPat()), name="x"), lambda x,x2: x.replace(src=(x2.src[0], x.src[1]))), -]) - pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p)), ]) def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None: if not should_resolve_call(c): return None @@ -112,7 +107,7 @@ def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None: if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}") return c.src[0].substitute(dict_map, walk=True) -earliest_rewrites = mop_cleanup+PatternMatcher([ +earliest_rewrites = PatternMatcher([ # early fixup const copy (UPat(Ops.COPY, src=(UPat.var("s"), UPat.var("d"))), lambda s,d: s.substitute({UOp(Ops.DEVICE, arg=s.device):d}) if s.base.op is Ops.CONST else None),