preserve metadata

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)