move graph/search to engine (#4596)

test/external/external_test_nv.py

@@ -3,7 +3,7 @@ from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.engine.schedule import create_schedule
 from tinygrad.runtime.ops_nv import NVDevice, HWComputeQueue
-from tinygrad.features.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from test.test_linearizer_failures import helper_test_lin
 
 from tinygrad.codegen.linearizer import Linearizer

test/external/external_test_uops_graphing.py

@@ -3,7 +3,7 @@ import unittest
 from tinygrad.tensor import Tensor
 from tinygrad.codegen.linearizer import Linearizer
 from tinygrad.renderer.cstyle import OpenCLRenderer
-from tinygrad.features.graph import graph_uops
+from tinygrad.engine.graph import graph_uops
 from tinygrad.engine.schedule import create_schedule
 from tinygrad.nn import Conv2d
 

test/external/fuzz_linearizer.py

@@ -7,8 +7,8 @@ from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.codegen.linearizer import Linearizer, UOp
 from tinygrad.codegen.kernel import Opt, OptOps
-from tinygrad.features.search import get_linearizer_actions, bufs_from_lin
-from tinygrad.features.graph import print_tree
+from tinygrad.engine.search import get_linearizer_actions, bufs_from_lin
+from tinygrad.engine.graph import print_tree
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG
 from tinygrad.ops import LazyOp, UnaryOps, BufferOps

test/external/replay_schedule.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import subprocess, pickle, shlex, sys
 from typing import Dict, List, Tuple
-from tinygrad.features.graph import print_tree
+from tinygrad.engine.graph import print_tree
 from tinygrad.helpers import colored
 from tinygrad.ops import LazyOp
 

test/external/speed_beam_v_hcopt.py

@@ -1,6 +1,6 @@
 from tinygrad import Device
 from tinygrad.helpers import getenv, DEBUG, BEAM
-from tinygrad.features.search import beam_search, time_linearizer, bufs_from_lin
+from tinygrad.engine.search import beam_search, time_linearizer, bufs_from_lin
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 
 if __name__ == "__main__":

test/external/speed_compare_cuda_nv.py

@@ -1,7 +1,7 @@
 from tinygrad import Device
 from tinygrad.helpers import getenv, colored
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.features.search import bufs_from_lin
+from tinygrad.engine.search import bufs_from_lin
 
 # move to helpers?
 def colorize_float(x):

test/external/speed_compare_cuda_ptx.py

@@ -3,7 +3,7 @@ from tinygrad import Device
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, colored
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.features.search import bufs_from_lin
+from tinygrad.engine.search import bufs_from_lin
 from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
 
 # move to helpers?

test/external/verify_kernel.py

@@ -4,8 +4,8 @@ from extra.optimization.helpers import kern_str_to_lin
 from test.external.fuzz_linearizer import compare_linearizer
 from tinygrad.helpers import colored
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.graph import print_tree
-from tinygrad.features.search import time_linearizer
+from tinygrad.engine.graph import print_tree
+from tinygrad.engine.search import time_linearizer
 
 # Use this with the LOGKERN options to verify that all executed kernels are valid and evaluate to the same ground truth results
 

test/test_linearizer.py

@@ -322,7 +322,7 @@ class TestLinearizer(unittest.TestCase):
         np.testing.assert_allclose(result, golden_result, atol=0.1, rtol=0.15)
 
       # check that get_linearizer_actions produces all 9 options
-      from tinygrad.features.search import get_linearizer_actions
+      from tinygrad.engine.search import get_linearizer_actions
       tc_actions = [k for i, k in get_linearizer_actions(Linearizer(realized_ast), False).items() if k.applied_opts[0].op == OptOps.TC]
       assert len(tc_actions) == 9, f"get_linearizer_actions should contain 9 possible TC actions, only got {len(tc_actions)}"
 

test/test_linearizer_failures.py

@@ -3,7 +3,7 @@ import unittest, random
 import numpy as np
 from tinygrad.codegen.kernel import KernelOptError
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import CI
 from test.external.fuzz_linearizer import compare_linearizer

test/test_linearizer_overflows.py

@@ -3,8 +3,8 @@ import unittest
 from tinygrad import dtypes, Device
 from tinygrad.helpers import CI
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.search import Opt, OptOps
-from tinygrad.features.search import time_linearizer, bufs_from_lin
+from tinygrad.engine.search import Opt, OptOps
+from tinygrad.engine.search import time_linearizer, bufs_from_lin
 
 # stuff needed to unpack a kernel
 from tinygrad.ops import LazyOp, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer

test/test_schedule.py

@@ -10,7 +10,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.ops import BinaryOps, LoadOps, ReduceOps
 from tinygrad.helpers import DEBUG, flatten
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.graph import print_tree
+from tinygrad.engine.graph import print_tree
 from tinygrad.engine.schedule import create_schedule
 from tinygrad import nn, dtypes
 from test.helpers import is_dtype_supported

test/test_search.py

@@ -3,7 +3,7 @@ import unittest
 from tinygrad.codegen.kernel import Opt, OptOps
 from tinygrad.codegen.linearizer import Linearizer
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.features.search import time_linearizer, bufs_from_lin, actions
+from tinygrad.engine.search import time_linearizer, bufs_from_lin, actions
 from tinygrad.device import Device, Buffer
 from tinygrad.ops import LoadOps, BufferOps
 from tinygrad.tensor import Tensor
@@ -48,7 +48,7 @@ class TestBEAM(unittest.TestCase):
     a = Tensor.rand(4, 3)
     b = Tensor.rand(3)
     realized_ast, _ = helper_realized_ast(a @ b)
-    from tinygrad.features.search import get_linearizer_actions
+    from tinygrad.engine.search import get_linearizer_actions
     lins = get_linearizer_actions(Linearizer(realized_ast), False).values()
 
     # ensure amt=0 are not duplicated