tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-09-02 06:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove mop_cleanup [pr] (#15217)

Browse Source 
no kernel diff, i think this was needed due to force_reshape?
test/external/external_benchmark_schedule.py is about the same speed
This commit is contained in:
chenyu
2026-03-11 03:54:42 -04:00
committed by GitHub
parent 2fb8a7f60f
commit 6b50df940a
2 changed files with 2 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/engine/jit.py
+1 -2
View File
@@ -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]
tinygrad/schedule/rangeify.py
+1 -6
View File
@@ -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),