move hand_coded_optimizations to heuristic.py [pr] (#9844)

* move hand_coded_optimizations to heuristic.py [pr]

also folded all long lines

* make a copy and rename self -> k

* fix test

extra/optimization/get_action_space.py

@@ -2,6 +2,7 @@ import random
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from tinygrad.engine.search import actions
 from tinygrad.codegen.kernel import Kernel
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 from tinygrad.helpers import tqdm
 
 tactions = set()
@@ -23,7 +24,7 @@ if __name__ == "__main__":
   for ast_str in tqdm(ast_strs):
     lin = ast_str_to_lin(ast_str)
     #if not lin.apply_tensor_cores():
-    lin.hand_coded_optimizations()
+    lin = hand_coded_optimizations(lin)
     test_rebuild(lin)
     # confirm linearize can be called twice
     uops1 = lin.linearize().uops

extra/optimization/test_net.py

@@ -7,6 +7,7 @@ from tinygrad.helpers import getenv, colored
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
 from tinygrad.engine.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.pretrain_valuenet import ValueNet
@@ -34,7 +35,7 @@ if __name__ == "__main__":
     rawbufs = bufs_from_lin(lin)
 
     linhc = deepcopy(lin)
-    linhc.hand_coded_optimizations()
+    linhc = hand_coded_optimizations(linhc)
     tmhc = time_linearizer(linhc, rawbufs)
     print(f"{tmhc*1e6:10.2f}     HC    ", linhc.colored_shape())