diff --git a/docs/abstractions2.py b/docs/abstractions2.py index 668bc9f21c..49222e3b3b 100644 --- a/docs/abstractions2.py +++ b/docs/abstractions2.py @@ -39,7 +39,7 @@ DEVICE = "CPU" # NOTE: you can change this! import struct from tinygrad.dtype import dtypes from tinygrad.device import Buffer, Device -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.shape.shapetracker import ShapeTracker # allocate some buffers + load in values diff --git a/docs/developer/uop.md b/docs/developer/uop.md index 063c475bc5..5562df2fbe 100644 --- a/docs/developer/uop.md +++ b/docs/developer/uop.md @@ -1,10 +1,10 @@ -::: tinygrad.ops.UOp +::: tinygrad.uop.ops.UOp options: members: false members_order: source show_labels: false -::: tinygrad.ops.Ops +::: tinygrad.uop.ops.Ops options: members: true members_order: source diff --git a/examples/gpt2.py b/examples/gpt2.py index d24afcc590..c3d933b2a9 100644 --- a/examples/gpt2.py +++ b/examples/gpt2.py @@ -3,7 +3,7 @@ import os, argparse, contextlib from typing import Optional, Union with contextlib.suppress(ImportError): import tiktoken from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange from tinygrad.nn import Embedding, Linear, LayerNorm from tinygrad.nn.state import gguf_load, torch_load, load_state_dict, get_state_dict diff --git a/examples/handcode_opt.py b/examples/handcode_opt.py index f6d834a92b..cdb7e32861 100644 --- a/examples/handcode_opt.py +++ b/examples/handcode_opt.py @@ -4,7 +4,7 @@ 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.ops import Ops, sym_infer +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.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA diff --git a/examples/llm.c/export.py b/examples/llm.c/export.py index b4b8a3ed38..bc13a09fbc 100755 --- a/examples/llm.c/export.py +++ b/examples/llm.c/export.py @@ -7,7 +7,7 @@ from train_gpt2 import GPT, GPTConfig from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name from tinygrad.engine.realize import get_kernel, run_schedule from tinygrad.engine.memory import memory_planner -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops TIMING = getenv("TIMING") diff --git a/extra/assembly/assembly.py b/extra/assembly/assembly.py index 666282c05a..1362d6deba 100644 --- a/extra/assembly/assembly.py +++ b/extra/assembly/assembly.py @@ -1,9 +1,9 @@ from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast from tinygrad.codegen.kernel import Ops, MemOp, UOp -from tinygrad.ops import BinaryOps, UnaryOps +from tinygrad.uop.ops import BinaryOps, UnaryOps from tinygrad.dtype import DType, dtypes from tinygrad.helpers import DEBUG -from tinygrad.ops import Variable, NumNode, MulNode, DivNode, ModNode, LtNode, SumNode, AndNode +from tinygrad.uop.ops import Variable, NumNode, MulNode, DivNode, ModNode, LtNode, SumNode, AndNode import functools import math from collections import defaultdict diff --git a/extra/assembly/assembly_arm64.py b/extra/assembly/assembly_arm64.py index bca479ab41..800f754705 100644 --- a/extra/assembly/assembly_arm64.py +++ b/extra/assembly/assembly_arm64.py @@ -2,7 +2,7 @@ import struct from platform import system from typing import Tuple, Dict, List, Optional from tinygrad import dtypes -from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps +from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps from tinygrad.codegen.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 7ac1ff5616..f3fce81db2 100644 --- a/extra/assembly/assembly_ptx.py +++ b/extra/assembly/assembly_ptx.py @@ -3,7 +3,7 @@ import struct from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage from tinygrad.codegen.kernel import Ops, UOp from tinygrad import dtypes -from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps +from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps from tinygrad.runtime.ops_cuda import arch dtype_to_nvtype = {dtypes.float32: "f32", dtypes.float16: "f16", dtypes.int64: "s64", dtypes.int32: "s32", dtypes.int8: "s8", dtypes.bool: "pred", dtypes.uint64: "u64", dtypes.uint32: "u32", dtypes.uint16: "u16", dtypes.uint8: "u8", "bits16": "b16", dtypes.float64: "f64"} diff --git a/extra/assembly/assembly_rdna.py b/extra/assembly/assembly_rdna.py index cb41ba8f22..6eae18b4b3 100644 --- a/extra/assembly/assembly_rdna.py +++ b/extra/assembly/assembly_rdna.py @@ -3,7 +3,7 @@ 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.ops import BinaryOps, UnaryOps, TernaryOps +from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH # ugh, is this really needed? diff --git a/extra/backends/clang_graph.py b/extra/backends/clang_graph.py index 5e558a3ef7..9c0cf3b3e0 100644 --- a/extra/backends/clang_graph.py +++ b/extra/backends/clang_graph.py @@ -4,7 +4,7 @@ from tinygrad.helpers import dedup, cpu_time_execution, DEBUG from tinygrad.engine.jit import GraphRunner, GraphException from tinygrad.device import Buffer, Device from tinygrad.engine.realize import ExecItem, CompiledRunner -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.runtime.ops_cpu import ClangProgram from tinygrad.renderer.cstyle import ClangRenderer render_dtype = ClangRenderer().render_dtype diff --git a/extra/backends/hsa_graph.py b/extra/backends/hsa_graph.py index 6541f3cdc8..aced2c028a 100644 --- a/extra/backends/hsa_graph.py +++ b/extra/backends/hsa_graph.py @@ -3,7 +3,7 @@ from typing import List, Any, Dict, cast, Optional, Tuple from tinygrad.helpers import init_c_var, round_up from tinygrad.device import Buffer, BufferSpec from tinygrad.device import Compiled, Device -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.runtime.ops_hsa import HSADevice, PROFILE, Profiler from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner from tinygrad.engine.jit import MultiGraphRunner, GraphException diff --git a/extra/backends/rdna.py b/extra/backends/rdna.py index 685a6b74e9..1e26e97783 100644 --- a/extra/backends/rdna.py +++ b/extra/backends/rdna.py @@ -1,7 +1,7 @@ from typing import Dict, Set import yaml from tinygrad.codegen.uops import UOpGraph, UOps, UOp -from tinygrad.ops import BinaryOps +from tinygrad.uop.ops import BinaryOps from tinygrad.dtype import dtypes def uops_to_rdna(function_name:str, uops:UOpGraph) -> str: diff --git a/extra/backends/triton.py b/extra/backends/triton.py index 9da79b108c..c1e81e4f0f 100644 --- a/extra/backends/triton.py +++ b/extra/backends/triton.py @@ -1,6 +1,6 @@ from typing import Dict, List, Final, Callable, DefaultDict from collections import defaultdict -from tinygrad.ops import UnaryOps, BinaryOps, TernaryOps, Op +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 triton.compiler import compile as triton_compile diff --git a/extra/export_model.py b/extra/export_model.py index 17cd7628d6..63b4098270 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -6,7 +6,7 @@ from tinygrad.engine.jit import TinyJit from tinygrad.nn.state import get_state_dict from tinygrad.helpers import Context from tinygrad.dtype import dtypes -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops import json from collections import OrderedDict diff --git a/extra/gemm/amd_matmul.py b/extra/gemm/amd_matmul.py index 613a4cfe81..c1d16cff23 100644 --- a/extra/gemm/amd_matmul.py +++ b/extra/gemm/amd_matmul.py @@ -6,7 +6,7 @@ from tinygrad import Tensor, Device, Context from tinygrad.helpers import getenv from tinygrad.codegen.kernel import Kernel, Opt, OptOps from tinygrad.engine.realize import CompiledRunner, ExecItem -from tinygrad.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp +from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp # TODO: on METAL for `DEBUG=4 python3 extra/gemm/amd_matmul.py` # * fix load grouping (like float4). idk why it's not working, need new devectorizer (this is a Monday project) diff --git a/extra/gemm/max_matmul.py b/extra/gemm/max_matmul.py index 5cf78fc0ef..975f4b1db7 100644 --- a/extra/gemm/max_matmul.py +++ b/extra/gemm/max_matmul.py @@ -5,7 +5,7 @@ from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self # for copied uops from tinygrad.codegen.kernel import Kernel, KernelOptError -from tinygrad.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo +from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo from tinygrad.engine.search import Opt, OptOps from tinygrad import Device, dtypes, Tensor from tinygrad.dtype import PtrDType, DType, DTYPES_DICT diff --git a/extra/introspection.py b/extra/introspection.py index 106c7c51fd..0686514725 100644 --- a/extra/introspection.py +++ b/extra/introspection.py @@ -1,6 +1,6 @@ import gc from tinygrad.helpers import prod -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp from tinygrad.device import Buffer from tinygrad import Tensor, GlobalCounters diff --git a/extra/optimization/extract_policynet.py b/extra/optimization/extract_policynet.py index 129208d4f7..cf231bf743 100644 --- a/extra/optimization/extract_policynet.py +++ b/extra/optimization/extract_policynet.py @@ -11,11 +11,11 @@ from tinygrad.codegen.kernel import Kernel from tinygrad.helpers import getenv # stuff needed to unpack a kernel -from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer +from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer from tinygrad.dtype import dtypes from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable inf, nan = float('inf'), float('nan') from tinygrad.codegen.kernel import Opt, OptOps diff --git a/extra/optimization/extract_sa_pairs.py b/extra/optimization/extract_sa_pairs.py index db4ac7e444..cae7019d53 100644 --- a/extra/optimization/extract_sa_pairs.py +++ b/extra/optimization/extract_sa_pairs.py @@ -4,11 +4,11 @@ from tqdm import tqdm, trange import numpy as np # stuff needed to unpack a kernel -from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer +from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer from tinygrad.dtype import dtypes from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable inf, nan = float('inf'), float('nan') from tinygrad.codegen.kernel import Opt, OptOps diff --git a/extra/optimization/helpers.py b/extra/optimization/helpers.py index e8e7c41c39..73006092e1 100644 --- a/extra/optimization/helpers.py +++ b/extra/optimization/helpers.py @@ -1,7 +1,7 @@ # stuff needed to unpack a kernel from tinygrad import Variable from tinygrad.codegen.kernel import Opt, OptOps -from tinygrad.ops import UOp, Ops, KernelInfo +from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.dtype import dtypes, PtrDType from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View diff --git a/extra/optimization/pretrain_valuenet.py b/extra/optimization/pretrain_valuenet.py index 2216b24d56..5fee43f956 100644 --- a/extra/optimization/pretrain_valuenet.py +++ b/extra/optimization/pretrain_valuenet.py @@ -8,11 +8,11 @@ from tinygrad.nn.optim import Adam from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict # stuff needed to unpack a kernel -from tinygrad.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer +from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer from tinygrad.dtype import dtypes from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable inf, nan = float('inf'), float('nan') from tinygrad.codegen.kernel import Opt, OptOps diff --git a/extra/to_movement_ops.py b/extra/to_movement_ops.py index fd306156f3..81afd306d8 100644 --- a/extra/to_movement_ops.py +++ b/extra/to_movement_ops.py @@ -4,9 +4,9 @@ from collections import defaultdict from typing import List, Tuple, DefaultDict from extra.optimization.helpers import load_worlds, ast_str_to_ast from tinygrad.helpers import prod, tqdm -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.ops import sym_infer +from tinygrad.uop.ops import sym_infer from tinygrad.tensor import Tensor class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto(); AS_STRIDED = auto() # noqa: E702 diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index 3d10577bdd..4a14a22ab7 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -3,7 +3,7 @@ # A002 Function argument `input` is shadowing a Python builtin # A006 Lambda argument `input` is shadowing a Python builtin from tinygrad import Tensor, dtypes, Device -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops from tinygrad.helpers import getenv, prod import torch.lib TORCH_DEBUG = getenv("TORCH_DEBUG") diff --git a/setup.py b/setup.py index 2a503fba71..a5fe80b3bd 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.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop'], package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'perfetto.html', 'assets/**/*', 'lib/**/*']}, classifiers=[ "Programming Language :: Python :: 3", diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py index 4219d696ec..6b606eeb23 100644 --- a/test/external/external_benchmark_schedule.py +++ b/test/external/external_benchmark_schedule.py @@ -2,7 +2,7 @@ from typing import List 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.ops import Ops +from tinygrad.uop.ops import Ops from tinygrad.codegen.kernel import Kernel from tinygrad.codegen.heuristic import hand_coded_optimizations from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites diff --git a/test/external/external_debug_metal_sd_conv.py b/test/external/external_debug_metal_sd_conv.py index 2577db097c..cbdd93f88b 100644 --- a/test/external/external_debug_metal_sd_conv.py +++ b/test/external/external_debug_metal_sd_conv.py @@ -3,7 +3,7 @@ from tinygrad.codegen.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.ops import UOp, Ops +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 dd3c1e5a46..09b3633d56 100644 --- a/test/external/external_test_hcq_fuzz_failures.py +++ b/test/external/external_test_hcq_fuzz_failures.py @@ -6,7 +6,7 @@ import unittest, random import numpy as np from tinygrad.codegen.kernel import Kernel, KernelOptError from tinygrad.device import is_dtype_supported -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.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_nv.py b/test/external/external_test_nv.py index 414ecece9b..b488e86725 100644 --- a/test/external/external_test_nv.py +++ b/test/external/external_test_nv.py @@ -8,7 +8,7 @@ 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.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer +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 7845308bca..6d3a1cb9c1 100644 --- a/test/external/external_test_train_gpt2.py +++ b/test/external/external_test_train_gpt2.py @@ -1,7 +1,7 @@ # ruff: noqa: E501 import unittest -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.search import Opt, OptOps from tinygrad.dtype import dtypes from tinygrad.shape.shapetracker import ShapeTracker diff --git a/test/external/external_test_valid_remove.py b/test/external/external_test_valid_remove.py index 3ccfa23e64..65f8112e96 100644 --- a/test/external/external_test_valid_remove.py +++ b/test/external/external_test_valid_remove.py @@ -2,7 +2,7 @@ import unittest from tinygrad import Device -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.search import Opt, OptOps from tinygrad.dtype import dtypes from tinygrad.shape.shapetracker import ShapeTracker diff --git a/test/external/fuzz_fast_idiv.py b/test/external/fuzz_fast_idiv.py index aef95f3c3a..c7bf90b841 100644 --- a/test/external/fuzz_fast_idiv.py +++ b/test/external/fuzz_fast_idiv.py @@ -1,7 +1,7 @@ import random from z3 import Int, Solver, sat from tinygrad import dtypes, Device -from tinygrad.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher +from tinygrad.uop.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher from tinygrad.codegen.devectorizer import fast_idiv random.seed(42) diff --git a/test/external/fuzz_linearizer.py b/test/external/fuzz_linearizer.py index 308fd008be..8a03176306 100644 --- a/test/external/fuzz_linearizer.py +++ b/test/external/fuzz_linearizer.py @@ -25,7 +25,7 @@ from tinygrad.codegen.kernel import Opt, OptOps from tinygrad.engine.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.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.device import is_dtype_supported def on_linearizer_will_run(): pass diff --git a/test/external/fuzz_schedule.py b/test/external/fuzz_schedule.py index 5703aa8b80..34f3313ada 100644 --- a/test/external/fuzz_schedule.py +++ b/test/external/fuzz_schedule.py @@ -5,7 +5,7 @@ from tinygrad.device import Buffer from tinygrad.engine.realize import capturing, lower_schedule_item from tinygrad.helpers import DEBUG, MULTIOUTPUT, colored, getenv from tinygrad.engine.schedule import LBScheduleItem, _graph_schedule, ScheduleItem -from tinygrad.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp from tinygrad.tensor import Tensor, _to_np_dtype ctx_vars = { MULTIOUTPUT: (0, 1) } diff --git a/test/external/fuzz_symbolic.py b/test/external/fuzz_symbolic.py index a0bde8e22d..8d2242a0c8 100644 --- a/test/external/fuzz_symbolic.py +++ b/test/external/fuzz_symbolic.py @@ -1,7 +1,7 @@ import itertools import random from tinygrad import Variable, dtypes -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp from tinygrad.helpers import DEBUG random.seed(42) @@ -42,7 +42,7 @@ def gt(expr, rng=None): return expr > rng, rng # NOTE: you have to replace these for this test to pass -from tinygrad.ops import python_alu, Ops +from tinygrad.uop.ops import python_alu, Ops python_alu[Ops.MOD] = lambda x,y: x%y python_alu[Ops.IDIV] = lambda x,y: x//y diff --git a/test/external/fuzz_uops.py b/test/external/fuzz_uops.py index 910d91f999..87d73e7cb5 100644 --- a/test/external/fuzz_uops.py +++ b/test/external/fuzz_uops.py @@ -3,11 +3,11 @@ from collections import defaultdict import numpy as np from dataclasses import replace from typing import DefaultDict, Dict, List, Tuple -from tinygrad.ops import UOp, print_uops, Ops +from tinygrad.uop.ops import UOp, print_uops, Ops from tinygrad.device import Buffer, Device from tinygrad.engine.realize import CompiledRunner from tinygrad.helpers import DEBUG, colored -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.tensor import _to_np_dtype from test.external.fuzz_schedule import FUZZ_SCHEDULE_MAX_PATHS, find_all_toposorts diff --git a/test/external/process_replay/process_replay.py b/test/external/process_replay/process_replay.py index a5074be1ce..81609d0b65 100755 --- a/test/external/process_replay/process_replay.py +++ b/test/external/process_replay/process_replay.py @@ -6,7 +6,7 @@ from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connectio from tinygrad.engine.grouper import get_becomes_map from tinygrad.codegen.kernel import Kernel, Opt from tinygrad.renderer import Renderer -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops # *** process replay settings diff --git a/test/helpers.py b/test/helpers.py index 7e1f13a818..9a34198a42 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -2,7 +2,7 @@ import time, struct from typing import Any, Callable, Optional import numpy as np from tinygrad import Tensor, dtypes, Device -from tinygrad.ops import UOp, Ops, sint +from tinygrad.uop.ops import UOp, Ops, sint from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.tensor import _to_np_dtype from tinygrad.engine.realize import Runner diff --git a/test/test_arange.py b/test/test_arange.py index 3109ad0581..0eec281f0c 100644 --- a/test/test_arange.py +++ b/test/test_arange.py @@ -6,7 +6,7 @@ from tinygrad.engine.realize import run_schedule from tinygrad.codegen.kernel import Opt, OptOps, Kernel, KernelOptError from tinygrad.engine.realize import CompiledRunner, ExecItem from tinygrad.engine.search import get_kernel_actions -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops class TestArange(unittest.TestCase): def _get_flops(self, N, opts=None): diff --git a/test/test_const_folding.py b/test/test_const_folding.py index 5ec0572fc7..f1f054da63 100644 --- a/test/test_const_folding.py +++ b/test/test_const_folding.py @@ -2,7 +2,7 @@ import unittest, itertools, math from typing import Any from tinygrad import Tensor, Device, dtypes from tinygrad.dtype import DType -from tinygrad.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp from tinygrad.codegen import full_rewrite_to_sink import numpy as np from tinygrad.device import is_dtype_supported diff --git a/test/test_conv_shapetracker.py b/test/test_conv_shapetracker.py index 7993775263..e812827ac1 100644 --- a/test/test_conv_shapetracker.py +++ b/test/test_conv_shapetracker.py @@ -1,6 +1,6 @@ #!/usr/bin/env python import unittest -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops from tinygrad.tensor import Tensor from tinygrad.nn import Conv2d from tinygrad.shape.shapetracker import ShapeTracker, View diff --git a/test/test_dtype_alu.py b/test/test_dtype_alu.py index 230f372d16..20d5f59b13 100644 --- a/test/test_dtype_alu.py +++ b/test/test_dtype_alu.py @@ -7,7 +7,7 @@ from hypothesis import given, strategies as strat, settings, HealthCheck from tinygrad.dtype import DType from tinygrad.helpers import CI, getenv from tinygrad.engine.realize import run_schedule -from tinygrad.ops import GroupOp +from tinygrad.uop.ops import GroupOp from tinygrad.tensor import _to_np_dtype from tinygrad.device import is_dtype_supported import pytest, math diff --git a/test/test_gc.py b/test/test_gc.py index f13c0c3635..e32053c8c8 100644 --- a/test/test_gc.py +++ b/test/test_gc.py @@ -4,7 +4,7 @@ import unittest import numpy as np from tinygrad.device import Buffer from tinygrad.engine.realize import run_schedule -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp from tinygrad.tensor import Tensor def tensors_allocated(): diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 9ffa28a21b..6466f7d609 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -6,7 +6,7 @@ from dataclasses import replace from test.helpers import ast_const from tinygrad.codegen.kernel import Opt, OptOps, KernelOptError, Kernel from tinygrad.codegen.lowerer import get_grouped_dims -from tinygrad.ops import UOp, Ops, GroupOp +from tinygrad.uop.ops import UOp, Ops, GroupOp from tinygrad.device import Device, Buffer, is_dtype_supported from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View diff --git a/test/test_linearizer_dumb.py b/test/test_linearizer_dumb.py index 23c538c0fb..0cbaea3cf8 100644 --- a/test/test_linearizer_dumb.py +++ b/test/test_linearizer_dumb.py @@ -6,7 +6,7 @@ import unittest from test.helpers import ast_const from tinygrad import Device, dtypes from tinygrad.device import is_dtype_supported -from tinygrad.ops import UOp, Ops +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 diff --git a/test/test_linearizer_failures.py b/test/test_linearizer_failures.py index 04871ae054..931b738845 100644 --- a/test/test_linearizer_failures.py +++ b/test/test_linearizer_failures.py @@ -3,7 +3,7 @@ import unittest, random import numpy as np from tinygrad.codegen.kernel import Kernel, KernelOptError from tinygrad.device import is_dtype_supported -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.search import Opt, OptOps from tinygrad import Device, dtypes, Tensor from tinygrad.helpers import CI, Context diff --git a/test/test_linearizer_overflows.py b/test/test_linearizer_overflows.py index 091d0758d4..502d544c36 100644 --- a/test/test_linearizer_overflows.py +++ b/test/test_linearizer_overflows.py @@ -8,7 +8,7 @@ from tinygrad.engine.search import Opt, OptOps, bufs_from_lin from extra.optimization.helpers import time_linearizer # stuff needed to unpack a kernel -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View diff --git a/test/test_multitensor.py b/test/test_multitensor.py index 7bc21753fb..cd455e74ea 100644 --- a/test/test_multitensor.py +++ b/test/test_multitensor.py @@ -1,7 +1,7 @@ import unittest, functools, random from typing import List from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable -from tinygrad.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp from tinygrad.helpers import CI, getenv, prod, Context, OSX from tinygrad.nn.state import get_parameters, get_state_dict from tinygrad.engine.realize import lower_schedule, BufferCopy, CompiledRunner, run_schedule diff --git a/test/test_nn.py b/test/test_nn.py index e987cdca44..f99b2d2cdb 100755 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -3,7 +3,7 @@ import unittest import numpy as np import torch from tinygrad import Tensor, Device, TinyJit -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops from tinygrad.helpers import GlobalCounters, CI, Context, OSX from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell diff --git a/test/test_pickle.py b/test/test_pickle.py index ffceae6130..e9d88e1817 100644 --- a/test/test_pickle.py +++ b/test/test_pickle.py @@ -2,7 +2,7 @@ import unittest, pickle, types import numpy as np from tinygrad import Tensor, TinyJit, Variable, dtypes from tinygrad.helpers import GlobalCounters, ContextVar, Context -from tinygrad.ops import PatternMatcher, UPat, UOp, Ops +from tinygrad.uop.ops import PatternMatcher, UPat, UOp, Ops class TestPickle(unittest.TestCase): def test_pickle_code_object(self): diff --git a/test/test_quantize_onnx.py b/test/test_quantize_onnx.py index 0c1492b562..9b0fd63f41 100644 --- a/test/test_quantize_onnx.py +++ b/test/test_quantize_onnx.py @@ -3,7 +3,7 @@ import numpy as np import unittest from dataclasses import replace from tinygrad import Tensor, Context, Device, dtypes -from tinygrad.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import +from tinygrad.uop.ops import Ops, UOp # noqa: F401 # pylint: disable=unused-import from tinygrad.codegen.kernel import Kernel, Opt, OptOps from tinygrad.engine.realize import CompiledRunner, ExecItem, lower_schedule_item from tinygrad.engine.search import bufs_from_lin diff --git a/test/test_renderer_failures.py b/test/test_renderer_failures.py index cf2cb08809..329975b94e 100644 --- a/test/test_renderer_failures.py +++ b/test/test_renderer_failures.py @@ -9,7 +9,7 @@ from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.renderer.ptx import PTXRenderer from tinygrad.renderer.wgsl import WGSLRenderer from tinygrad.runtime.ops_python import PythonRenderer -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.renderer import ProgramSpec from tinygrad.tensor import Tensor, _to_np_dtype from tinygrad.codegen import full_rewrite diff --git a/test/test_rewrite_tracked_childen.py b/test/test_rewrite_tracked_childen.py index f409881cef..a1dea7955d 100644 --- a/test/test_rewrite_tracked_childen.py +++ b/test/test_rewrite_tracked_childen.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor -from tinygrad.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp +from tinygrad.uop.ops import PatternMatcher, Ops, UPat, graph_rewrite, RewriteContext, UOp from tinygrad.engine.grouper import sym, merge_views class TestRewriteTrackedChildren(unittest.TestCase): diff --git a/test/test_schedule.py b/test/test_schedule.py index 45d1e7f61b..7d198e83ae 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -11,9 +11,9 @@ from tinygrad import nn, dtypes, Device, Tensor from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType, ImageDType from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites +from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites from tinygrad.codegen.symbolic import symbolic_simple -from tinygrad.spec import type_verify, shape_spec +from tinygrad.uop.spec import type_verify, shape_spec from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp from tinygrad.engine.grouper import view_left, view_right, sym, get_becomes_map, Kernel, create_ast, merge_views, create_kernels from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars diff --git a/test/test_search.py b/test/test_search.py index aa54cbe3cf..19923a6953 100644 --- a/test/test_search.py +++ b/test/test_search.py @@ -3,7 +3,7 @@ import unittest from test.helpers import ast_const from tinygrad.codegen.kernel import Opt, OptOps from tinygrad.codegen.kernel import Kernel -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.engine.search import bufs_from_lin, actions, beam_search from tinygrad.device import Device, Buffer from tinygrad.tensor import Tensor diff --git a/test/test_symbolic_ops.py b/test/test_symbolic_ops.py index 766c4c901b..76bb3e87ef 100644 --- a/test/test_symbolic_ops.py +++ b/test/test_symbolic_ops.py @@ -2,7 +2,7 @@ import unittest from tinygrad import Tensor, Variable from tinygrad.shape.shapetracker import View from tinygrad.helpers import Context, GlobalCounters -from tinygrad.ops import sym_infer +from tinygrad.uop.ops import sym_infer from examples.gpt2 import Attention import numpy as np diff --git a/test/test_tensor.py b/test/test_tensor.py index f10c365e18..8f4403139a 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -8,7 +8,7 @@ from tinygrad.helpers import getenv, temp, mv_address from extra.gradcheck import numerical_jacobian, jacobian, gradcheck from hypothesis import given, settings, strategies as strat from tinygrad.device import is_dtype_supported -from tinygrad.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp from tinygrad.runtime.support.compiler_cuda import PTX from tinygrad.codegen import full_rewrite from tinygrad.dtype import DType diff --git a/test/test_tensor_uop.py b/test/test_tensor_uop.py index 653c17e4e5..d6fc085e06 100644 --- a/test/test_tensor_uop.py +++ b/test/test_tensor_uop.py @@ -3,7 +3,7 @@ import numpy as np import unittest from tinygrad import Tensor, Device, dtypes from tinygrad.engine.realize import run_schedule -from tinygrad.ops import Ops, UOp, UPat +from tinygrad.uop.ops import Ops, UOp, UPat class TestTensorUOp(unittest.TestCase): def test_fromcpu_shape_tracker(self): diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 2c8c8852f1..38418936b6 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -2,7 +2,7 @@ from typing import List import unittest, pytest from tinygrad import dtypes, Variable from tinygrad.helpers import DEBUG, Context -from tinygrad.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite +from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite from tinygrad.codegen.symbolic import sym from tinygrad.codegen import full_rewrite, full_rewrite_to_sink from tinygrad.codegen.expander import expander diff --git a/test/test_uops.py b/test/test_uops.py index dd111be637..ebee1f659c 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -7,8 +7,8 @@ from tinygrad.tensor import Tensor, _to_np_dtype from tinygrad.helpers import CI, DEBUG, getenv, Context, Timing from tinygrad.dtype import dtypes, DType from tinygrad.device import Buffer, Device -from tinygrad.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401 -from tinygrad.spec import spec +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.grouper import fix_kernel_ops from tinygrad.engine.realize import CompiledRunner, get_kernel diff --git a/test/test_uops_stats.py b/test/test_uops_stats.py index a5369bce51..fdaf21c4f7 100644 --- a/test/test_uops_stats.py +++ b/test/test_uops_stats.py @@ -4,7 +4,7 @@ from tinygrad.helpers import getenv, GlobalCounters from tinygrad.engine.realize import lower_schedule_item, ProgramSpec from tinygrad.renderer import Estimates from tinygrad.codegen import full_rewrite -from tinygrad.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp from tinygrad.dtype import dtypes from tinygrad.codegen.kernel import Kernel, Opt, OptOps, KernelOptError from tinygrad.device import Device diff --git a/test/test_winograd.py b/test/test_winograd.py index 26ba62272e..975fe88b69 100644 --- a/test/test_winograd.py +++ b/test/test_winograd.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor, GlobalCounters, dtypes -from tinygrad.ops import Ops +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 diff --git a/test/unit/test_allreduce.py b/test/unit/test_allreduce.py index e99a44b1c6..7eb54bd9d3 100644 --- a/test/unit/test_allreduce.py +++ b/test/unit/test_allreduce.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Tensor from tinygrad.helpers import Context -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops class TestRingAllReduce(unittest.TestCase): @unittest.skip("still broken") diff --git a/test/unit/test_block_reorder.py b/test/unit/test_block_reorder.py index 8459ccb28e..24488d23c8 100644 --- a/test/unit/test_block_reorder.py +++ b/test/unit/test_block_reorder.py @@ -1,6 +1,6 @@ import unittest, random from tinygrad.dtype import dtypes -from tinygrad.ops import print_uops, UOp, Ops +from tinygrad.uop.ops import print_uops, UOp, Ops from tinygrad.codegen.linearize import block_reorder from tinygrad.renderer.cstyle import OpenCLRenderer diff --git a/test/unit/test_gradient.py b/test/unit/test_gradient.py index e82d401d54..3c3290864d 100644 --- a/test/unit/test_gradient.py +++ b/test/unit/test_gradient.py @@ -3,7 +3,7 @@ import unittest, math import torch from tinygrad import Tensor from tinygrad.dtype import dtypes -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp from tinygrad.gradient import compute_gradient class TestGradient(unittest.TestCase): diff --git a/test/unit/test_graph_rewrite.py b/test/unit/test_graph_rewrite.py index ed33c9e764..97816f2e89 100644 --- a/test/unit/test_graph_rewrite.py +++ b/test/unit/test_graph_rewrite.py @@ -1,7 +1,7 @@ import unittest, math from tinygrad import dtypes from tinygrad.helpers import all_same -from tinygrad.ops import GroupOp, UOp, Ops, exec_alu +from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu from tinygrad.codegen import full_rewrite_to_sink # Helper function to apply the graph rewrite @@ -203,7 +203,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase): import inspect -from tinygrad.ops import graph_rewrite, _substitute, track_rewrites +from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites from tinygrad.codegen.symbolic import symbolic_simple class TestBottomUpRewrite(unittest.TestCase): diff --git a/test/unit/test_helpers.py b/test/unit/test_helpers.py index 192e709680..3114e5202f 100644 --- a/test/unit/test_helpers.py +++ b/test/unit/test_helpers.py @@ -290,7 +290,7 @@ class TestPolyN(unittest.TestCase): def test_uop(self): from tinygrad.dtype import dtypes - from tinygrad.ops import UOp + from tinygrad.uop.ops import UOp from test.helpers import eval_uop np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0) np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0) diff --git a/test/unit/test_microbenchmarks.py b/test/unit/test_microbenchmarks.py index 1e9f5112e3..71b43c880c 100644 --- a/test/unit/test_microbenchmarks.py +++ b/test/unit/test_microbenchmarks.py @@ -1,5 +1,5 @@ import unittest, time -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp from tinygrad.dtype import dtypes # it's about 1 ms per 1k UOps on M3 diff --git a/test/unit/test_pattern_matcher.py b/test/unit/test_pattern_matcher.py index cfbd02770c..d1dfffaadd 100644 --- a/test/unit/test_pattern_matcher.py +++ b/test/unit/test_pattern_matcher.py @@ -1,7 +1,7 @@ import unittest, itertools from tinygrad.dtype import dtypes -from tinygrad.ops import Ops, UOp, GroupOp # noqa: F401 -from tinygrad.ops import PatternMatcher, UPat +from tinygrad.uop.ops import Ops, UOp, GroupOp # noqa: F401 +from tinygrad.uop.ops import PatternMatcher, UPat class TestPatternMatcher(unittest.TestCase): def test_simple_match(self): diff --git a/test/unit/test_rewrite_map.py b/test/unit/test_rewrite_map.py index 8220351683..0eed947727 100644 --- a/test/unit/test_rewrite_map.py +++ b/test/unit/test_rewrite_map.py @@ -1,6 +1,6 @@ import unittest from tinygrad import dtypes -from tinygrad.ops import UOp, graph_rewrite_map, _substitute +from tinygrad.uop.ops import UOp, graph_rewrite_map, _substitute from tinygrad.codegen.symbolic import symbolic class TestRewriteMap(unittest.TestCase): diff --git a/test/unit/test_shapetracker.py b/test/unit/test_shapetracker.py index 919766681f..223acceb04 100644 --- a/test/unit/test_shapetracker.py +++ b/test/unit/test_shapetracker.py @@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes from tinygrad.helpers import prod from tinygrad.shape.shapetracker import ShapeTracker, View from tinygrad import Variable -from tinygrad.ops import UOp, Ops, graph_rewrite +from tinygrad.uop.ops import UOp, Ops, graph_rewrite from tinygrad.codegen.devectorizer import sym from itertools import product diff --git a/test/unit/test_simplify_valid_idx.py b/test/unit/test_simplify_valid_idx.py index 5073dbdabe..34377903de 100644 --- a/test/unit/test_simplify_valid_idx.py +++ b/test/unit/test_simplify_valid_idx.py @@ -2,7 +2,7 @@ import unittest, itertools from tinygrad.codegen import full_rewrite_to_sink from tinygrad.dtype import dtypes -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.codegen.symbolic import simplify_valid def get_gated_load_uop(valid:UOp, idx:UOp): diff --git a/test/unit/test_tensor_uop_representation.py b/test/unit/test_tensor_uop_representation.py index fb9245015d..9156d2b6fb 100644 --- a/test/unit/test_tensor_uop_representation.py +++ b/test/unit/test_tensor_uop_representation.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor -from tinygrad.ops import UPat, Ops, UOp +from tinygrad.uop.ops import UPat, Ops, UOp # NOTE: unlike before base for a realized tensor is always a BUFFER realized_pattern = UPat(Ops.BUFFER) diff --git a/test/unit/test_transcendental_helpers.py b/test/unit/test_transcendental_helpers.py index 33cf6c0c20..16a97ad48b 100644 --- a/test/unit/test_transcendental_helpers.py +++ b/test/unit/test_transcendental_helpers.py @@ -1,7 +1,7 @@ import unittest, math import numpy as np from tinygrad import dtypes -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.codegen.transcendental import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction from tinygrad.codegen.transcendental import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if from test.helpers import eval_uop diff --git a/test/unit/test_uop_resolve.py b/test/unit/test_uop_resolve.py index a4bb4889d4..83db03317b 100644 --- a/test/unit/test_uop_resolve.py +++ b/test/unit/test_uop_resolve.py @@ -1,6 +1,6 @@ import unittest from tinygrad.dtype import dtypes -from tinygrad.ops import UOp, resolve +from tinygrad.uop.ops import UOp, resolve class TestUOpResolve(unittest.TestCase): def test_simple_int(self): diff --git a/test/unit/test_uop_symbolic.py b/test/unit/test_uop_symbolic.py index 24f039c6cb..f8482f2468 100644 --- a/test/unit/test_uop_symbolic.py +++ b/test/unit/test_uop_symbolic.py @@ -4,7 +4,7 @@ import unittest, pickle, functools from tinygrad.dtype import dtypes, ConstType from tinygrad.codegen import full_rewrite from tinygrad.codegen.devectorizer import sym -from tinygrad.ops import UOp, Ops, graph_rewrite, sym_infer +from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer from tinygrad import Variable def render(self) -> tuple[str, ConstType, ConstType]: diff --git a/test/unit/test_uop_vmin_vmax.py b/test/unit/test_uop_vmin_vmax.py index c4e71bd06c..cc1ca49dd3 100644 --- a/test/unit/test_uop_vmin_vmax.py +++ b/test/unit/test_uop_vmin_vmax.py @@ -1,5 +1,5 @@ import unittest, math -from tinygrad.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops from tinygrad.dtype import dtypes class TestVminVmaxProperties(unittest.TestCase): diff --git a/test/unit/test_upat_compile.py b/test/unit/test_upat_compile.py index 3fdf8db1ec..1f68a8830b 100644 --- a/test/unit/test_upat_compile.py +++ b/test/unit/test_upat_compile.py @@ -1,8 +1,8 @@ import unittest from tinygrad.helpers import DEBUG from tinygrad.dtype import dtypes -from tinygrad.ops import UPat, track_rewrites, GroupOp, Ops -from tinygrad.upat import _get_code, upat_compile +from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops +from tinygrad.uop.upat import _get_code, upat_compile import dis @track_rewrites() diff --git a/test/unit/test_verify_ast.py b/test/unit/test_verify_ast.py index 907345d13a..eaa23ebcb9 100644 --- a/test/unit/test_verify_ast.py +++ b/test/unit/test_verify_ast.py @@ -4,8 +4,8 @@ import unittest from tinygrad import Tensor from tinygrad.codegen.kernel import Kernel from tinygrad.helpers import DEBUG -from tinygrad.ops import UOp, Ops, print_uops -from tinygrad.spec import type_verify, shape_spec +from tinygrad.uop.ops import UOp, Ops, print_uops +from tinygrad.uop.spec import type_verify, shape_spec from tinygrad.shape.shapetracker import ShapeTracker from tinygrad import dtypes from tinygrad.shape.view import View diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index cb4085efee..d7409c4cc0 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -1,8 +1,8 @@ import unittest, decimal, json from tinygrad.dtype import dtypes -from tinygrad.ops import TRACK_MATCH_STATS, TrackedPatternMatcher, UOp, graph_rewrite, track_rewrites, UPat, Ops +from tinygrad.uop.ops import TRACK_MATCH_STATS, TrackedPatternMatcher, UOp, graph_rewrite, track_rewrites, UPat, Ops from tinygrad.codegen.symbolic import symbolic -from tinygrad.ops import tracked_ctxs as contexts, tracked_keys as keys, _name_cnt, _substitute +from tinygrad.uop.ops import tracked_ctxs as contexts, tracked_keys as keys, _name_cnt, _substitute from tinygrad.device import ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent, ProfileGraphEntry from tinygrad.viz.serve import get_metadata, get_details, uop_to_json, to_perfetto diff --git a/tinygrad/__init__.py b/tinygrad/__init__.py index 947d03c9ed..5bc12f9116 100644 --- a/tinygrad/__init__.py +++ b/tinygrad/__init__.py @@ -4,7 +4,7 @@ if int(os.getenv("TYPED", "0")): install_import_hook(__name__) from tinygrad.tensor import Tensor # noqa: F401 from tinygrad.engine.jit import TinyJit # noqa: F401 -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp Variable = UOp.variable from tinygrad.dtype import dtypes # noqa: F401 from tinygrad.helpers import GlobalCounters, fetch, Context, getenv # noqa: F401 diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 528577adc5..b89549c322 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -2,7 +2,7 @@ from typing import Any, Callable import functools from dataclasses import dataclass from tinygrad.helpers import QUANTIZE, DEVECTORIZE, TRANSCENDENTAL -from tinygrad.ops import PatternMatcher, graph_rewrite, UOp +from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp from tinygrad.renderer import Renderer # import all pattern matchers here diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index 50ee88069f..c3a4911ec3 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -4,7 +4,7 @@ from collections import defaultdict from dataclasses import dataclass from tinygrad.device import is_dtype_supported from tinygrad.dtype import dtypes, ImageDType, PtrDType, promo_lattice, DType -from tinygrad.ops import UOp, Ops, UPat, PatternMatcher, resolve, graph_rewrite, GroupOp, identity_element +from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, resolve, graph_rewrite, GroupOp, identity_element from tinygrad.codegen.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat from tinygrad.helpers import getenv, flatten, AMX, prod, partition from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES diff --git a/tinygrad/codegen/expander.py b/tinygrad/codegen/expander.py index 11f33afe4a..be59b4cdc8 100644 --- a/tinygrad/codegen/expander.py +++ b/tinygrad/codegen/expander.py @@ -2,7 +2,7 @@ import functools, itertools, operator from tinygrad.helpers import AMX, dedup, flatten, all_same, prod -from tinygrad.ops import UOp, Ops, UPat, PatternMatcher, GroupOp +from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp def _expand_arg_to_idx(args:tuple[tuple[int, int], ...], rpk:dict[int, int]) -> int: idx, mul = 0, 1 diff --git a/tinygrad/codegen/heuristic.py b/tinygrad/codegen/heuristic.py index 1b6c28db23..1f77339003 100644 --- a/tinygrad/codegen/heuristic.py +++ b/tinygrad/codegen/heuristic.py @@ -2,7 +2,7 @@ import itertools from tinygrad.codegen.kernel import Kernel, Opt, OptOps, KernelOptError from tinygrad.helpers import getenv, DEBUG, all_int, prod from tinygrad.dtype import ImageDType -from tinygrad.ops import Ops, resolve +from tinygrad.uop.ops import Ops, resolve def hand_coded_optimizations(k:Kernel) -> list[Opt]: # make a copy so it does not mutate the input diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index c7a66a188e..9da24668db 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -4,9 +4,9 @@ from dataclasses import dataclass from collections import defaultdict from typing import Optional, cast, Final, Callable, Sequence -from tinygrad.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, track_rewrites, print_uops, PatternMatcher -from tinygrad.ops import smax -from tinygrad.spec import type_verify, shape_spec +from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, track_rewrites, print_uops +from tinygrad.uop.ops import PatternMatcher, smax +from tinygrad.uop.spec import type_verify, shape_spec from tinygrad.device import Device from tinygrad.renderer import Renderer, TensorCore, ProgramSpec, Opt, OptOps from tinygrad.dtype import ImageDType diff --git a/tinygrad/codegen/linearize.py b/tinygrad/codegen/linearize.py index 6cea5bb8ea..bd9231e157 100644 --- a/tinygrad/codegen/linearize.py +++ b/tinygrad/codegen/linearize.py @@ -2,9 +2,9 @@ from __future__ import annotations import heapq from collections import defaultdict from dataclasses import dataclass, replace -from tinygrad.ops import UOp, Ops, PatternMatcher, UPat, GroupOp +from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp from tinygrad.helpers import dedup, partition, all_same, flatten -from tinygrad.spec import type_verify +from tinygrad.uop.spec import type_verify # NOTE: any toposort should be valid here, unlike last time this isn't required, it's just for speed def block_reorder(lst:list[UOp]) -> list[UOp]: diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index 38ac4203a9..5790fc120f 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -3,7 +3,7 @@ import itertools, operator, math from dataclasses import dataclass from typing import cast from tinygrad.dtype import dtypes, PtrDType, least_upper_dtype -from tinygrad.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint, sint_to_uop +from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint, sint_to_uop from tinygrad.renderer import Renderer from tinygrad.helpers import all_int, prod, partition, flatten, unwrap from tinygrad.codegen.symbolic import symbolic diff --git a/tinygrad/codegen/symbolic.py b/tinygrad/codegen/symbolic.py index 19c78ffa56..661ca60afc 100644 --- a/tinygrad/codegen/symbolic.py +++ b/tinygrad/codegen/symbolic.py @@ -2,7 +2,7 @@ from typing import Any, Literal, cast import math, operator, struct, functools from collections import defaultdict -from tinygrad.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu +from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu from tinygrad.dtype import ConstType, dtypes, PtrDType from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod from tinygrad.codegen.transcendental import xpow diff --git a/tinygrad/codegen/transcendental.py b/tinygrad/codegen/transcendental.py index e2757937ef..3fff6e09bf 100644 --- a/tinygrad/codegen/transcendental.py +++ b/tinygrad/codegen/transcendental.py @@ -1,7 +1,7 @@ import math from tinygrad.dtype import dtypes, DType from tinygrad.helpers import polyN -from tinygrad.ops import UOp +from tinygrad.uop.ops import UOp TRANSCENDENTAL_SUPPORTED_DTYPES = (dtypes.float16, dtypes.float32, dtypes.float64) diff --git a/tinygrad/device.py b/tinygrad/device.py index 325dc817d5..7a3b183dc5 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -382,7 +382,7 @@ if PROFILE: with open(fn:=temp("profile.pkl", append_user=True), "wb") as f: pickle.dump(Compiled.profile_events, f) if not getenv("SQTT", 0): - from tinygrad.ops import launch_viz + from tinygrad.uop.ops import launch_viz launch_viz("PROFILE", fn) if __name__ == "__main__": diff --git a/tinygrad/engine/grouper.py b/tinygrad/engine/grouper.py index 607fc14c35..c684da41a0 100644 --- a/tinygrad/engine/grouper.py +++ b/tinygrad/engine/grouper.py @@ -1,7 +1,7 @@ from collections import defaultdict, deque from dataclasses import dataclass -from tinygrad.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, identity_element, resolve -from tinygrad.ops import can_pad, sint, track_rewrites, _substitute +from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, identity_element, resolve +from tinygrad.uop.ops import can_pad, sint, track_rewrites, _substitute from tinygrad.codegen.lowerer import get_contraction_with_reduce, get_contraction from tinygrad.codegen.symbolic import symbolic_simple from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize, ContextVar, Context, diskcache_put @@ -10,7 +10,7 @@ from tinygrad.dtype import ImageDType from tinygrad.engine.multi import replace_allreduce from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View, strides_for_shape -from tinygrad.spec import type_verify, sched_spec +from tinygrad.uop.spec import type_verify, sched_spec # creation can recurse a lot import sys diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 8e2a75756f..23f24250f7 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, dedup, partition, unwrap from tinygrad.device import Buffer, Compiled, Device, MultiBuffer from tinygrad.dtype import DType -from tinygrad.ops import UOp, Variable, sym_infer, Ops +from tinygrad.uop.ops import UOp, Variable, sym_infer, Ops from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.engine.realize import ExecItem, capturing, ViewOp, BufferCopy, BufferXfer, CompiledRunner, Runner, Estimates from tinygrad.engine.memory import _internal_memory_planner diff --git a/tinygrad/engine/memory.py b/tinygrad/engine/memory.py index c102d9a52e..e925972d1d 100644 --- a/tinygrad/engine/memory.py +++ b/tinygrad/engine/memory.py @@ -3,7 +3,7 @@ from collections import defaultdict from tinygrad.engine.schedule import ScheduleItem from tinygrad.device import Device, Buffer from tinygrad.helpers import NO_MEMORY_PLANNER, dedup, DEBUG, round_up -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops from tinygrad.dtype import dtypes, ImageDType from tinygrad.runtime.support.allocator import TLSFAllocator diff --git a/tinygrad/engine/multi.py b/tinygrad/engine/multi.py index c3013047ea..cbafe4ec57 100644 --- a/tinygrad/engine/multi.py +++ b/tinygrad/engine/multi.py @@ -1,6 +1,6 @@ import functools, itertools, operator from tinygrad.helpers import all_same, all_int, prod, DEBUG, RING, getenv -from tinygrad.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, graph_rewrite_map, track_rewrites +from tinygrad.uop.ops import Ops, UOp, sint, PatternMatcher, UPat, GroupOp, graph_rewrite_map, track_rewrites # *** allreduce implementation *** diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 9fdf596c17..fc4352d11e 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -3,7 +3,7 @@ import time, pprint from dataclasses import dataclass, replace, field from tinygrad.helpers import all_same, colored, getenv, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU -from tinygrad.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer +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 diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index fb650f87b3..829d56266a 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -1,7 +1,7 @@ from typing import cast from dataclasses import dataclass, field from collections import deque, defaultdict -from tinygrad.ops import UOp, Variable, Ops, UPat, PatternMatcher, graph_rewrite, buffers +from tinygrad.uop.ops import UOp, Variable, Ops, UPat, PatternMatcher, graph_rewrite, buffers from tinygrad.device import Buffer, MultiBuffer from tinygrad.helpers import Metadata, unwrap, merge_dicts diff --git a/tinygrad/engine/search.py b/tinygrad/engine/search.py index cd7f31aa1d..739b74df13 100644 --- a/tinygrad/engine/search.py +++ b/tinygrad/engine/search.py @@ -2,7 +2,7 @@ from typing import cast, Optional, Callable import itertools, functools, random, math, time, multiprocessing, traceback, signal, atexit from collections import defaultdict from dataclasses import replace -from tinygrad.ops import UOp, Ops, Variable, sym_infer +from tinygrad.uop.ops import UOp, Ops, Variable, sym_infer 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 diff --git a/tinygrad/gradient.py b/tinygrad/gradient.py index fe6347ce2c..cebaad1b71 100644 --- a/tinygrad/gradient.py +++ b/tinygrad/gradient.py @@ -1,7 +1,7 @@ from typing import cast import math, dataclasses from tinygrad.dtype import dtypes, sum_acc_dtype -from tinygrad.ops import UOp, PatternMatcher, UPat, Ops, all_metadata +from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata from tinygrad.helpers import argsort def reduce_gradient(ctx:UOp, ret:UOp): diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 51717def11..24b536ff8d 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -4,7 +4,7 @@ import functools, math from enum import Enum, auto from dataclasses import dataclass, field, replace from tinygrad.helpers import to_function_name, dedup, prod -from tinygrad.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher +from tinygrad.uop.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp, PatternMatcher from tinygrad.dtype import DType class OptOps(Enum): diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index d2d1e730ff..efbffe1330 100644 --- a/tinygrad/renderer/cstyle.py +++ b/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.ops import GroupOp, Ops, UOp, PatternMatcher, UPat +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 from tinygrad.renderer import Renderer, TensorCore diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 0292663d3b..69b645256c 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -2,7 +2,7 @@ from typing import cast import math, struct, sys from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import ClangRenderer, AMDRenderer -from tinygrad.ops import UOp, PatternMatcher, UPat, Ops, GroupOp +from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp from tinygrad.dtype import dtypes, DType, PtrDType, truncate from tinygrad.helpers import prod, AMX diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 585009c538..c47074d682 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -1,7 +1,7 @@ from typing import cast, Callable import struct from collections import defaultdict -from tinygrad.ops import Ops, UOp, PatternMatcher, UPat, GroupOp +from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp from tinygrad.dtype import dtypes, DType, PtrDType from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import CUDARenderer diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 4a6e1c6736..03337f8e3e 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -1,5 +1,5 @@ from tinygrad.dtype import DType, PtrDType, dtypes -from tinygrad.ops import UOp, Ops, PatternMatcher, UPat +from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite, extra_pm from tinygrad.helpers import strip_parens import math diff --git a/tinygrad/runtime/graph/cpu.py b/tinygrad/runtime/graph/cpu.py index f2479cfaa0..390c35a204 100644 --- a/tinygrad/runtime/graph/cpu.py +++ b/tinygrad/runtime/graph/cpu.py @@ -4,7 +4,7 @@ from tinygrad.helpers import dedup, flatten, DEBUG, to_function_name from tinygrad.engine.jit import GraphRunner, GraphException from tinygrad.device import Buffer from tinygrad.engine.realize import ExecItem, CompiledRunner -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.dtype import DType, dtypes from tinygrad.renderer.cstyle import ClangRenderer from tinygrad.renderer.llvmir import LLVMRenderer, ldt diff --git a/tinygrad/runtime/graph/cuda.py b/tinygrad/runtime/graph/cuda.py index c7f6708b24..7ab99ad0bd 100644 --- a/tinygrad/runtime/graph/cuda.py +++ b/tinygrad/runtime/graph/cuda.py @@ -4,7 +4,7 @@ import tinygrad.runtime.autogen.cuda as cuda from tinygrad.helpers import init_c_var, dedup from tinygrad.device import Buffer, Device from tinygrad.runtime.ops_cuda import CUDADevice, check, encode_args, cu_time_execution -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner from tinygrad.engine.jit import MultiGraphRunner, GraphException diff --git a/tinygrad/runtime/graph/hcq.py b/tinygrad/runtime/graph/hcq.py index af8bf36296..4bcf17dba5 100644 --- a/tinygrad/runtime/graph/hcq.py +++ b/tinygrad/runtime/graph/hcq.py @@ -4,7 +4,7 @@ from tinygrad.helpers import round_up, PROFILE, merge_dicts from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent from tinygrad.dtype import dtypes -from tinygrad.ops import UOp, Variable +from tinygrad.uop.ops import UOp, Variable from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner from tinygrad.engine.jit import MultiGraphRunner diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index 0d4b42afc8..d8fd14abec 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -5,7 +5,7 @@ from tinygrad.helpers import dedup, getenv, merge_dicts from tinygrad.device import Buffer from tinygrad.engine.realize import ExecItem, CompiledRunner from tinygrad.engine.jit import GraphRunner, GraphException -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.runtime.ops_metal import wait_check, msg, libobjc, to_struct, objc_instance,\ MTLResourceOptions, cmdbuf_st_time, cmdbuf_en_time, objc_id, to_ns_str diff --git a/tinygrad/runtime/graph/remote.py b/tinygrad/runtime/graph/remote.py index f3ffc3be7a..04c6dbdbc9 100644 --- a/tinygrad/runtime/graph/remote.py +++ b/tinygrad/runtime/graph/remote.py @@ -1,4 +1,4 @@ -from tinygrad.ops import Variable +from tinygrad.uop.ops import Variable from tinygrad.engine.jit import GraphRunner, MultiGraphRunner, GraphException from tinygrad.engine.realize import CompiledRunner, BufferXfer, ExecItem from tinygrad.device import Device, Buffer diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index 2461ef04ab..e0f99014fa 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -5,7 +5,7 @@ assert sys.platform != 'win32' from dataclasses import dataclass from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface from tinygrad.runtime.support.hcq import MMIOInterface -from tinygrad.ops import sint +from tinygrad.uop.ops import sint from tinygrad.device import Compiled, ProfileEvent, BufferSpec, CPUProgram, PROFILE from tinygrad.helpers import getenv, to_mv, round_up, data64_le, all_same, flatten, DEBUG, OSX from tinygrad.renderer.cstyle import AMDRenderer diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 274a40d7e8..3264be212f 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -3,13 +3,13 @@ import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, context assert sys.platform != 'win32' from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, MallocAllocator from tinygrad.dtype import dtypes, DType, PtrDType -from tinygrad.ops import Ops, UOp +from tinygrad.uop.ops import Ops, UOp from tinygrad.helpers import from_mv, getenv, round_up, mv_address, to_mv, cpu_objdump, DEBUG from tinygrad.renderer.cstyle import ClangRenderer from tinygrad.runtime.autogen import libc, qcom_dsp if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import -from tinygrad.ops import PatternMatcher, UPat +from tinygrad.uop.ops import PatternMatcher, UPat dsp_pm = PatternMatcher([ (((UPat.var('x').maximum(0) ^ -1).maximum(-256) ^ -1).cast(dtypes.uchar.vec(128)), diff --git a/tinygrad/runtime/ops_null.py b/tinygrad/runtime/ops_null.py index ed2a657407..569888d8c5 100644 --- a/tinygrad/runtime/ops_null.py +++ b/tinygrad/runtime/ops_null.py @@ -1,5 +1,5 @@ from tinygrad.device import Compiled, Compiler, Renderer, Allocator -from tinygrad.ops import Ops +from tinygrad.uop.ops import Ops from tinygrad.engine.jit import MultiGraphRunner class NullRenderer(Renderer): diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index c45341408b..c2d1d09b14 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -5,7 +5,7 @@ from typing import Any, cast, Union, Type, ClassVar from dataclasses import dataclass from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, MOCKGPU -from tinygrad.ops import sint +from tinygrad.uop.ops import sint from tinygrad.device import BufferSpec, CPUProgram from tinygrad.helpers import getenv, mv_address, init_c_struct_t, round_up, data64, data64_le, DEBUG, prod, OSX from tinygrad.renderer.ptx import PTXRenderer diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 9ebb3be764..0cdc598507 100644 --- a/tinygrad/runtime/ops_python.py +++ b/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.ops import exec_alu, Ops, UOp, GroupOp +from tinygrad.uop.ops import exec_alu, Ops, UOp, GroupOp from tinygrad.renderer import Renderer from tinygrad.renderer.cstyle import CUDARenderer, MetalRenderer, AMDRenderer, IntelRenderer, ClangRenderer diff --git a/tinygrad/runtime/ops_remote.py b/tinygrad/runtime/ops_remote.py index 54d783c65f..2ef89667e1 100644 --- a/tinygrad/runtime/ops_remote.py +++ b/tinygrad/runtime/ops_remote.py @@ -11,7 +11,7 @@ from dataclasses import dataclass, field, replace import multiprocessing, functools, itertools, asyncio, http, http.client, hashlib, time, os, binascii, struct, ast, contextlib, weakref from tinygrad.renderer import Renderer, ProgramSpec from tinygrad.dtype import DTYPES_DICT, dtypes -from tinygrad.ops import UOp, Ops, Variable, sint +from tinygrad.uop.ops import UOp, Ops, Variable, sint from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, Timing from tinygrad.engine.jit import GraphRunner, MultiGraphRunner, ExecItem, graph_class from tinygrad.engine.realize import CompiledRunner, BufferXfer diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index d82a3d1daa..d24c2fb963 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -4,7 +4,7 @@ import contextlib, decimal, statistics, time, ctypes, array, os, fcntl, struct from tinygrad.helpers import PROFILE, getenv, to_mv, round_up from tinygrad.renderer import Renderer from tinygrad.device import BufferSpec, Compiler, Compiled, LRUAllocator, ProfileRangeEvent, ProfileDeviceEvent, ProfileProgramEvent -from tinygrad.ops import sym_infer, sint, Variable, UOp +from tinygrad.uop.ops import sym_infer, sint, Variable, UOp from tinygrad.runtime.autogen import libc class MMIOInterface: diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 4dd8b8ce3c..031fb64e2b 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -6,7 +6,7 @@ from typing import Optional, Callable from tinygrad.helpers import merge_dicts, getenv from tinygrad.shape.view import View, strides_for_shape, unravel from tinygrad.dtype import dtypes -from tinygrad.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context, PatternMatcher, UPat, GroupOp +from tinygrad.uop.ops import UOp, Ops, graph_rewrite, Variable, sint, sint_to_uop, Context, PatternMatcher, UPat, GroupOp from tinygrad.codegen.symbolic import split_uop, symbolic_flat, uop_given_valid, simplify_valid # If a node overflow, its srcs need to be checked to see if this overflow is the result of an ALU operation, diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index a7bd6710ec..bc7ba0716c 100644 --- a/tinygrad/shape/view.py +++ b/tinygrad/shape/view.py @@ -3,7 +3,7 @@ import functools, operator, itertools from dataclasses import dataclass from typing import Optional, cast, Sequence from tinygrad.dtype import dtypes -from tinygrad.ops import resolve, UOp, Variable, sint, sym_infer, smax, smin, sint_to_uop, Ops +from tinygrad.uop.ops import resolve, UOp, Variable, sint, sym_infer, smax, smin, sint_to_uop, Ops from tinygrad.helpers import prod, all_int, argsort, flatten, ceildiv @functools.cache diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index d04a9712e7..b1a55df4b3 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -9,8 +9,8 @@ from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_u from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG from tinygrad.engine.multi import get_multi_map from tinygrad.gradient import compute_gradient -from tinygrad.ops import smax, smin, resolve, UOp, Ops, sint, Variable, SimpleMathTrait, identity_element, all_metadata -from tinygrad.spec import tensor_uop_spec, type_verify +from tinygrad.uop.ops import smax, smin, resolve, UOp, Ops, sint, Variable, SimpleMathTrait, identity_element, all_metadata +from tinygrad.uop.spec import tensor_uop_spec, type_verify from tinygrad.device import Device, Buffer from tinygrad.engine.realize import run_schedule from tinygrad.engine.memory import memory_planner diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tinygrad/ops.py b/tinygrad/uop/ops.py similarity index 99% rename from tinygrad/ops.py rename to tinygrad/uop/ops.py index cc6437a31b..dcee9d8f92 100644 --- a/tinygrad/ops.py +++ b/tinygrad/uop/ops.py @@ -845,7 +845,7 @@ def upat_interpret(p:UPat, fxn:Callable) -> Callable: class PatternMatcher: def __init__(self, patterns:Sequence[tuple[UPat, Callable|tuple]], compiled=bool(getenv("UPAT_COMPILE", 1))): - if compiled: from tinygrad.upat import upat_compile + if compiled: from tinygrad.uop.upat import upat_compile # if this comes from a pickle, we reconstruct the lambda functions here self.patterns:list[tuple[UPat, Callable]] = [(p,types.FunctionType(*fxn) if isinstance(fxn, tuple) else fxn) for p,fxn in patterns] # NOTE: use of DefaultDict here is very dangerous! all keys will live for the lifetime of the PatternMatcher! @@ -954,7 +954,7 @@ def launch_viz(env_str:str, data:str): if not int(os.getenv("VIZ", "0")) and not int(os.getenv("PROFILE", "0")): args = ['--kernels', getenv("VIZ_DATA", "")] if getenv("VIZ_DATA", "") else [] args += ['--profile', getenv("PROFILE_DATA", "")] if getenv("PROFILE_DATA", "") else [] - os.execv(sys.executable, [sys.executable] + [os.path.join(os.path.dirname(__file__), ".", "viz", "serve.py")] + args) + os.execv(sys.executable, [sys.executable] + [os.path.join(os.path.dirname(__file__), "../", "viz", "serve.py")] + args) # *** simple graph rewrite engine *** diff --git a/tinygrad/spec.py b/tinygrad/uop/spec.py similarity index 99% rename from tinygrad/spec.py rename to tinygrad/uop/spec.py index e13757678a..9a5e810f63 100644 --- a/tinygrad/spec.py +++ b/tinygrad/uop/spec.py @@ -1,5 +1,5 @@ from typing import cast, Callable -from tinygrad.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, resolve +from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, resolve from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType from tinygrad.helpers import all_same, prod, DEBUG, IGNORE_OOB, Context try: diff --git a/tinygrad/upat.py b/tinygrad/uop/upat.py similarity index 98% rename from tinygrad/upat.py rename to tinygrad/uop/upat.py index 17007f6492..090de21357 100644 --- a/tinygrad/upat.py +++ b/tinygrad/uop/upat.py @@ -1,7 +1,7 @@ from typing import Any, Callable import itertools, inspect, functools, types from tinygrad.helpers import partition, dedup, Context -from tinygrad.ops import UPat, UPatAny, UOp, Ops, PatternMatcher, graph_rewrite, deconstruct_function +from tinygrad.uop.ops import UPat, UPatAny, UOp, Ops, PatternMatcher, graph_rewrite, deconstruct_function class UPatCompileError(Exception): pass diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 9305fca8bd..d33952b26f 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -4,7 +4,7 @@ from http.server import BaseHTTPRequestHandler from urllib.parse import parse_qs, urlparse from typing import Any, Callable, TypedDict, Generator from tinygrad.helpers import colored, getenv, tqdm, unwrap, word_wrap, TRACEMETA -from tinygrad.ops import TrackedGraphRewrite, UOp, Ops, lines, GroupOp, srender, sint +from tinygrad.uop.ops import TrackedGraphRewrite, UOp, Ops, lines, GroupOp, srender, sint from tinygrad.codegen.kernel import Kernel from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent from tinygrad.dtype import dtypes