From 92678e59ee2d1d6e59c45ff7a7cb115d3ae8d877 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 20 Jun 2025 15:22:28 -0700 Subject: [PATCH] move kernel to opt (#10899) --- examples/handcode_opt.py | 6 +++--- extra/assembly/assembly.py | 2 +- extra/assembly/assembly_arm64.py | 2 +- extra/assembly/assembly_ptx.py | 2 +- extra/assembly/assembly_rdna.py | 2 +- extra/backends/triton.py | 2 +- extra/gemm/amd_matmul.py | 2 +- extra/gemm/max_matmul.py | 4 ++-- extra/gemm/simple_matmul.py | 2 +- extra/gemm/tinygrad_nv_matmul.py | 2 +- extra/gemm/tvm_gemm.py | 2 +- extra/mcts_search.py | 4 ++-- extra/optimization/extract_policynet.py | 6 +++--- extra/optimization/extract_sa_pairs.py | 6 +++--- extra/optimization/get_action_space.py | 6 +++--- extra/optimization/helpers.py | 6 +++--- extra/optimization/pretrain_valuenet.py | 4 ++-- extra/optimization/rl.py | 2 +- extra/optimization/run_qnet.py | 4 ++-- extra/optimization/search.py | 4 ++-- extra/optimization/test_net.py | 4 ++-- extra/optimization/test_time_linearizer.py | 2 +- extra/replay_pkl.py | 4 ++-- setup.py | 2 +- test/external/external_benchmark_hcopt.py | 4 ++-- test/external/external_benchmark_schedule.py | 6 +++--- test/external/external_benchmark_sdxl_softmax.py | 2 +- test/external/external_debug_metal_sd_conv.py | 4 ++-- test/external/external_test_hcq_fuzz_failures.py | 4 ++-- test/external/external_test_hip_compile.py | 2 +- test/external/external_test_nv.py | 4 ++-- test/external/external_test_train_gpt2.py | 4 ++-- test/external/external_test_valid_remove.py | 4 ++-- test/external/fuzz_linearizer.py | 6 +++--- test/external/process_replay/process_replay.py | 2 +- test/external/speed_beam_v_hcopt.py | 4 ++-- test/external/speed_compare_amd_am.py | 4 ++-- test/external/speed_compare_cuda_nv.py | 4 ++-- test/external/speed_compare_cuda_ptx.py | 4 ++-- test/external/verify_kernel.py | 2 +- test/test_arange.py | 4 ++-- test/test_hcq.py | 2 +- test/test_linearizer.py | 6 +++--- test/test_linearizer_dumb.py | 4 ++-- test/test_linearizer_failures.py | 4 ++-- test/test_linearizer_overflows.py | 4 ++-- test/test_opt_gemm.py | 2 +- test/test_quantize_onnx.py | 4 ++-- test/test_search.py | 10 +++++----- test/test_uops.py | 2 +- test/test_uops_stats.py | 2 +- test/test_winograd.py | 4 ++-- test/unit/test_linearizer_rewrite.py | 2 +- test/unit/test_search.py | 4 ++-- test/unit/test_uop_spec.py | 2 +- tinygrad/engine/realize.py | 8 ++++---- tinygrad/{codegen => opt}/heuristic.py | 2 +- tinygrad/{codegen => opt}/kernel.py | 0 tinygrad/{engine => opt}/search.py | 2 +- tinygrad/viz/serve.py | 2 +- 60 files changed, 106 insertions(+), 106 deletions(-) rename tinygrad/{codegen => opt}/heuristic.py (99%) rename tinygrad/{codegen => opt}/kernel.py (100%) rename tinygrad/{engine => opt}/search.py (99%) diff --git a/examples/handcode_opt.py b/examples/handcode_opt.py index cdb7e32861..f48d8a3812 100644 --- a/examples/handcode_opt.py +++ b/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 diff --git a/extra/assembly/assembly.py b/extra/assembly/assembly.py index 1362d6deba..688f2b9aac 100644 --- a/extra/assembly/assembly.py +++ b/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 diff --git a/extra/assembly/assembly_arm64.py b/extra/assembly/assembly_arm64.py index 800f754705..1b61b8e5af 100644 --- a/extra/assembly/assembly_arm64.py +++ b/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 diff --git a/extra/assembly/assembly_ptx.py b/extra/assembly/assembly_ptx.py index f3fce81db2..5535d4ccf5 100644 --- a/extra/assembly/assembly_ptx.py +++ b/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 diff --git a/extra/assembly/assembly_rdna.py b/extra/assembly/assembly_rdna.py index 6eae18b4b3..d57ba52456 100644 --- a/extra/assembly/assembly_rdna.py +++ b/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 diff --git a/extra/backends/triton.py b/extra/backends/triton.py index c1e81e4f0f..f653341623 100644 --- a/extra/backends/triton.py +++ b/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 diff --git a/extra/gemm/amd_matmul.py b/extra/gemm/amd_matmul.py index c1d16cff23..3c4cfbd25c 100644 --- a/extra/gemm/amd_matmul.py +++ b/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 diff --git a/extra/gemm/max_matmul.py b/extra/gemm/max_matmul.py index 975f4b1db7..9d093c2fbe 100644 --- a/extra/gemm/max_matmul.py +++ b/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 diff --git a/extra/gemm/simple_matmul.py b/extra/gemm/simple_matmul.py index 29f26d7e80..8b857a8826 100644 --- a/extra/gemm/simple_matmul.py +++ b/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 diff --git a/extra/gemm/tinygrad_nv_matmul.py b/extra/gemm/tinygrad_nv_matmul.py index ab8f078ade..f81877a41a 100644 --- a/extra/gemm/tinygrad_nv_matmul.py +++ b/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 diff --git a/extra/gemm/tvm_gemm.py b/extra/gemm/tvm_gemm.py index fa6f661557..47796023f8 100644 --- a/extra/gemm/tvm_gemm.py +++ b/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() diff --git a/extra/mcts_search.py b/extra/mcts_search.py index 9090c902ff..562d11bc18 100644 --- a/extra/mcts_search.py +++ b/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): diff --git a/extra/optimization/extract_policynet.py b/extra/optimization/extract_policynet.py index da8aa0d42c..663107f3f7 100644 --- a/extra/optimization/extract_policynet.py +++ b/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: diff --git a/extra/optimization/extract_sa_pairs.py b/extra/optimization/extract_sa_pairs.py index 224d9c6302..935860b168 100644 --- a/extra/optimization/extract_sa_pairs.py +++ b/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 diff --git a/extra/optimization/get_action_space.py b/extra/optimization/get_action_space.py index f443d7c6d9..d27c568230 100644 --- a/extra/optimization/get_action_space.py +++ b/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() diff --git a/extra/optimization/helpers.py b/extra/optimization/helpers.py index 73006092e1..bca2eb6b85 100644 --- a/extra/optimization/helpers.py +++ b/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 diff --git a/extra/optimization/pretrain_valuenet.py b/extra/optimization/pretrain_valuenet.py index 312aec6fc6..2dd0f90c6d 100644 --- a/extra/optimization/pretrain_valuenet.py +++ b/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 diff --git a/extra/optimization/rl.py b/extra/optimization/rl.py index f05070957c..d45e84b14b 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.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 diff --git a/extra/optimization/run_qnet.py b/extra/optimization/run_qnet.py index 65acae193c..6dfca1b379 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.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]]: diff --git a/extra/optimization/search.py b/extra/optimization/search.py index 469487da58..6b49f3622c 100644 --- a/extra/optimization/search.py +++ b/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__': diff --git a/extra/optimization/test_net.py b/extra/optimization/test_net.py index 68a023f0ee..efd2284753 100644 --- a/extra/optimization/test_net.py +++ b/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 diff --git a/extra/optimization/test_time_linearizer.py b/extra/optimization/test_time_linearizer.py index 4bfb1f0f03..32d27b2f74 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, 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() diff --git a/extra/replay_pkl.py b/extra/replay_pkl.py index 9f02492d6a..62f991132a 100644 --- a/extra/replay_pkl.py +++ b/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"): diff --git a/setup.py b/setup.py index 9765ed31e3..9c13c35d05 100644 --- a/setup.py +++ b/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", diff --git a/test/external/external_benchmark_hcopt.py b/test/external/external_benchmark_hcopt.py index c2ed276e2e..c10c0e89cf 100644 --- a/test/external/external_benchmark_hcopt.py +++ b/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): diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py index 67d7846ca8..2a47061e69 100644 --- a/test/external/external_benchmark_schedule.py +++ b/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__": diff --git a/test/external/external_benchmark_sdxl_softmax.py b/test/external/external_benchmark_sdxl_softmax.py index 35cf06b0b0..639d59e415 100644 --- a/test/external/external_benchmark_sdxl_softmax.py +++ b/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()) diff --git a/test/external/external_debug_metal_sd_conv.py b/test/external/external_debug_metal_sd_conv.py index cbdd93f88b..6c2d24b5a3 100644 --- a/test/external/external_debug_metal_sd_conv.py +++ b/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 diff --git a/test/external/external_test_hcq_fuzz_failures.py b/test/external/external_test_hcq_fuzz_failures.py index 09b3633d56..2ac55c93c7 100644 --- a/test/external/external_test_hcq_fuzz_failures.py +++ b/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 diff --git a/test/external/external_test_hip_compile.py b/test/external/external_test_hip_compile.py index 0edef8730b..53ddda0b03 100644 --- a/test/external/external_test_hip_compile.py +++ b/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") diff --git a/test/external/external_test_nv.py b/test/external/external_test_nv.py index f061975e44..1e8aaa377d 100644 --- a/test/external/external_test_nv.py +++ b/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 diff --git a/test/external/external_test_train_gpt2.py b/test/external/external_test_train_gpt2.py index 6d3a1cb9c1..2cac47941c 100644 --- a/test/external/external_test_train_gpt2.py +++ b/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 diff --git a/test/external/external_test_valid_remove.py b/test/external/external_test_valid_remove.py index 65f8112e96..9734b1e39b 100644 --- a/test/external/external_test_valid_remove.py +++ b/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): diff --git a/test/external/fuzz_linearizer.py b/test/external/fuzz_linearizer.py index 1bf269ce56..b9711b8d88 100644 --- a/test/external/fuzz_linearizer.py +++ b/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 diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index d072d8b490..34a89438ee 100755 --- a/test/external/process_replay/process_replay.py +++ b/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 diff --git a/test/external/speed_beam_v_hcopt.py b/test/external/speed_beam_v_hcopt.py index 1bd107f6d7..e73c6466da 100644 --- a/test/external/speed_beam_v_hcopt.py +++ b/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__": diff --git a/test/external/speed_compare_amd_am.py b/test/external/speed_compare_amd_am.py index 760833d48c..3a422d4ca0 100644 --- a/test/external/speed_compare_amd_am.py +++ b/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 diff --git a/test/external/speed_compare_cuda_nv.py b/test/external/speed_compare_cuda_nv.py index 3ce1a74a31..f9185a94bf 100644 --- a/test/external/speed_compare_cuda_nv.py +++ b/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 diff --git a/test/external/speed_compare_cuda_ptx.py b/test/external/speed_compare_cuda_ptx.py index 32b586a1d5..068b7619ba 100644 --- a/test/external/speed_compare_cuda_ptx.py +++ b/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__": diff --git a/test/external/verify_kernel.py b/test/external/verify_kernel.py index 55ebb6849b..5a251ce8a8 100644 --- a/test/external/verify_kernel.py +++ b/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 diff --git a/test/test_arange.py b/test/test_arange.py index e9bc9983b2..308c741e52 100644 --- a/test/test_arange.py +++ b/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): diff --git a/test/test_hcq.py b/test/test_hcq.py index 884c5fc21a..90a90a7b43 100644 --- a/test/test_hcq.py +++ b/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") diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 6a9dfe2049..9f41bcdfae 100644 --- a/test/test_linearizer.py +++ b/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]) diff --git a/test/test_linearizer_dumb.py b/test/test_linearizer_dumb.py index 0c56dc9e0e..669d0a5d02 100644 --- a/test/test_linearizer_dumb.py +++ b/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") diff --git a/test/test_linearizer_failures.py b/test/test_linearizer_failures.py index ec4f2ae665..9a682a73f6 100644 --- a/test/test_linearizer_failures.py +++ b/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 diff --git a/test/test_linearizer_overflows.py b/test/test_linearizer_overflows.py index 0bd8ef0785..9192d6145a 100644 --- a/test/test_linearizer_overflows.py +++ b/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 diff --git a/test/test_opt_gemm.py b/test/test_opt_gemm.py index 862d977895..f38caf2f33 100644 --- a/test/test_opt_gemm.py +++ b/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): diff --git a/test/test_quantize_onnx.py b/test/test_quantize_onnx.py index 7626aad8d2..7b1dd2747e 100644 --- a/test/test_quantize_onnx.py +++ b/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 diff --git a/test/test_search.py b/test/test_search.py index a3f459e06b..af2c4f0a39 100644 --- a/test/test_search.py +++ b/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) diff --git a/test/test_uops.py b/test/test_uops.py index d6c7d99bb6..c7ccbb7713 100644 --- a/test/test_uops.py +++ 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) diff --git a/test/test_uops_stats.py b/test/test_uops_stats.py index fdaf21c4f7..b534c0b0e1 100644 --- a/test/test_uops_stats.py +++ b/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): diff --git a/test/test_winograd.py b/test/test_winograd.py index e4ef3caa41..fa34146a8d 100644 --- a/test/test_winograd.py +++ b/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): diff --git a/test/unit/test_linearizer_rewrite.py b/test/unit/test_linearizer_rewrite.py index 358c43cec6..6db29cfd3f 100644 --- a/test/unit/test_linearizer_rewrite.py +++ b/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): diff --git a/test/unit/test_search.py b/test/unit/test_search.py index be16686cde..cb9eb3d3c6 100644 --- a/test/unit/test_search.py +++ b/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 diff --git a/test/unit/test_uop_spec.py b/test/unit/test_uop_spec.py index dfc9bca106..f5227e5894 100644 --- a/test/unit/test_uop_spec.py +++ b/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 diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index db0b0fc7fe..3c9298fdc2 100644 --- a/tinygrad/engine/realize.py +++ b/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) diff --git a/tinygrad/codegen/heuristic.py b/tinygrad/opt/heuristic.py similarity index 99% rename from tinygrad/codegen/heuristic.py rename to tinygrad/opt/heuristic.py index 1f77339003..b12f13135b 100644 --- a/tinygrad/codegen/heuristic.py +++ b/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 diff --git a/tinygrad/codegen/kernel.py b/tinygrad/opt/kernel.py similarity index 100% rename from tinygrad/codegen/kernel.py rename to tinygrad/opt/kernel.py diff --git a/tinygrad/engine/search.py b/tinygrad/opt/search.py similarity index 99% rename from tinygrad/engine/search.py rename to tinygrad/opt/search.py index b9b78910d8..27c9e5e1eb 100644 --- a/tinygrad/engine/search.py +++ b/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 diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index e2ebbe086f..1ff7faf133 100755 --- a/tinygrad/viz/serve.py +++ b/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