move graph/search to engine (#4596)

docs/abstractions2.py

@@ -93,7 +93,7 @@ sched = create_schedule([out])
 for si in sched: print(si.ast[0].op)  # NOTE: the first two convert it to CLANG
 
 # DEBUGGING: print the compute ast as a tree
-from tinygrad.features.graph import print_tree
+from tinygrad.engine.graph import print_tree
 print_tree(sched[-1].ast[0])
 # NOTE: sched[-1].ast is the same as st_0 above
 

examples/handcode_resnet50_opt.py

@@ -4,12 +4,12 @@ from tinygrad import Tensor, nn
 from tinygrad.ops import LoadOps, get_lazyop_info
 from tinygrad.device import Device, Compiled
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.search import time_linearizer, beam_search, bufs_from_lin
+from tinygrad.engine.search import time_linearizer, beam_search, bufs_from_lin
 from tinygrad.helpers import ansilen, DEBUG, getenv
 from tinygrad.shape.symbolic import sym_infer
 from tinygrad.dtype import dtypes
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.features.graph import print_tree
+from tinygrad.engine.graph import print_tree
 
 if __name__ == "__main__":
   if getenv("HALF"):

extra/optimization/extract_policynet.py

@@ -5,7 +5,7 @@ from tinygrad.nn import Linear
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import Adam
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.features.search import actions
+from tinygrad.engine.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
 from tinygrad.codegen.linearizer import Linearizer
 from tinygrad.helpers import getenv

extra/optimization/extract_sa_pairs.py

@@ -14,7 +14,7 @@ from tinygrad.codegen.kernel import Opt, OptOps
 
 # more stuff
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.search import actions
+from tinygrad.engine.search import actions
 from extra.optimization.helpers import lin_to_feats
 from extra.optimization.pretrain_valuenet import ValueNet
 from tinygrad.nn.optim import Adam

extra/optimization/get_action_space.py

@@ -1,7 +1,7 @@
 import random
 from tqdm import tqdm
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.features.search import actions
+from tinygrad.engine.search import actions
 from tinygrad.codegen.linearizer import Linearizer
 
 tactions = set()

extra/optimization/rl.py

@@ -3,7 +3,7 @@ import numpy as np
 import math, random
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.features.search import actions, bufs_from_lin, time_linearizer, get_linearizer_actions
+from tinygrad.engine.search import actions, bufs_from_lin, time_linearizer, get_linearizer_actions
 from tinygrad.nn.optim import Adam
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats

extra/optimization/run_qnet.py

@@ -1,6 +1,6 @@
 from typing import List, Tuple
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.search import get_linearizer_actions, actions
+from tinygrad.engine.search import get_linearizer_actions, actions
 
 _net = None
 def beam_q_estimate(beam:List[Tuple[Linearizer, float]]) -> List[Tuple[Linearizer, float]]:

extra/optimization/search.py

@@ -5,7 +5,7 @@ from tinygrad import dtypes
 from tinygrad.helpers import BEAM, getenv
 from tinygrad.device import Device, Compiled
 from tinygrad.codegen.linearizer import Linearizer
-from tinygrad.features.search import time_linearizer, beam_search, bufs_from_lin
+from tinygrad.engine.search import time_linearizer, beam_search, bufs_from_lin
 
 
 if __name__ == '__main__':

extra/optimization/test_net.py

@@ -6,7 +6,7 @@ from copy import deepcopy
 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.features.search import bufs_from_lin, time_linearizer, actions, get_linearizer_actions
+from tinygrad.engine.search import bufs_from_lin, time_linearizer, actions, get_linearizer_actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.pretrain_valuenet import ValueNet

extra/optimization/test_time_linearizer.py

@@ -1,5 +1,5 @@
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.features.search import bufs_from_lin, time_linearizer, get_linearizer_actions
+from tinygrad.engine.search import bufs_from_lin, time_linearizer, get_linearizer_actions
 
 if __name__ == "__main__":
   ast_strs = load_worlds()

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

tinygrad/codegen/uops.py

@@ -141,7 +141,7 @@ class UOpGraph:
   def globals(self) -> List[Tuple[int, bool]]: return [x.arg for x in self.uops if x.uop is UOps.DEFINE_GLOBAL]
 
   def graph(self):
-    from tinygrad.features.graph import graph_uops
+    from tinygrad.engine.graph import graph_uops
     graph_uops(self.uops)
 
   def print(self):

tinygrad/engine/realize.py

@@ -14,14 +14,14 @@ from tinygrad.engine.schedule import ScheduleItem
 logkerns, logkerns_level = open(getenv("LOGKERNS", ""), "a") if getenv("LOGKERNS", "") else None, getenv("LOGKERNS_LEVEL", 1)
 def get_linearizer(renderer:Renderer, ast:Tuple[LazyOp, ...]) -> Linearizer:
   if DEBUG >= 3:
-    from tinygrad.features.graph import print_tree
+    from tinygrad.engine.graph import print_tree
     for op in ast: print_tree(op)
   k = Linearizer(*ast, opts=renderer)
   k.required_optimizations()
   if not NOOPT:
     if not (used_tensor_cores:=k.apply_tensor_cores(getenv("TC", 1))): k.hand_coded_optimizations()
     if BEAM >= 1:
-      from tinygrad.features.search import beam_search, time_linearizer, bufs_from_lin
+      from tinygrad.engine.search import beam_search, time_linearizer, bufs_from_lin
       kb, k_opt = Linearizer(*ast, opts=renderer), k
       kb.required_optimizations()
       rawbufs = bufs_from_lin(kb, allocate=False)
@@ -67,7 +67,7 @@ class CompiledRunner(Runner):
     global_size, local_size = self.p.launch_dims(var_vals)
     if global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type]
       # TODO: this is copied from get_program
-      from tinygrad.features.search import optimize_local_size
+      from tinygrad.engine.search import optimize_local_size
       local_size = optimize_local_size(self.clprg, global_size, rawbufs)
       global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
       self.p = replace(self.p, global_size=global_size, local_size=local_size)

tinygrad/engine/schedule.py

@@ -3,7 +3,7 @@ from collections import defaultdict, deque
 from dataclasses import dataclass
 from typing import Tuple, List, Dict, Optional, Set, DefaultDict
 from tinygrad.ops import LoadOps, BufferOps, LazyOp, ReduceOps, ConstBuffer, MemBuffer, UNSAFE_PAD_OPS, UnaryOps
-from tinygrad.features.graph import log_lazybuffer, realized_lazybuffer
+from tinygrad.engine.graph import log_lazybuffer, realized_lazybuffer
 from tinygrad.helpers import GRAPH, DEBUG, MULTIOUTPUT, SAVE_SCHEDULE, GlobalCounters, prod, dedup, all_int, merge_dicts, getenv
 from tinygrad.shape.symbolic import Variable
 from tinygrad.dtype import ImageDType, dtypes