diff --git a/docs/abstractions2.py b/docs/abstractions2.py index f3f3723ff4..e7dfc6313b 100644 --- a/docs/abstractions2.py +++ b/docs/abstractions2.py @@ -78,7 +78,7 @@ print("******** third, the UOp ***********") from tinygrad.engine.realize import run_schedule from tinygrad.engine.schedule import create_schedule_with_vars -from tinygrad.engine.kernelize import get_kernelize_map +from tinygrad.kernelize.kernelize import get_kernelize_map # allocate some values + load in values a = UOp.new_buffer(DEVICE, 1, dtypes.int32) diff --git a/examples/openpilot/compile4.py b/examples/openpilot/compile4.py index 7beb8ad029..be537f3cdf 100644 --- a/examples/openpilot/compile4.py +++ b/examples/openpilot/compile4.py @@ -2,7 +2,7 @@ import sys, onnx from tinygrad import Tensor, fetch, GlobalCounters from tinygrad.uop.ops import UOp from tinygrad.frontend.onnx import OnnxRunner -from tinygrad.engine.kernelize import get_kernelize_map +from tinygrad.kernelize.kernelize import get_kernelize_map from tinygrad.engine.schedule import create_schedule_with_vars from tinygrad.engine.realize import run_schedule diff --git a/setup.py b/setup.py index 9c13c35d05..ecff13461f 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,7 @@ setup(name='tinygrad', long_description=long_description, 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.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.kernelize', '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=[ diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index 34a89438ee..3253a5c7b4 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -3,7 +3,7 @@ import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools 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.kernelize.kernelize import get_kernelize_map from tinygrad.opt.kernel import Kernel from tinygrad.uop.ops import UOp, Ops diff --git a/test/helpers.py b/test/helpers.py index e8eeea9bf8..91667e0787 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -6,7 +6,7 @@ from tinygrad.uop.ops import UOp, Ops, sint, graph_rewrite from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.tensor import _to_np_dtype from tinygrad.engine.realize import Runner -from tinygrad.engine.kernelize import view_left +from tinygrad.kernelize.kernelize import view_left from tinygrad.dtype import ConstType, DType from tinygrad.nn.state import get_parameters from tinygrad.helpers import T, unwrap, CI diff --git a/test/test_schedule.py b/test/test_schedule.py index 01d2e4c803..de63056016 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -15,7 +15,7 @@ from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites from tinygrad.uop.symbolic import symbolic_simple from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp -from tinygrad.engine.kernelize import merge_views, get_kernelize_map, Kernel +from tinygrad.kernelize.kernelize import merge_views, get_kernelize_map, Kernel from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule diff --git a/test/test_uops.py b/test/test_uops.py index c7ccbb7713..9c794de8c7 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -10,7 +10,7 @@ from tinygrad.device import Buffer, Device from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401 from tinygrad.uop.spec import spec from tinygrad.renderer import ProgramSpec -from tinygrad.engine.kernelize import fix_kernel_ops +from tinygrad.kernelize.kernelize import fix_kernel_ops from tinygrad.engine.realize import CompiledRunner from tinygrad.codegen import full_rewrite from tinygrad.uop.symbolic import sym diff --git a/test/unit/test_rewrite_tracked_childen.py b/test/unit/test_rewrite_tracked_childen.py index 8d3d4c46e9..6cddd6f289 100644 --- a/test/unit/test_rewrite_tracked_childen.py +++ b/test/unit/test_rewrite_tracked_childen.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Tensor from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp -from tinygrad.engine.kernelize import sym, merge_views +from tinygrad.kernelize.kernelize import sym, merge_views class TestRewriteTrackedChildren(unittest.TestCase): @unittest.skip("track_children no longer supported") diff --git a/tinygrad/kernelize/__init__.py b/tinygrad/kernelize/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tinygrad/engine/grouper.py b/tinygrad/kernelize/grouper.py similarity index 100% rename from tinygrad/engine/grouper.py rename to tinygrad/kernelize/grouper.py diff --git a/tinygrad/engine/kernelize.py b/tinygrad/kernelize/kernelize.py similarity index 99% rename from tinygrad/engine/kernelize.py rename to tinygrad/kernelize/kernelize.py index 55a0430c3c..dd2141bf07 100644 --- a/tinygrad/engine/kernelize.py +++ b/tinygrad/kernelize/kernelize.py @@ -6,10 +6,10 @@ from tinygrad.codegen.lowerer import get_contraction_with_reduce from tinygrad.uop.symbolic import symbolic_simple from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize, FUSE_ARANGE, DEBUG, SPLIT_REDUCEOP from tinygrad.dtype import ImageDType -from tinygrad.engine.multi import multi_pm +from tinygrad.kernelize.multi import multi_pm from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View, strides_for_shape -from tinygrad.engine.grouper import group_realizes, ALWAYS_CONTIGUOUS +from tinygrad.kernelize.grouper import group_realizes, ALWAYS_CONTIGUOUS # creation can recurse a lot import sys diff --git a/tinygrad/engine/multi.py b/tinygrad/kernelize/multi.py similarity index 100% rename from tinygrad/engine/multi.py rename to tinygrad/kernelize/multi.py diff --git a/tinygrad/opt/kernel.py b/tinygrad/opt/kernel.py index b475647731..8380f984c5 100644 --- a/tinygrad/opt/kernel.py +++ b/tinygrad/opt/kernel.py @@ -15,7 +15,7 @@ from tinygrad.helpers import DEBUG, TC_SELECT, TC_OPT, AMX from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import strides_for_shape from tinygrad.codegen.lowerer import get_contraction -from tinygrad.engine.kernelize import view_left +from tinygrad.kernelize.kernelize import view_left from tinygrad.codegen import full_rewrite class KernelOptError(Exception): pass diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 65b039224e..1b77a9fae8 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -14,7 +14,7 @@ from tinygrad.device import Device, Buffer from tinygrad.engine.realize import run_schedule from tinygrad.engine.memory import memory_planner from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars -from tinygrad.engine.kernelize import get_kernelize_map +from tinygrad.kernelize.kernelize import get_kernelize_map # *** all in scope Tensors are here. this gets relevant UOps ***