move kernel to opt (#10899)

examples/handcode_opt.py

@@ -2,11 +2,11 @@ from extra.models.resnet import ResNet50
 from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import Ops, sym_infer
 from tinygrad.device import Compiled
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
 from extra.optimization.helpers import time_linearizer
 

extra/assembly/assembly.py

@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.codegen.kernel import Ops, MemOp, UOp
+from tinygrad.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG

extra/assembly/assembly_arm64.py

@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.codegen.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
 

extra/assembly/assembly_ptx.py

@@ -1,7 +1,7 @@
 from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.codegen.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch

extra/assembly/assembly_rdna.py

@@ -2,7 +2,7 @@ import yaml
 from typing import Tuple, Set, Dict
 from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.codegen.kernel import Ops
+from tinygrad.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
 

extra/backends/triton.py

@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
 from collections import defaultdict
 from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.codegen.kernel import  UOp, Ops
+from tinygrad.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math

extra/gemm/amd_matmul.py

@@ -4,7 +4,7 @@ import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor, Device, Context
 from tinygrad.helpers import getenv
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 

extra/gemm/max_matmul.py

@@ -4,9 +4,9 @@ from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 
 # for copied uops
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
 from tinygrad.shape.shapetracker import ShapeTracker

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.kernel import OptOps
+from tinygrad.opt.kernel import OptOps
 from tinygrad.engine.realize import lower_schedule
 
 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float

extra/gemm/tinygrad_nv_matmul.py

@@ -1,6 +1,6 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from dataclasses import replace
 

extra/gemm/tvm_gemm.py

@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
 
 sched = C.schedule()
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.linearize()

extra/mcts_search.py

@@ -4,9 +4,9 @@ import numpy as np
 np.set_printoptions(suppress=True)
 import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.engine.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
 
 class MCTSNode:
   def __init__(self, kernel:Kernel, parent=None):

extra/optimization/extract_policynet.py

@@ -5,9 +5,9 @@ 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.engine.search import actions
+from tinygrad.opt.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import getenv
 
 # stuff needed to unpack a kernel
@@ -17,7 +17,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 INNER = 256
 class PolicyNet:

extra/optimization/extract_sa_pairs.py

@@ -10,11 +10,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 # more stuff
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.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,8 +1,8 @@
 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.opt.search import actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import tqdm
 
 tactions = set()

extra/optimization/helpers.py

@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -9,7 +9,7 @@ inf, nan = float('inf'), float('nan')
 UOps = Ops
 
 # kernel unpacker
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
 def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
 def kern_str_to_lin(kern_str:str, opts=None):
@@ -101,7 +101,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
   return ret
 
 from tinygrad.device import Device, Buffer
-from tinygrad.engine.search import _ensure_buffer_alloc, _time_program
+from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
 from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
 
 def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501

extra/optimization/pretrain_valuenet.py

@@ -1,4 +1,4 @@
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tqdm import tqdm, trange
 import math
 import random
@@ -14,7 +14,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 from extra.optimization.helpers import lin_to_feats, MAX_DIMS
 

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.engine.search import actions, bufs_from_lin, get_kernel_actions
+from tinygrad.opt.search import actions, bufs_from_lin, get_kernel_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, time_linearizer

extra/optimization/run_qnet.py

@@ -1,6 +1,6 @@
 from typing import List, Tuple
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import get_kernel_actions, actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import get_kernel_actions, actions
 
 _net = None
 def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:

extra/optimization/search.py

@@ -4,8 +4,8 @@ from extra.optimization.helpers import ast_str_to_lin, time_linearizer
 from tinygrad import dtypes
 from tinygrad.helpers import BEAM, getenv
 from tinygrad.device import Device, Compiled
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import beam_search, bufs_from_lin
 
 
 if __name__ == '__main__':

extra/optimization/test_net.py

@@ -6,8 +6,8 @@ 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.engine.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.opt.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

extra/optimization/test_time_linearizer.py

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

extra/replay_pkl.py

@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
 from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
 from tinygrad.renderer import ProgramSpec
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.heuristic import hand_coded_optimizations
 import numpy as np
 
 def move_jit_captured_to_dev(captured, device="DSP"):

setup.py

@@ -26,7 +26,7 @@ setup(name='tinygrad',
       long_description_content_type='text/markdown',
       packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
                   'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop'],
+                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'perfetto.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",

test/external/external_benchmark_hcopt.py

@@ -1,7 +1,7 @@
 import random
 from tinygrad.helpers import getenv
-from tinygrad.engine.search import beam_search, bufs_from_lin
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
 
 def optimize_kernel(k):

test/external/external_benchmark_schedule.py

@@ -3,10 +3,10 @@ from extra.models.resnet import ResNet50
 from tinygrad import Tensor, nn
 from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.uop.spec import type_verify
 
 if __name__ == "__main__":

test/external/external_benchmark_sdxl_softmax.py

@@ -7,7 +7,7 @@ if __name__ == "__main__":
   GlobalCounters.reset()
   t.softmax(-1, dtype="half", _single_kernel=True).realize()
 
-  from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+  from tinygrad.opt.kernel import Kernel, Opt, OptOps
   from tinygrad.helpers import get_single_element
   GlobalCounters.reset()
   si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())

test/external/external_debug_metal_sd_conv.py

@@ -1,8 +1,8 @@
 # ruff: noqa: E501
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

test/external/external_test_hcq_fuzz_failures.py

@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"
 
 import unittest, random
 import numpy as np
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from test.external.fuzz_linearizer import compare_linearizer, compare_states, get_fuzz_rawbuf_like
 

test/external/external_test_hip_compile.py

@@ -3,7 +3,7 @@ from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 class TestHIPCompileSpeed(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")

test/external/external_test_nv.py

@@ -2,12 +2,12 @@ import unittest, struct, array, ctypes
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.runtime.ops_nv import NVDevice, HWQueue
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from test.test_linearizer_failures import helper_test_lin
 from tinygrad.engine.realize import get_runner, CompiledRunner
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

test/external/external_test_train_gpt2.py

@@ -2,11 +2,11 @@
 import unittest
 
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 from test.external.fuzz_linearizer import run_linearizer
 

test/external/external_test_valid_remove.py

@@ -3,11 +3,11 @@ import unittest
 
 from tinygrad import Device
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 class TestOpenpilotValidhack(unittest.TestCase):
   def test_valid_removal(self):

test/external/fuzz_linearizer.py

@@ -20,9 +20,9 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
 
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.kernel import Opt, OptOps
-from tinygrad.engine.search import get_kernel_actions, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
 from tinygrad.uop.ops import UOp, Ops

test/external/process_replay/process_replay.py

@@ -4,7 +4,7 @@ import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertoo
 from typing import Callable, Any
 from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, to_function_name
 from tinygrad.engine.kernelize import get_kernelize_map
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.uop.ops import UOp, Ops
 
 # *** process replay settings

test/external/speed_beam_v_hcopt.py

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

test/external/speed_compare_amd_am.py

@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
 from tinygrad.helpers import getenv, colorize_float, DEBUG
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.heuristic import hand_coded_optimizations
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_amd import AMDDevice

test/external/speed_compare_cuda_nv.py

@@ -2,8 +2,8 @@ from tinygrad import Device, dtypes
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
-from tinygrad.codegen.heuristic import hand_coded_optimizations
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.tensor import _to_np_dtype
 import numpy as np

test/external/speed_compare_cuda_ptx.py

@@ -1,10 +1,10 @@
 import itertools
 from tinygrad import Device
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import getenv, colorize_float
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.runtime.ops_cuda import PTXCompiler, PTXRenderer, CUDACompiler
 
 if __name__ == "__main__":

test/external/verify_kernel.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from extra.optimization.helpers import kern_str_to_lin, time_linearizer
 from test.external.fuzz_linearizer import compare_linearizer
 from tinygrad.helpers import colored
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 # Use this with the LOGKERNS options to verify that all executed kernels are valid and evaluate to the same ground truth results
 

test/test_arange.py

@@ -3,9 +3,9 @@ import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
-from tinygrad.codegen.kernel import Opt, OptOps, Kernel, KernelOptError
+from tinygrad.opt.kernel import Opt, OptOps, Kernel, KernelOptError
 from tinygrad.engine.realize import CompiledRunner, ExecItem
-from tinygrad.engine.search import get_kernel_actions
+from tinygrad.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops
 
 class TestArange(unittest.TestCase):

test/test_hcq.py

@@ -5,7 +5,7 @@ from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.engine.realize import get_runner, CompiledRunner
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad import Variable
 
 MOCKGPU = getenv("MOCKGPU")

test/test_linearizer.py

@@ -4,7 +4,7 @@ import unittest
 from dataclasses import replace
 
 from test.helpers import ast_const
-from tinygrad.codegen.kernel import Opt, OptOps, KernelOptError, Kernel
+from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel
 from tinygrad.codegen.lowerer import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp
 from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -12,7 +12,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX
 from tinygrad.dtype import DType, dtypes
 
@@ -1136,7 +1136,7 @@ class TestLinearizer(unittest.TestCase):
         np.testing.assert_allclose(result, golden_result, atol=0.1, rtol=0.2)
 
       # check that get_kernel_actions produces all 9 options
-      from tinygrad.engine.search import get_kernel_actions
+      from tinygrad.opt.search import get_kernel_actions
       tc_actions = [k for i, k in get_kernel_actions(Kernel(realized_ast), False).items() if k.applied_opts[0].op == OptOps.TC]
 
       available_tc = len([x for x in Device[Device.DEFAULT].renderer.tensor_cores if x.dtype_in == tc.dtype_in and x.dtype_out == tc.dtype_out])

test/test_linearizer_dumb.py

@@ -8,8 +8,8 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.helpers import getenv
 from tinygrad.shape.shapetracker import ShapeTracker, View
-from tinygrad.engine.search import Opt, OptOps
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.search import Opt, OptOps
+from tinygrad.opt.kernel import Kernel
 
 class TestLinearizerDumb(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")

test/test_linearizer_failures.py

@@ -1,10 +1,10 @@
 # ruff: noqa: E501
 import unittest, random
 import numpy as np
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import CI, Context
 from test.external.fuzz_linearizer import compare_linearizer

test/test_linearizer_overflows.py

@@ -2,8 +2,8 @@
 import unittest
 from tinygrad import dtypes, Device
 from tinygrad.helpers import CI
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import Opt, OptOps, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
 from extra.optimization.helpers import time_linearizer
 
 # stuff needed to unpack a kernel

test/test_opt_gemm.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad import Tensor
 from tinygrad.helpers import get_single_element
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 
 class TestOptGemm(unittest.TestCase):

test/test_quantize_onnx.py

@@ -4,9 +4,9 @@ import unittest
 from dataclasses import replace
 from tinygrad import Tensor, Context, Device, dtypes
 from tinygrad.uop.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.opt.search import bufs_from_lin
 from tinygrad.shape.shapetracker import ShapeTracker, View # noqa: F401 # pylint: disable=unused-import
 
 N = 512

test/test_search.py

@@ -1,8 +1,8 @@
 import unittest
 
-from tinygrad.codegen.kernel import Opt, OptOps, Kernel
+from tinygrad.opt.kernel import Opt, OptOps, Kernel
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import bufs_from_lin, actions, beam_search
+from tinygrad.opt.search import bufs_from_lin, actions, beam_search
 from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
@@ -39,7 +39,7 @@ class TestBEAM(unittest.TestCase):
     a = Tensor.rand(4, 3)
     b = Tensor.rand(3)
     realized_ast, _ = helper_realized_ast(a @ b)
-    from tinygrad.engine.search import get_kernel_actions
+    from tinygrad.opt.search import get_kernel_actions
     lins = get_kernel_actions(Kernel(realized_ast), False).values()
 
     # ensure amt=0 are not duplicated
@@ -57,7 +57,7 @@ class TestBEAM(unittest.TestCase):
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_search_over_shape(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.engine.search import get_kernel_actions
+    from tinygrad.opt.search import get_kernel_actions
 
     dtype_pairs = [(tc.dtype_in, tc.dtype_out) for tc in Device[Device.DEFAULT].renderer.tensor_cores]
     multi_shape_dtype_pairs = [dts for dts in dtype_pairs if dtype_pairs.count(dts) > 1]
@@ -74,7 +74,7 @@ class TestBEAM(unittest.TestCase):
 
   def test_get_kernel_actions_preserves_actions_state(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.engine.search import get_kernel_actions
+    from tinygrad.opt.search import get_kernel_actions
     a = Tensor.rand(16, 16)
     b = Tensor.rand(16, 16)
     realized_ast, _ = helper_realized_ast(a @ b)

test/test_uops.py

@@ -15,7 +15,7 @@ from tinygrad.engine.realize import CompiledRunner
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 
 def to_uops_list(u:list[UOp], opts=None, skip_check=False) -> list[UOp]: return full_rewrite(UOp.sink(*u), opts)
 

test/test_uops_stats.py

@@ -6,7 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
 from tinygrad.device import Device
 
 def flops_mem(uops, ignore_indexing=False):

test/test_winograd.py

@@ -3,8 +3,8 @@ import numpy as np
 from tinygrad import Tensor, GlobalCounters, dtypes, Context, nn
 from tinygrad.uop.ops import Ops
 from tinygrad.helpers import Timing, CI, Profiling, WINO, DEBUG, getenv
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 
 class TestWinogradClose(unittest.TestCase):
   def test_close(self):

test/unit/test_linearizer_rewrite.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Context, Device
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 
 class TestLinearizerRewrite(unittest.TestCase):
   def test_reduction(self):

test/unit/test_search.py

@@ -1,8 +1,8 @@
 import unittest
 from tinygrad import Tensor, Device
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import Buffer
-from tinygrad.engine.search import get_test_global_size, bufs_from_lin
+from tinygrad.opt.search import get_test_global_size, bufs_from_lin
 from tinygrad.helpers import GlobalCounters
 from extra.optimization.helpers import time_linearizer
 

test/unit/test_uop_spec.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import unittest
 
 from tinygrad import Tensor
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import DEBUG
 from tinygrad.uop.ops import UOp, Ops, print_uops
 from tinygrad.uop.spec import type_verify, ast_spec, tensor_uop_spec

tinygrad/engine/realize.py

@@ -6,8 +6,8 @@ from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.engine.schedule import ScheduleItem
 
 # **************** Program Creation ****************
@@ -18,7 +18,7 @@ def get_program(renderer:Renderer, ast:UOp) -> ProgramSpec:
   if not NOOPT:
     if not k.apply_tensor_cores(getenv("TC", 1)): k.apply_opts(hand_coded_optimizations(k))
     if BEAM >= 1:
-      from tinygrad.engine.search import beam_search, bufs_from_lin
+      from tinygrad.opt.search import beam_search, bufs_from_lin
       kb = Kernel(ast, opts=renderer)
       rawbufs = bufs_from_lin(kb, allocate=False)
       k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
@@ -52,7 +52,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.engine.search import optimize_local_size
+      from tinygrad.opt.search import optimize_local_size
       local_size = optimize_local_size(self._prg, 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/opt/heuristic.py

@@ -1,5 +1,5 @@
 import itertools
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
 from tinygrad.helpers import getenv, DEBUG, all_int, prod
 from tinygrad.dtype import ImageDType
 from tinygrad.uop.ops import Ops, resolve

tinygrad/opt/search.py

@@ -7,7 +7,7 @@ from tinygrad.device import Device, Buffer, Compiler
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE, TC_SEARCH_OVER_SHAPE
 from tinygrad.dtype import ImageDType, PtrDType
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.opt.kernel import Kernel, Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.renderer import ProgramSpec

tinygrad/viz/serve.py

@@ -5,7 +5,7 @@ from urllib.parse import parse_qs, urlparse
 from typing import Any, TypedDict, Generator
 from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA
 from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, lines, GroupOp, srender, sint
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent
 from tinygrad.dtype import dtypes