diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3b00c4c9..befd883d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -348,12 +348,11 @@ jobs: run: PYTHONPATH="." METAL=1 CACHELEVEL=0 FUZZ_ALL_ACTIONS=1 DEPTH=2 FUZZ_N=48 FUZZ_MAX_SIZE=10000000 python test/external/fuzz_linearizer.py - name: Fuzz Test models schedule run: FUZZ_SCHEDULE=1 FUZZ_SCHEDULE_MAX_PATHS=5 python -m pytest test/models/test_train.py test/models/test_end2end.py + - name: Run TRANSCENDENTAL math + run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 - name: Run process replay tests if: env.RUN_PROCESS_REPLAY == '1' run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git checkout origin/master && PYTHONPATH=. python3 process_replay.py - - name: Run TRANSCENDENTAL math - # put this after process replay since it has same ast - run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 # testwebgl: # name: WebGL Tests @@ -544,12 +543,11 @@ jobs: PYTHONPATH="." python examples/compile_efficientnet.py > recognize.c clang -O2 recognize.c -lm -o recognize cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock + - name: Run TRANSCENDENTAL math + run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 - name: Run process replay tests if: env.RUN_PROCESS_REPLAY == '1' run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git checkout origin/master && PYTHONPATH=. python3 process_replay.py - - name: Run TRANSCENDENTAL math - # put this after process replay since it has same ast - run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20 #testunicorn: # name: ARM64 unicorn Test diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index c7dc944f..e81cc242 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -2,7 +2,7 @@ # compare kernels created by HEAD against master import difflib, pickle from tinygrad.codegen.linearizer import Linearizer -from tinygrad.helpers import colored, db_connection, VERSION, getenv, tqdm +from tinygrad.helpers import Context, ContextVar, colored, db_connection, VERSION, getenv, tqdm page_size = 100 conn = db_connection() @@ -11,16 +11,17 @@ row_count = cur.execute(f"select count(*) from 'process_replay_{VERSION}'").fetc for offset in tqdm(range(0, row_count, page_size)): cur.execute(f"SELECT val FROM 'process_replay_{VERSION}' LIMIT ? OFFSET ?", (page_size, offset)) for row in cur.fetchall(): - ast, opts, applied_opts, name, compare_src = pickle.loads(row[0]) - k = Linearizer(*ast, opts=opts) - for opt in applied_opts: k.apply_opt(opt) - good_src = k.opts.render(name, k.linearize().uops) - try: assert compare_src == good_src - except AssertionError as e: - print("PROCESS REPLAY DETECTED CHANGE") - print(ast) - print(applied_opts) - diff = list(difflib.unified_diff(good_src.splitlines(), compare_src.splitlines())) - for line in diff: - print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None)) - if getenv("ASSERT_PROCESS_REPLAY", 1): raise e + ast, opts, applied_opts, name, compare_src, ctx = pickle.loads(row[0]) + with Context(**{k:v for k,v in ctx.items() if k in ContextVar._cache}): + k = Linearizer(*ast, opts=opts) + for opt in applied_opts: k.apply_opt(opt) + good_src = k.opts.render(name, k.linearize().uops) + try: assert compare_src == good_src + except AssertionError as e: + print("PROCESS REPLAY DETECTED CHANGE") + print(ast) + print(applied_opts) + diff = list(difflib.unified_diff(good_src.splitlines(), compare_src.splitlines())) + for line in diff: + print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None)) + if getenv("ASSERT_PROCESS_REPLAY", 1): raise e diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index 805f3990..e8678ba1 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -8,7 +8,7 @@ from tinygrad.ops import BufferOps, LazyOp, TernaryOps, ReduceOps, UnaryOps, get from tinygrad.codegen.uops import UOp, flops_mem, UOps from tinygrad.codegen.uopgraph import UOpGraph from tinygrad.renderer import Program -from tinygrad.helpers import to_function_name, DEBUG, getenv, prod, diskcache_put +from tinygrad.helpers import to_function_name, DEBUG, getenv, prod, diskcache_put, ContextVar # TODO: this needs to be replaced, there shouldn't be variables in the shapetracker def variable_to_uop(x, ctx=None) -> UOp: @@ -136,7 +136,8 @@ class Lowerer(Kernel): def to_program(self) -> Program: self.linearize() src = self.opts.render(name:=to_function_name(self.name), self.uops) - if getenv("RUN_PROCESS_REPLAY"): diskcache_put("process_replay", id(self), (self.ast, self.opts, self.applied_opts, name, src)) + if getenv("RUN_PROCESS_REPLAY"): + diskcache_put("process_replay", id(self), (self.ast, self.opts, self.applied_opts, name, src, {k:v.value for k,v in ContextVar._cache.items()})) info = get_lazyop_info(self.ast[0]) ops, mem = flops_mem(self.uops.uops) run_count = prod((self.global_size or []) + (self.local_size or []))