From ff64bcab69014689eae148fb2809f763547f4aa0 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Tue, 14 May 2024 23:12:59 -0700
Subject: [PATCH] move graph/search to engine (#4596)

---
 docs/abstractions2.py                        | 2 +-
 examples/handcode_resnet50_opt.py            | 4 ++--
 extra/optimization/extract_policynet.py      | 2 +-
 extra/optimization/extract_sa_pairs.py       | 2 +-
 extra/optimization/get_action_space.py       | 2 +-
 extra/optimization/rl.py                     | 2 +-
 extra/optimization/run_qnet.py               | 2 +-
 extra/optimization/search.py                 | 2 +-
 extra/optimization/test_net.py               | 2 +-
 extra/optimization/test_time_linearizer.py   | 2 +-
 test/external/external_test_nv.py            | 2 +-
 test/external/external_test_uops_graphing.py | 2 +-
 test/external/fuzz_linearizer.py             | 4 ++--
 test/external/replay_schedule.py             | 2 +-
 test/external/speed_beam_v_hcopt.py          | 2 +-
 test/external/speed_compare_cuda_nv.py       | 2 +-
 test/external/speed_compare_cuda_ptx.py      | 2 +-
 test/external/verify_kernel.py               | 4 ++--
 test/test_linearizer.py                      | 2 +-
 test/test_linearizer_failures.py             | 2 +-
 test/test_linearizer_overflows.py            | 4 ++--
 test/test_schedule.py                        | 2 +-
 test/test_search.py                          | 4 ++--
 tinygrad/codegen/uops.py                     | 2 +-
 tinygrad/{features => engine}/graph.py       | 0
 tinygrad/engine/realize.py                   | 6 +++---
 tinygrad/engine/schedule.py                  | 2 +-
 tinygrad/{features => engine}/search.py      | 0
 28 files changed, 33 insertions(+), 33 deletions(-)
 rename tinygrad/{features => engine}/graph.py (100%)
 rename tinygrad/{features => engine}/search.py (100%)

diff --git a/docs/abstractions2.py b/docs/abstractions2.py
index 9b95e787a4..0caea2d326 100644
--- a/docs/abstractions2.py
+++ b/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
 
diff --git a/examples/handcode_resnet50_opt.py b/examples/handcode_resnet50_opt.py
index da92147825..2050018953 100644
--- a/examples/handcode_resnet50_opt.py
+++ b/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"):
diff --git a/extra/optimization/extract_policynet.py b/extra/optimization/extract_policynet.py
index c751d74e56..84b17b4ee0 100644
--- a/extra/optimization/extract_policynet.py
+++ b/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
diff --git a/extra/optimization/extract_sa_pairs.py b/extra/optimization/extract_sa_pairs.py
index 6e9de99977..53df9ec34f 100644
--- a/extra/optimization/extract_sa_pairs.py
+++ b/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
diff --git a/extra/optimization/get_action_space.py b/extra/optimization/get_action_space.py
index e2ded91382..b1fe39c141 100644
--- a/extra/optimization/get_action_space.py
+++ b/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()
diff --git a/extra/optimization/rl.py b/extra/optimization/rl.py
index a088ae337f..801df84cec 100644
--- a/extra/optimization/rl.py
+++ b/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
diff --git a/extra/optimization/run_qnet.py b/extra/optimization/run_qnet.py
index f456f8adf2..c1f51d4a1f 100644
--- a/extra/optimization/run_qnet.py
+++ b/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]]:
diff --git a/extra/optimization/search.py b/extra/optimization/search.py
index df62feca3a..36ec8366c7 100644
--- a/extra/optimization/search.py
+++ b/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__':
diff --git a/extra/optimization/test_net.py b/extra/optimization/test_net.py
index 851c2a85bd..4adf817b20 100644
--- a/extra/optimization/test_net.py
+++ b/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
diff --git a/extra/optimization/test_time_linearizer.py b/extra/optimization/test_time_linearizer.py
index 47af6c52b7..3b5d672918 100644
--- a/extra/optimization/test_time_linearizer.py
+++ b/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()
diff --git a/test/external/external_test_nv.py b/test/external/external_test_nv.py
index 525af6bc61..996a3b8738 100644
--- a/test/external/external_test_nv.py
+++ b/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
diff --git a/test/external/external_test_uops_graphing.py b/test/external/external_test_uops_graphing.py
index 2f64b6f274..cd4c20ede2 100644
--- a/test/external/external_test_uops_graphing.py
+++ b/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
 
diff --git a/test/external/fuzz_linearizer.py b/test/external/fuzz_linearizer.py
index 4c19807350..d237af50df 100644
--- a/test/external/fuzz_linearizer.py
+++ b/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
diff --git a/test/external/replay_schedule.py b/test/external/replay_schedule.py
index 53f97fc000..b8f1a76861 100755
--- a/test/external/replay_schedule.py
+++ b/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
 
diff --git a/test/external/speed_beam_v_hcopt.py b/test/external/speed_beam_v_hcopt.py
index bfe14cf322..447a891ea7 100644
--- a/test/external/speed_beam_v_hcopt.py
+++ b/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__":
diff --git a/test/external/speed_compare_cuda_nv.py b/test/external/speed_compare_cuda_nv.py
index 07953bc479..9a6b0cc1a1 100644
--- a/test/external/speed_compare_cuda_nv.py
+++ b/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):
diff --git a/test/external/speed_compare_cuda_ptx.py b/test/external/speed_compare_cuda_ptx.py
index aeb19a1aa2..854e1db595 100644
--- a/test/external/speed_compare_cuda_ptx.py
+++ b/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?
diff --git a/test/external/verify_kernel.py b/test/external/verify_kernel.py
index 24dfe9af4a..16399c9864 100644
--- a/test/external/verify_kernel.py
+++ b/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
 
diff --git a/test/test_linearizer.py b/test/test_linearizer.py
index 9c9654fff8..e8ad02e639 100644
--- a/test/test_linearizer.py
+++ b/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)}"
 
diff --git a/test/test_linearizer_failures.py b/test/test_linearizer_failures.py
index 31d4c6eb75..f16b6538d2 100644
--- a/test/test_linearizer_failures.py
+++ b/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
diff --git a/test/test_linearizer_overflows.py b/test/test_linearizer_overflows.py
index 3202fe4cd5..21373ad8b3 100644
--- a/test/test_linearizer_overflows.py
+++ b/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
diff --git a/test/test_schedule.py b/test/test_schedule.py
index 74aefad7f6..cedecf5616 100644
--- a/test/test_schedule.py
+++ b/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
diff --git a/test/test_search.py b/test/test_search.py
index f376a1a4c0..4ef969af3f 100644
--- a/test/test_search.py
+++ b/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
diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py
index 92a2c21642..f663cef93a 100644
--- a/tinygrad/codegen/uops.py
+++ b/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):
diff --git a/tinygrad/features/graph.py b/tinygrad/engine/graph.py
similarity index 100%
rename from tinygrad/features/graph.py
rename to tinygrad/engine/graph.py
diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py
index dc41444381..2da25cda50 100644
--- a/tinygrad/engine/realize.py
+++ b/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)
diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py
index 8dd1f94bea..8a0e00a0f9 100644
--- a/tinygrad/engine/schedule.py
+++ b/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
diff --git a/tinygrad/features/search.py b/tinygrad/engine/search.py
similarity index 100%
rename from tinygrad/features/search.py
rename to tinygrad/engine/search.py