41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
|
import random
|
||
|
|
from tinygrad.helpers import getenv
|
||
|
|
from tinygrad.engine.search import time_linearizer, beam_search, bufs_from_lin
|
||
|
|
from extra.optimization.helpers import load_worlds, ast_str_to_lin
|
||
|
|
|
||
|
|
def optimize_kernel(k):
|
||
|
|
# TODO: update this
|
||
|
|
return k.hand_coded_optimizations()
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
hcopt_wins = beam_wins = tie = 0
|
||
|
|
hcopt_total = beam_total = 0.0
|
||
|
|
|
||
|
|
worlds = load_worlds(filter_reduce=False, filter_noimage=True, filter_novariable=False)
|
||
|
|
random.seed(0)
|
||
|
|
random.shuffle(worlds)
|
||
|
|
|
||
|
|
for world in worlds[:500]:
|
||
|
|
k = ast_str_to_lin(world)
|
||
|
|
rawbufs = bufs_from_lin(k)
|
||
|
|
|
||
|
|
k_hcopt = optimize_kernel(k.copy())
|
||
|
|
k_beam = beam_search(k.copy(), rawbufs, getenv("BEAM", 2))
|
||
|
|
|
||
|
|
disable_cache = bool(getenv("NOCACHE", 0))
|
||
|
|
t_hcopt = time_linearizer(k_hcopt, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
|
||
|
|
t_beam = time_linearizer(k_beam, rawbufs, allow_test_size=False, cnt=10, disable_cache=disable_cache, clear_l2=True) * 1e6
|
||
|
|
|
||
|
|
if t_hcopt == t_beam: tie += 1
|
||
|
|
elif t_hcopt < t_beam: hcopt_wins += 1
|
||
|
|
else: beam_wins += 1
|
||
|
|
hcopt_total += t_hcopt
|
||
|
|
beam_total += t_beam
|
||
|
|
|
||
|
|
print(f"{t_hcopt=:5.2f} {k_hcopt.applied_opts=}")
|
||
|
|
print("")
|
||
|
|
print(f"{t_beam=:5.2f} {k_beam.applied_opts=}")
|
||
|
|
print("*"*20)
|
||
|
|
|
||
|
|
print(f"{hcopt_wins=}, {beam_wins=}, {tie=}")
|
||
|
|
print(f"{hcopt_total=:.2f}, {beam_total=:.2f}")
|