tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-30 11:06:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

viz: execv after all producers complete (#15696)

Browse Source
This commit is contained in:
qazal
2026-04-13 08:15:47 +09:00
committed by GitHub
parent f7ff480fa6
commit 2ada38f777
2 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/device.py
+3 -2
View File
@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
from tinygrad.helpers import select_first_inited, DEV, VIZ, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target
from tinygrad.helpers import select_first_inited, DEV, EMULATED_DTYPES, IMAGE, FLOAT16, TracingKey, size_to_str, Target
from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype
if TYPE_CHECKING: from tinygrad.renderer import Renderer
@@ -364,7 +364,8 @@ if PROFILE:
with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(cpu_events+Compiled.profile_events+Buffer.profile_events, f)
if VIZ:
PROFILE.value = 0
if getenv("VIZ") > 0:
from tinygrad.uop.ops import launch_viz
launch_viz("PROFILE", fn)
tinygrad/uop/ops.py
+7 -5
View File
@@ -1300,7 +1300,9 @@ if TRACK_MATCH_STATS or PROFILE:
with open(fn:=temp("rewrites.pkl", append_user=True), "wb") as f:
print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}")
pickle.dump(RewriteTrace(tracked_keys, tracked_ctxs, uop_fields), f)
if VIZ > 0: return launch_viz("VIZ", temp("rewrites.pkl", append_user=True))
if getenv("VIZ") > 0:
VIZ.value = 0
return launch_viz("VIZ", temp("rewrites.pkl", append_user=True))
if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value and VIZ.value>=0):
ret = [0,0,0.0,0.0]
for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]+x[1][3]):
@@ -1311,11 +1313,11 @@ if TRACK_MATCH_STATS or PROFILE:
print(f"{len(match_stats)} rules, {sum(v[0] > 0 for v in match_stats.values())} matched once")
def launch_viz(env_str:str, data:str):
os.environ[env_str] = "0"
os.environ[f"{env_str}_DATA"] = data
if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")):
args = ['--rewrites-path', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else []
args += ['--profile-path', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else []
if not VIZ and not PROFILE:
os.environ["VIZ"], os.environ["PROFILE"] = "0", "0"
args = ['--rewrites-path', os.getenv("VIZ_DATA", "")] if os.getenv("VIZ_DATA", "") else []
args += ['--profile-path', os.getenv("PROFILE_DATA", "")] if os.getenv("PROFILE_DATA", "") else []
viz_path = pathlib.Path(__file__).resolve().parent.parent / "viz" / "serve.py"
os.execv(sys.executable, [sys.executable, viz_path.as_posix()] + args)