diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index 8d96794496..2a040998e1 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -1,4 +1,4 @@ -import sys, functools +import sys, atexit, functools, pickle from collections import defaultdict, deque from dataclasses import dataclass, field from typing import FrozenSet, Set, Tuple, List, Dict, Optional, DefaultDict @@ -211,9 +211,15 @@ def full_ast_rewrite(pre:UOp, ctx:ScheduleContext) -> Tuple[UOp, ScheduleItemCon and ShapeTracker.from_shape(s.shape).shrink(m) == s.shrink(m)) for x in ops): raise RuntimeError("self operand of augmented assign must be contiguous.\nhelp: consider using .contiguous():\n" +colored(" - a += a.T\n", "red")+colored(" + a += a.T.contiguous()", "green")) - if getenv("RUN_PROCESS_REPLAY"): diskcache_put("schedule_process_replay", str(pre.key), (pre, si_ctx.assigns, {}, sink)) + if getenv("RUN_PROCESS_REPLAY"): PROCESS_REPLAY_CAPTURE[str(pre.key)] = pickle.dumps((pre, si_ctx.assigns, {}, sink)) return sink, si_ctx +PROCESS_REPLAY_CAPTURE: Dict[str, bytes] = {} +if getenv("RUN_PROCESS_REPLAY"): + @atexit.register + def save_process_replay() -> None: + for k,v in PROCESS_REPLAY_CAPTURE.items(): diskcache_put("schedule_process_replay", k, v, prepickled=True) + # **** Schedule grouping def uval(u:UOp) -> UOp: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 5278696b9c..35262eb633 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -191,7 +191,7 @@ def diskcache_get(table:str, key:Union[Dict, str, int]) -> Any: return None _db_tables = set() -def diskcache_put(table:str, key:Union[Dict, str, int], val:Any): +def diskcache_put(table:str, key:Union[Dict, str, int], val:Any, prepickled=False): if CACHELEVEL == 0: return val if isinstance(key, (str,int)): key = {"key": key} conn = db_connection() @@ -201,7 +201,7 @@ def diskcache_put(table:str, key:Union[Dict, str, int], val:Any): ltypes = ', '.join(f"{k} {TYPES[type(key[k])]}" for k in key.keys()) cur.execute(f"CREATE TABLE IF NOT EXISTS '{table}_{VERSION}' ({ltypes}, val blob, PRIMARY KEY ({', '.join(key.keys())}))") _db_tables.add(table) - cur.execute(f"REPLACE INTO '{table}_{VERSION}' ({', '.join(key.keys())}, val) VALUES ({', '.join(['?']*len(key.keys()))}, ?)", tuple(key.values()) + (pickle.dumps(val), )) # noqa: E501 + cur.execute(f"REPLACE INTO '{table}_{VERSION}' ({', '.join(key.keys())}, val) VALUES ({', '.join(['?']*len(key.keys()))}, ?)", tuple(key.values()) + (val if prepickled else pickle.dumps(val), )) # noqa: E501 conn.commit() cur.close() return val