diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9149f22f85..33fb2955c9 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -837,6 +837,8 @@ jobs: NV_PTX: 1 NV: 1 FORWARD_ONLY: 1 + # TODO: failing due to library loading error + CAPTURE_PROCESS_REPLAY: 0 run: | python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20 - name: Run process replay tests diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index 48ef3284ea..fe66c04f44 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -88,7 +88,8 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None: name, loc = "", "" try: name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0]) - ctx_vars = {k:v.value for k,v in ctx_vals.items() if k != "DEBUG" and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value} + ctx_vars = {k:v.value for k,v in ctx_vals.items() if k not in ("DEBUG", "CAPTURE_PROCESS_REPLAY") + and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value} if (replayer:=fxns.get(name)) is None: continue with Context(**ctx_vars): if (ret:=replayer(ret, *args, **kwargs)) is None: continue diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index b21250bb74..72ebb240e2 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -1178,7 +1178,8 @@ def _get_runner(inst_bytes: bytes, arch: str = "rdna3"): canonical_name = f"{_op_name(inst).lower()}_{base.to_bytes(size, 'little').hex()}" sink = sink.replace(arg=KernelInfo(name=canonical_name)).rtag(1) - with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES=""): + # NOTE: renderer output is not reproducible because of _MXCSRContext + with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES="", CAPTURE_PROCESS_REPLAY=0): runner = get_runner('CPU', sink) _canonical_runner_cache.append((base, mask, size, runner)) return runner diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 6f34133833..b22c8169d9 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -184,6 +184,7 @@ CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), Co ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0) MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0) EMULATE, EMULATED_DTYPES = ContextVar("EMULATE", ""), ContextVar("EMULATED_DTYPES", "") +CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0) CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1))) # Compilers CPU_CC, CPU_LLVM, CPU_LVP = ContextVar("CPU_CC", ""), ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index cf90d1429e..0039d54060 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -7,7 +7,7 @@ from tinygrad.uop import Ops, GroupOp from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, AddrSpace, ConstFloat, PyConst from tinygrad.dtype import storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA -from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC +from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, VIZ, SPEC, CAPTURE_PROCESS_REPLAY from tinygrad.helpers import strip_parens, colored, ansilen, printable, panic if TYPE_CHECKING: from tinygrad.device import Buffer, MultiBuffer @@ -1071,7 +1071,7 @@ tracked_keys:list[TracingKey] = [] tracked_ctxs:list[list[TrackedGraphRewrite]] = [] _name_cnt:dict[str, itertools.count] = {} -if getenv("CAPTURE_PROCESS_REPLAY"): +if CAPTURE_PROCESS_REPLAY: replay_capture: list[bytes] = [] import atexit, uuid @atexit.register @@ -1095,7 +1095,7 @@ def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=Fal assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}" tracked_keys[-1] = k = TracingKey(n:=tracked_keys[-1].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys) - if getenv("CAPTURE_PROCESS_REPLAY") and replay: + if CAPTURE_PROCESS_REPLAY and replay: # find the unittest frame we're capturing in frm = sys._getframe(1) while (f_back:=frm.f_back) is not None and "unittest" not in f_back.f_code.co_filename: frm = f_back