tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-09-01 21:26:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

print best in MCTS + light up the winner in hcopt

Browse Source
This commit is contained in:
geohot
2024-07-20 09:39:36 -07:00
parent 0de5812032
commit 1113e47f96
2 changed files with 11 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/handcode_opt.py
+8 -5
View File
@@ -7,7 +7,7 @@ from tinygrad.codegen.kernel import Kernel
from tinygrad.device import Compiled
from tinygrad.engine.schedule import create_schedule
from tinygrad.engine.search import time_linearizer, beam_search, bufs_from_lin
from tinygrad.helpers import DEBUG, ansilen, getenv
from tinygrad.helpers import DEBUG, ansilen, getenv, colored
from tinygrad.ops import MetaOps
from tinygrad.shape.symbolic import sym_infer
@@ -113,11 +113,14 @@ if __name__ == "__main__":
tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
ops = lin.to_program().op_estimate
gflops = sym_infer(ops, {k:k.min for k in lin.ast.vars()})*1e-9/tm
choices.append((tm, gflops, lin.linearize()))
choices.append((tm, gflops, lin.linearize(), nm))
# print all kernels
if DEBUG >= 1: print(f" kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {nm}")
tm, gflops, lin = sorted(choices, key=lambda x: x[0])[0]
sorted_choices = sorted(choices, key=lambda x: x[0])
if DEBUG >= 1: # print all kernels
for tm, gflops, lin, nm in choices:
print(f" kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(lin.global_size):18s} {str(lin.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")
tm, gflops, lin, nm = sorted_choices[0]
if getenv("SRC"): print(lin.to_program().src)
total_tm += tm
running_gflops += gflops * tm
extra/mcts_search.py
+3 -3
View File
@@ -39,7 +39,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
node.parent.children.remove(node)
st = time.perf_counter()
best, best_tm = lin, math.inf
best, best_idx, best_tm = lin, 0, math.inf
for i in range(amt):
# tree traversal
node = root
@@ -66,8 +66,8 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
remove_node(node)
continue
if DEBUG>=2: print(f"\r{time.perf_counter() - st:7.2f}s: {tm:12.2f} us best: {best_tm:12.2f} us {i+1:4d}/{amt:4d} {node.kernel.colored_shape()}\033[K", end="") # noqa: E501
if tm < best_tm: best, best_tm = node.kernel, tm
if DEBUG>=2: print(f"\r{time.perf_counter() - st:7.2f}s: {tm:12.2f} us best: {best_tm:12.2f} us @ {best_idx+1:4d} {i+1:4d}/{amt:4d} {node.kernel.colored_shape()}\033[K", end="") # noqa: E501
if tm < best_tm: best, best_idx, best_tm = node.kernel, i, tm
# backprop
bnode: Optional[MCTSNode] = node