From 37a930591f374c17ab8723ab8ad9ef07dd916962 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 3 Dec 2025 14:03:39 -0800 Subject: [PATCH] preserve metadata --- tinygrad/engine/schedule.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index b9e26977cf..fe09a31fcf 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -3,7 +3,7 @@ from typing import cast from dataclasses import dataclass, field, replace from collections import deque from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites -from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite +from tinygrad.uop.ops import PatternMatcher, UPat, graph_rewrite, graph_rewrite_map from tinygrad.uop.spec import type_verify, tensor_spec from tinygrad.device import Buffer, MultiBuffer from tinygrad.helpers import Metadata, DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize @@ -150,25 +150,28 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li # big_sink srcs are all the Tensors st = time.perf_counter() - # verify Tensors match the spec - if SPEC: type_verify(big_sink, tensor_spec) - # replace all UNIQUE buffers with LUNIQUE input_buffers: dict[UOp, UOp] = {} - big_sink = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache") + big_sink_cache = graph_rewrite(big_sink, pm_pre_sched_cache, ctx=input_buffers, name="rewrite for sched cache") sched_cache_key = big_sink.key if (sc_ret:=schedule_cache.get(sched_cache_key, None)) is None: + # verify Tensors match the spec (on big_sink, we only need to do this if cache misses) + if SPEC: type_verify(big_sink, tensor_spec) + + # hack to preserve metadata + graph_rewrite_map(big_sink, pm_pre_sched_cache, ctx={}, name="preserve metadata") + # tensor map is what we return tensor_map: dict[UOp, UOp] = {} - if any(isinstance(x._device, tuple) for x in big_sink.toposort()): - tensor_map |= get_multi_map(big_sink) - big_sink = big_sink.substitute(tensor_map, name="Apply Multi Map") - big_sink = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink.src])) + if any(isinstance(x._device, tuple) for x in big_sink_cache.toposort()): + tensor_map |= get_multi_map(big_sink_cache) + big_sink_cache = big_sink_cache.substitute(tensor_map, name="Apply Multi Map") + big_sink_cache = UOp.sink(*flatten([x.src if x.op is Ops.MULTI else [x] for x in big_sink_cache.src])) - tensor_map |= get_rangeify_map(big_sink) - big_sink = big_sink.substitute(tensor_map, name="Apply Kernelize Map") + tensor_map |= get_rangeify_map(big_sink_cache) + big_sink = big_sink_cache.substitute(tensor_map, name="Apply Kernelize Map") # save in schedule cache schedule_cache[sched_cache_key] = (big_sink, tensor_map)