mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-02-11 15:15:13 -05:00
50 lines
1.9 KiB
Python
Executable File
50 lines
1.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# compare kernels created by HEAD against master
|
|
import difflib, pickle, multiprocessing
|
|
from tinygrad.codegen.kernel import Kernel
|
|
from tinygrad.helpers import Context, ContextVar, colored, db_connection, VERSION, getenv, tqdm
|
|
|
|
page_size = 100
|
|
table_name = f"process_replay_{getenv('GITHUB_RUN_ID', 'HEAD')}_{VERSION}"
|
|
|
|
def process_replay(offset:int):
|
|
conn = db_connection()
|
|
cur = conn.cursor()
|
|
cur.execute(f"SELECT val FROM '{table_name}' LIMIT ? OFFSET ?", (page_size, offset))
|
|
for row in cur.fetchall():
|
|
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}):
|
|
# try linearize
|
|
try:
|
|
k = Kernel(ast, opts=opts)
|
|
for opt in applied_opts: k.apply_opt(opt)
|
|
good_src = k.opts.render(name, k.linearize().uops)
|
|
except Exception as e:
|
|
print("FAILED TO RECREATE KERNEL")
|
|
print(ast)
|
|
print(applied_opts)
|
|
print(e)
|
|
if getenv("ASSERT_PROCESS_REPLAY", 1): raise e
|
|
continue
|
|
# try compare
|
|
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
|
|
conn.commit()
|
|
cur.close()
|
|
|
|
if __name__ == "__main__":
|
|
conn = db_connection()
|
|
cur = conn.cursor()
|
|
row_count = cur.execute(f"select count(*) from '{table_name}'").fetchone()[0]
|
|
conn.commit()
|
|
cur.close()
|
|
offsets = range(0, row_count, page_size)
|
|
with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool: list(tqdm(pool.imap(process_replay, offsets), total=len(offsets)))
|