move opt under codegen (#11569)

docs/developer/layout.md

@@ -18,11 +18,11 @@ Group UOps into kernels.
 
 ---
 
-## tinygrad/opt
+## tinygrad/codegen/opt
 
 Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
 
-::: tinygrad.opt.get_optimized_ast
+::: tinygrad.codegen.opt.get_optimized_ast
     options:
         members: false
         show_labels: false

extra/assembly/assembly.py

@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.opt.kernel import Ops, MemOp, UOp
+from tinygrad.codegen.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.opt.kernel import Ops, UOp
+from tinygrad.codegen.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.opt.kernel import Ops, UOp
+from tinygrad.codegen.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.opt.kernel import Ops
+from tinygrad.codegen.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.opt.kernel import  UOp, Ops
+from tinygrad.codegen.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math

extra/gemm/amd_uop_matmul.py

@@ -6,7 +6,7 @@ from tinygrad.schedule.kernelize import merge_views, view_left
 from tinygrad.helpers import getenv, colored, prod, unwrap
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.view import strides_for_shape
-from tinygrad.opt.kernel import axis_colors
+from tinygrad.codegen.opt.kernel import axis_colors
 
 def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
 

extra/gemm/max_matmul.py

@@ -5,9 +5,9 @@ from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program
 
 # for copied uops
-from tinygrad.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.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.opt.kernel import OptOps
+from tinygrad.codegen.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.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 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.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.to_program()

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.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.codegen.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
 from tinygrad.engine.realize import get_program
 
 class MCTSNode:

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.opt.search import actions
+from tinygrad.codegen.opt.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 
 # more stuff
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.search import actions
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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
@@ -11,7 +11,7 @@ inf, nan = float('inf'), float('nan')
 UOps = Ops
 
 # kernel unpacker
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.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):
@@ -103,7 +103,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
   return ret
 
 from tinygrad.device import Device, Buffer
-from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
+from tinygrad.codegen.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.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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.opt.search import actions, bufs_from_lin, get_kernel_actions
+from tinygrad.codegen.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.opt.kernel import Kernel
-from tinygrad.opt.search import get_kernel_actions, actions
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Kernel
-from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.codegen.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.opt.search import bufs_from_lin, get_kernel_actions
+from tinygrad.codegen.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, get_program
 from tinygrad.renderer import ProgramSpec
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.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.schedule',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
+                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.codegen.opt',
                   'tinygrad.runtime.support.nv', 'tinygrad.apps'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[

test/external/external_benchmark_hcopt.py

@@ -1,7 +1,7 @@
 import random
 from tinygrad.helpers import getenv
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.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_sdxl_softmax.py

@@ -8,7 +8,7 @@ if __name__ == "__main__":
   GlobalCounters.reset()
   t.softmax(-1, dtype="half", _single_kernel=True).realize()
 
-  from tinygrad.opt.kernel import Kernel, Opt, OptOps
+  from tinygrad.codegen.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.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.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.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.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.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 
 class TestHIPCompileSpeed(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")

test/external/external_test_nv.py

@@ -2,11 +2,11 @@ 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.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 
 from test.external.fuzz_linearizer import run_linearizer
 

test/external/fuzz_linearizer.py

@@ -21,9 +21,9 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
 
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-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.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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/speed_beam_v_hcopt.py

@@ -1,7 +1,7 @@
 from tinygrad import Device
 from tinygrad.helpers import getenv, DEBUG, BEAM
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.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.opt.heuristic import hand_coded_optimizations
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner, get_program
 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.opt.heuristic import hand_coded_optimizations
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner, get_program
 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, get_program
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.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.opt.search import bufs_from_lin
+from tinygrad.codegen.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.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps, Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel, KernelOptError
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
-from tinygrad.opt.search import get_kernel_actions
+from tinygrad.codegen.opt.search import get_kernel_actions
 from tinygrad.uop.ops import Ops
 from tinygrad.codegen import apply_rewrites, rewrites_for_views
 

test/test_hcq.py

@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad import Variable
 
 MOCKGPU = getenv("MOCKGPU")

test/test_linearizer.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from dataclasses import replace
 
-from tinygrad.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
+from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError, Kernel, AxisType
 from tinygrad.codegen.gpudims import get_grouped_dims
 from tinygrad.uop.ops import UOp, Ops, GroupOp, KernelInfo
 from tinygrad.device import Device, Buffer, is_dtype_supported
@@ -10,7 +10,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, get_program
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM
 from tinygrad.dtype import DType, dtypes, AddrSpace
 from tinygrad.codegen import apply_rewrites, rewrites_for_views
@@ -410,7 +410,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.opt.search import get_kernel_actions
+      from tinygrad.codegen.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.opt.search import Opt, OptOps
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.engine.realize import get_program
 
 class TestLinearizerDumb(unittest.TestCase):

test/test_linearizer_overflows.py

@@ -1,8 +1,8 @@
 # ruff: noqa: E501
 import unittest
 from tinygrad import dtypes
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import Opt, OptOps, bufs_from_lin
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 
 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.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item, get_program
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps, Kernel
+from tinygrad.codegen.opt.kernel import Opt, OptOps, Kernel
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import bufs_from_lin, actions, beam_search
+from tinygrad.codegen.opt.search import bufs_from_lin, actions, beam_search
 from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
@@ -36,7 +36,7 @@ class TestBEAM(unittest.TestCase):
 
   def test_get_kernel_actions_dedup(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.opt.search import get_kernel_actions
+    from tinygrad.codegen.opt.search import get_kernel_actions
     a = Tensor.empty(4, 3)
     b = Tensor.empty(3)
     realized_ast, _ = helper_realized_ast(a @ b)
@@ -61,7 +61,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.opt.search import get_kernel_actions
+    from tinygrad.codegen.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]
@@ -78,7 +78,7 @@ class TestBEAM(unittest.TestCase):
 
   def test_get_kernel_actions_preserves_actions_state(self):
     from test.test_linearizer import helper_realized_ast
-    from tinygrad.opt.search import get_kernel_actions
+    from tinygrad.codegen.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

@@ -14,7 +14,7 @@ from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen import full_rewrite
 from tinygrad.uop.symbolic import sym
 from tinygrad.device import is_dtype_supported
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import 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.opt.kernel import Opt, OptOps, KernelOptError
+from tinygrad.codegen.opt.kernel import Opt, OptOps, KernelOptError
 from tinygrad.device import Device
 
 def flops_mem(uops, ignore_indexing=False):

test/unit/test_linearizer_rewrite.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Context, Device
 from tinygrad.engine.realize import get_program
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import KernelInfo
 
 class TestLinearizerRewrite(unittest.TestCase):

test/unit/test_search.py

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

tinygrad/codegen/__init__.py

@@ -16,8 +16,8 @@ from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexin
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.optional import get_late_rewrite_patterns
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.opt import pm_optimize
-from tinygrad.opt.swizzler import view_left, view_right, fix_kernel_ops
+from tinygrad.codegen.opt import pm_optimize
+from tinygrad.codegen.opt.swizzler import view_left, view_right, fix_kernel_ops
 
 @dataclass
 class RewriteStep:

tinygrad/codegen/opt/__init__.py

@@ -1,7 +1,7 @@
 # opt opinionatedly transforms an ast into an optimized ast using either heuristics or beam search
 
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.opt.kernel import Kernel
+from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops
 from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
 from tinygrad.renderer import Renderer
@@ -24,7 +24,7 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
   elif not NOOPT:
     if not k.apply_tensor_cores(USE_TC.value): k.apply_opts(hand_coded_optimizations(k))
     if BEAM >= 1:
-      from tinygrad.opt.search import beam_search, bufs_from_lin
+      from tinygrad.codegen.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)))

tinygrad/codegen/opt/heuristic.py

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

tinygrad/codegen/opt/kernel.py

@@ -8,13 +8,13 @@ from enum import Enum, auto
 from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, AxisType
 from tinygrad.uop.spec import type_verify, ast_spec
 from tinygrad.device import Device
-from tinygrad.opt.tc import TensorCore
+from tinygrad.codegen.opt.tc import TensorCore
 from tinygrad.renderer import Renderer
 from tinygrad.dtype import ImageDType, AddrSpace
 from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG, TC_SELECT, TC_OPT, AMX
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import strides_for_shape, get_contraction
-from tinygrad.opt.swizzler import view_left, view_left_through_load
+from tinygrad.codegen.opt.swizzler import view_left, view_left_through_load
 
 class OptOps(Enum):
   TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702

tinygrad/codegen/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.opt.kernel import Kernel, Opt, OptOps, KernelOptError
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps, KernelOptError
 from tinygrad.tensor import Tensor
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.renderer import ProgramSpec

tinygrad/engine/realize.py

@@ -8,7 +8,7 @@ from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
 from tinygrad.engine.schedule import ScheduleItem
 from tinygrad.codegen import full_rewrite
-from tinygrad.opt.kernel import Opt
+from tinygrad.codegen.opt.kernel import Opt
 
 # **************** Program Creation ****************
 
@@ -77,7 +77,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.opt.search import optimize_local_size
+      from tinygrad.codegen.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/helpers.py

@@ -218,7 +218,7 @@ def cpu_profile(name:str|TracingKey, device="CPU", is_copy=False, display=True)
 cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Library/Caches" if OSX else "~/.cache")), "tinygrad")
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))
 
-VERSION = 21
+VERSION = 22
 _db_connection = None
 def db_connection():
   global _db_connection

tinygrad/renderer/__init__.py

@@ -6,8 +6,8 @@ from tinygrad.helpers import to_function_name, dedup, prod
 from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher
 from tinygrad.dtype import AddrSpace, PtrDType
 if TYPE_CHECKING:
-  from tinygrad.opt.tc import TensorCore
-  from tinygrad.opt.kernel import Opt
+  from tinygrad.codegen.opt.tc import TensorCore
+  from tinygrad.codegen.opt.kernel import Opt
 
 @dataclass(frozen=True)
 class Estimates:

tinygrad/renderer/cstyle.py

@@ -1,7 +1,7 @@
 from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
-from tinygrad.opt import tc
+from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate

tinygrad/renderer/llvmir.py

@@ -1,6 +1,6 @@
 from typing import cast
 import math, struct, sys
-from tinygrad.opt import tc
+from tinygrad.codegen.opt import tc
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp

tinygrad/renderer/ptx.py

@@ -1,7 +1,7 @@
 from typing import cast, Callable
 import struct
 from collections import defaultdict
-from tinygrad.opt import tc
+from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
 from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace
 from tinygrad.renderer import Renderer

tinygrad/runtime/ops_python.py

@@ -7,7 +7,7 @@ import pickle, base64, itertools, time, struct, sys
 from tinygrad.dtype import DType, dtypes, ImageDType, PtrDType, truncate
 from tinygrad.helpers import all_same, getenv, flatten, get_single_element
 from tinygrad.device import Compiled, Compiler, Allocator
-from tinygrad.opt import tc
+from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
 

tinygrad/schedule/kernelize.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import Metadata, all_int, all_same, prod, dedup, unwrap, g
 from tinygrad.dtype import ImageDType
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.grouper import group_realizes, ALWAYS_CONTIGUOUS
-from tinygrad.opt.swizzler import merge_views, apply_swizzle, swizzle_reduceop
+from tinygrad.codegen.opt.swizzler import merge_views, apply_swizzle, swizzle_reduceop
 
 # creation can recurse a lot
 import sys