no global kernel stuff [run_process_replay] (#6808)

* use traceback instead of global metadata crap [run_process_replay]

* save the kernel

* correct, imports clean, no device

* UNPARENTED

* speed

* proudly unparented

* Update ops.py

* update tests for unparented

---------

Co-authored-by: qazal <[email protected]>

test/external/external_benchmark_schedule.py

@@ -1,7 +1,7 @@
 from typing import List
 from extra.models.resnet import ResNet50
 from tinygrad import Tensor, Device
-from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen, _CURRENT_KERNEL
+from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
 from tinygrad.ops import UOps
 from tinygrad.codegen.kernel import Kernel
 from tinygrad.codegen.lowerer import ast_to_uop
@@ -43,7 +43,6 @@ if __name__ == "__main__":
             rewritten_uops = []
             for i,(k,u) in enumerate(zip(kernels, uops)):
               with Timing(f"rewrite {i:2d} {k.name}{' '*(50-ansilen(k.name))}", enabled=getenv("VERBOSE", 0)):
-                if getenv("VIZ"): _CURRENT_KERNEL.set(k.name)
                 rewritten_uops.append(full_graph_rewrite(u, k.opts))
             uops = rewritten_uops
         if getenv("LINEARIZE", 1):

tinygrad/codegen/kernel.py

@@ -5,13 +5,12 @@ from collections import defaultdict
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict
 from enum import Enum, auto
 
-from tinygrad.ops import TRACK_MATCH_STATS, BinaryOps, UNSAFE_PAD_OPS, KernelInfo, BUFFER_UOPS, UOp, UOps, print_uops, type_verify, \
-  graph_rewrite, PatternMatcher
+from tinygrad.ops import BinaryOps, UNSAFE_PAD_OPS, KernelInfo, BUFFER_UOPS, UOp, UOps, print_uops, type_verify, graph_rewrite, PatternMatcher
 from tinygrad.device import Device
 from tinygrad.renderer import Renderer, TensorCore, Program
 from tinygrad.dtype import ImageDType, PtrDType
 from tinygrad.helpers import all_same, colored, ansilen, dedup, getenv, prod, round_up, all_int, get_contraction, to_function_name, diskcache_put
-from tinygrad.helpers import _CURRENT_KERNEL, DEBUG, TC_OPT, USE_TC, AMX
+from tinygrad.helpers import DEBUG, TC_OPT, USE_TC, AMX
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.symbolic import Variable, sint
 from tinygrad.shape.view import strides_for_shape
@@ -704,7 +703,6 @@ class Kernel:
   # **** this is the lowerer ****
 
   def linearize(self) -> Kernel:
-    if TRACK_MATCH_STATS >= 2: _CURRENT_KERNEL.set(self.name)
     modified_ast = self.get_optimized_ast()
 
     if DEBUG >= 3:
@@ -715,7 +713,6 @@ class Kernel:
     verify_ast(modified_ast)
 
     self.uops:List[UOp] = linearize_uop(full_graph_rewrite(ast_to_uop(modified_ast, self.opts), self.opts))
-    if TRACK_MATCH_STATS >= 2: _CURRENT_KERNEL.set(None)
     if DEBUG >= 5: print_uops(self.uops)
     if getenv("GRAPHUOPS"):
       from tinygrad.engine.graph import graph_uops

tinygrad/helpers.py

@@ -131,8 +131,6 @@ class Metadata:
   def __str__(self): return self.name + (" bw" if self.backward else "")
 _METADATA: contextvars.ContextVar[Optional[Metadata]] = contextvars.ContextVar("_METADATA", default=None)
 
-_CURRENT_KERNEL: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar("_CURRENT_KERNEL", default=None)
-
 # **************** global state Counters ****************
 
 class GlobalCounters:

tinygrad/ops.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 from typing import Any, List, Optional, Set, Union, Tuple, Dict, Callable, cast, TYPE_CHECKING, TypeVar
+from types import FrameType
 import sys, time, functools, itertools, math, operator, hashlib, os
 from enum import auto, IntEnum, Enum
 from dataclasses import dataclass, field
 from tinygrad.dtype import ConstType, ImageDType, PtrDType, dtypes, DType, truncate
-from tinygrad.helpers import _CURRENT_KERNEL, ContextVar, pretty_print, prod, getenv, all_same
+from tinygrad.helpers import ContextVar, pretty_print, prod, getenv, all_same
 from tinygrad.shape.symbolic import Variable, sint
 if TYPE_CHECKING:
   from tinygrad.shape.shapetracker import ShapeTracker
+  from tinygrad.codegen.kernel import Kernel
 
 # wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses
 class FastEnum(IntEnum):
@@ -499,7 +501,7 @@ match_stats:Dict[UPat, List[Union[int, float]]] = dict()
 class TrackedRewriteContext:
   loc: Tuple[str, int]                                                    # location that called graph_rewrite
   sink: UOp                                                               # the sink passed into the rewrite
-  kernel_name: Optional[str] = None                                       # the name of the kernel being rewritten
+  kernel: Optional[Kernel] = None                                         # the kernel being rewritten
   rewrites: List[Tuple[UOp, UOp, UPat]] = field(default_factory=list)     # all rewrites of sparents. (before, after, UPat)
 contexts: List[TrackedRewriteContext] = []
 class TrackedPatternMatcher(PatternMatcher):
@@ -565,7 +567,12 @@ class RewriteContext:
     return found
 def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None) -> UOp:
   if TRACK_MATCH_STATS >= 2:
-    contexts.append(TrackedRewriteContext(((f:=sys._getframe(1)).f_code.co_filename, f.f_lineno), sink, _CURRENT_KERNEL.get()))
+    from tinygrad.codegen.kernel import Kernel
+    frm = sys._getframe(1)
+    # get Kernel we are rewriting in the context of
+    frm_walk: Optional[FrameType] = frm
+    while frm_walk is not None and not isinstance(kernel:=frm_walk.f_locals.get("self", None), Kernel): kernel, frm_walk = None, frm_walk.f_back
+    contexts.append(TrackedRewriteContext((frm.f_code.co_filename, frm.f_lineno), sink, kernel))
   return RewriteContext(pm, ctx).rewrite(sink)
 
 # ***** uop type spec *****

viz/serve.py

@@ -1,16 +1,13 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 from typing import Dict, List, Optional, Tuple
-import pickle, os, sys, time, threading, webbrowser, json, difflib, contextlib, re, traceback
+import pickle, os, sys, time, threading, webbrowser, json, difflib, contextlib, re, multiprocessing
 from dataclasses import dataclass, asdict
 from urllib.parse import parse_qs, urlparse
 from http.server import HTTPServer, BaseHTTPRequestHandler
-from tinygrad import Device
-from tinygrad.helpers import Context, getenv, to_function_name
+from tinygrad.helpers import getenv, to_function_name
 from tinygrad.ops import TrackedRewriteContext, UOp, UOps, lines
 from tinygrad.engine.graph import uops_colors, word_wrap
-from tinygrad.engine.realize import get_runner
-from tinygrad.engine.schedule import ScheduleItemContext, full_ast_rewrite
 
 # **** /graph - detailed UOp + rewrites
 
@@ -93,16 +90,10 @@ class KernelRet:
 
 def load_kernels(contexts:List[TrackedRewriteContext]) -> List[KernelRet]:
   ret: Dict[str, KernelRet] = {}
-  kernel_name = ""
-  code = ""
   for ctx in contexts:
-    try:
-      if ctx.loc[0].split("/")[-1] == "schedule.py":
-        si_ctx = ScheduleItemContext(bufs=tuple(x.arg for x in ctx.sink.sparents if x.op is UOps.BUFFER))
-        with Context(TRACK_MATCH_STATS=0): kernel_name, code = (prg:=get_runner(Device.DEFAULT, full_ast_rewrite(ctx.sink, si_ctx)).p).name, prg.src
-      elif ctx.kernel_name is not None: kernel_name, code = ctx.kernel_name, ""
-    except Exception: kernel_name, code = "RENDERING_ERROR", traceback.format_exc()
-    if ret.get(k:=to_function_name(kernel_name)) is None: ret[k] = KernelRet(k, code, [])
+    name = ctx.kernel.name if ctx.kernel is not None else "UNPARENTED"
+    if ret.get(k:=to_function_name(name)) is None:
+      ret[k] = KernelRet(k, ctx.kernel.to_program().src if ctx.kernel is not None else "", [])
     ret[k].ctxs.append(ctx)
   return list(ret.values())
 
@@ -149,6 +140,7 @@ def reloader():
     time.sleep(0.1)
 
 if __name__ == "__main__":
+  multiprocessing.current_process().name = "VizProcess"    # disallow opening of devices
   print("*** viz is starting")
   with open("/tmp/rewrites.pkl", "rb") as f: contexts: List[TrackedRewriteContext] = pickle.load(f)
   print("*** unpickled saved rewrites")

viz/test_viz.py

@@ -48,9 +48,9 @@ class TestViz(unittest.TestCase):
     list(lower_schedule(schedule1))
     list(lower_schedule(schedule2))
     ret = load_kernels(contexts)
-    assert len(ret) == 2
-    assert all(len([x for x in y.ctxs if "schedule" in x.loc[0]]) != 0 for y in ret)
-    assert all(len([x for x in y.ctxs if "uopgraph" in x.loc[0]]) != 0 for y in ret)
+    assert len(ret) == 3
+    assert all(len([x for x in y.ctxs if "schedule" in x.loc[0]]) == 0 for y in ret[1:])
+    assert all(len([x for x in y.ctxs if "uopgraph" in x.loc[0]]) != 0 for y in ret[1:])
 
   def test_gemm_diff(self):
     x = Tensor.empty(64, 64).realize()
@@ -132,10 +132,9 @@ class TestViz(unittest.TestCase):
     s = a.schedule()
     with Context(NOOPT=1): list(lower_schedule(s.copy()))
     with Context(NOOPT=0): list(lower_schedule(s.copy()))
-    kernels = load_kernels(contexts)
+    kernels = load_kernels(contexts)[1:]
     self.assertEqual(len(kernels), 2)
     assert all(len(v) == 1 for _,v in group_rewrites(kernels[0]).items())
-    assert all(len(v) == 0 for k,v in group_rewrites(kernels[1]).items() if "schedule.py" in k)
 
   def test_fold_const_nodes(self):
     a = Tensor.empty(4, 4)+2