From a38947b4bb08ccae14183baa6d0453906ef74f17 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 10 Jun 2025 20:51:22 -0700 Subject: [PATCH] move symbolic and transcendental to uop [pr] (#10771) --- .pre-commit-config.yaml | 2 +- test/test_schedule.py | 2 +- test/test_uop_graph.py | 2 +- test/test_uops.py | 2 +- test/unit/test_graph_rewrite.py | 2 +- test/unit/test_rewrite_map.py | 2 +- test/unit/test_simplify_valid_idx.py | 2 +- test/unit/test_transcendental_helpers.py | 4 ++-- test/unit/test_viz.py | 2 +- tinygrad/codegen/__init__.py | 2 +- tinygrad/codegen/devectorizer.py | 4 ++-- tinygrad/codegen/lowerer.py | 2 +- tinygrad/engine/grouper.py | 2 +- tinygrad/shape/shapetracker.py | 2 +- tinygrad/uop/ops.py | 2 +- tinygrad/{codegen => uop}/symbolic.py | 2 +- tinygrad/{codegen => uop}/transcendental.py | 0 17 files changed, 18 insertions(+), 18 deletions(-) rename tinygrad/{codegen => uop}/symbolic.py (99%) rename tinygrad/{codegen => uop}/transcendental.py (100%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0e23d11069..b6c7ef1779 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -28,7 +28,7 @@ repos: pass_filenames: false - id: tests name: subset of tests - entry: env MAX_BUFFER_SIZE=300000000 PYTHONPATH="." python3 -m pytest -n=4 --ignore=test/unit/test_keccak.py --ignore=test/unit/test_indexing.py test/unit/ test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py test/test_symbolic_shapetracker.py + entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py language: system always_run: true pass_filenames: false diff --git a/test/test_schedule.py b/test/test_schedule.py index cbcb6a6ca4..7a8a28f8cc 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -13,7 +13,7 @@ from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType, ImageDType from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites -from tinygrad.codegen.symbolic import symbolic_simple +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.grouper import view_left, view_right, sym, get_kernelize_map, Kernel, create_ast, merge_views, create_kernels from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 0ad6f072a2..d9089cb09c 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -3,7 +3,7 @@ import unittest, pytest from tinygrad import dtypes, Variable from tinygrad.helpers import DEBUG, Context from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp -from tinygrad.codegen.symbolic import sym +from tinygrad.uop.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 779cac713d..c5bf5e4183 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -13,7 +13,7 @@ from tinygrad.renderer import ProgramSpec from tinygrad.engine.grouper import fix_kernel_ops from tinygrad.engine.realize import CompiledRunner, get_kernel from tinygrad.codegen import full_rewrite -from tinygrad.codegen.symbolic import sym +from tinygrad.uop.symbolic import sym from tinygrad.device import is_dtype_supported from tinygrad.codegen.kernel import Kernel, Opt, OptOps diff --git a/test/unit/test_graph_rewrite.py b/test/unit/test_graph_rewrite.py index 50a1f93db3..327140b4a1 100644 --- a/test/unit/test_graph_rewrite.py +++ b/test/unit/test_graph_rewrite.py @@ -204,7 +204,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase): import inspect from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites -from tinygrad.codegen.symbolic import symbolic_simple +from tinygrad.uop.symbolic import symbolic_simple class TestBottomUpRewrite(unittest.TestCase): def test_const_folding(self): diff --git a/test/unit/test_rewrite_map.py b/test/unit/test_rewrite_map.py index 0eed947727..adc8cd2f8a 100644 --- a/test/unit/test_rewrite_map.py +++ b/test/unit/test_rewrite_map.py @@ -1,7 +1,7 @@ import unittest from tinygrad import dtypes from tinygrad.uop.ops import UOp, graph_rewrite_map, _substitute -from tinygrad.codegen.symbolic import symbolic +from tinygrad.uop.symbolic import symbolic class TestRewriteMap(unittest.TestCase): def test_substitute(self): diff --git a/test/unit/test_simplify_valid_idx.py b/test/unit/test_simplify_valid_idx.py index 656821ba12..8ef89d56bf 100644 --- a/test/unit/test_simplify_valid_idx.py +++ b/test/unit/test_simplify_valid_idx.py @@ -3,7 +3,7 @@ import unittest, itertools from tinygrad.codegen import full_rewrite_to_sink from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops -from tinygrad.codegen.symbolic import simplify_valid +from tinygrad.uop.symbolic import simplify_valid def get_gated_load_uop(valid:UOp, idx:UOp): return UOp(Ops.LOAD, dtypes.float, ( diff --git a/test/unit/test_transcendental_helpers.py b/test/unit/test_transcendental_helpers.py index 16a97ad48b..5a14b381de 100644 --- a/test/unit/test_transcendental_helpers.py +++ b/test/unit/test_transcendental_helpers.py @@ -2,8 +2,8 @@ import unittest, math import numpy as np from tinygrad import dtypes 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 tinygrad.uop.transcendental import TRANSCENDENTAL_SUPPORTED_DTYPES, payne_hanek_reduction, cody_waite_reduction +from tinygrad.uop.transcendental import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if from test.helpers import eval_uop class TestTranscendentalFunctions(unittest.TestCase): diff --git a/test/unit/test_viz.py b/test/unit/test_viz.py index 90e7924847..6f413b7924 100644 --- a/test/unit/test_viz.py +++ b/test/unit/test_viz.py @@ -1,7 +1,7 @@ import unittest, decimal, json from tinygrad.dtype import dtypes from tinygrad.uop.ops import TRACK_MATCH_STATS, TrackedPatternMatcher, UOp, graph_rewrite, track_rewrites, UPat, Ops -from tinygrad.codegen.symbolic import symbolic +from tinygrad.uop.symbolic import symbolic 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/codegen/__init__.py b/tinygrad/codegen/__init__.py index fc3eb412bd..24b4ce72e9 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -7,7 +7,7 @@ from tinygrad.renderer import Renderer # import all pattern matchers here from tinygrad.codegen.lowerer import pm_quant, pm_lowerer, get_index -from tinygrad.codegen.symbolic import sym, symbolic_simple, gep_pushing +from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing from tinygrad.codegen.expander import migrate_indexing, pm_store_ignore, pm_move_ignore, pm_delete_ignore, expander from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, \ pm_reduce, ReduceContext, correct_load_store, pm_render, get_late_rewrite_patterns diff --git a/tinygrad/codegen/devectorizer.py b/tinygrad/codegen/devectorizer.py index c3a4911ec3..a6aae9d261 100644 --- a/tinygrad/codegen/devectorizer.py +++ b/tinygrad/codegen/devectorizer.py @@ -5,9 +5,9 @@ from dataclasses import dataclass from tinygrad.device import is_dtype_supported from tinygrad.dtype import dtypes, ImageDType, PtrDType, promo_lattice, DType 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.uop.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 +from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES from tinygrad.renderer import Renderer # ***** image load valid simplification ***** diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index 7922330075..41fd650933 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -6,7 +6,7 @@ from tinygrad.dtype import dtypes, PtrDType, least_upper_dtype 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 +from tinygrad.uop.symbolic import symbolic # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None: diff --git a/tinygrad/engine/grouper.py b/tinygrad/engine/grouper.py index d04669db27..b901f06226 100644 --- a/tinygrad/engine/grouper.py +++ b/tinygrad/engine/grouper.py @@ -3,7 +3,7 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewr from tinygrad.uop.ops import track_rewrites, _substitute from tinygrad.uop.spec import type_verify, tensor_uop_spec from tinygrad.codegen.lowerer import get_contraction_with_reduce -from tinygrad.codegen.symbolic import symbolic_simple +from tinygrad.uop.symbolic import symbolic_simple from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize from tinygrad.helpers import FUSE_CONV_BW, FUSE_ARANGE, DEBUG, DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES, SPLIT_REDUCEOP from tinygrad.dtype import ImageDType diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 031fb64e2b..ea716d9193 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -7,7 +7,7 @@ from tinygrad.helpers import merge_dicts, getenv from tinygrad.shape.view import View, strides_for_shape, unravel from tinygrad.dtype import dtypes 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 +from tinygrad.uop.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, # or that the node simply inherits the dtype from srcs. Upcast is either `Ops.CAST`+`replace` or just `replace`. diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 40122499c9..d0e71030b6 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -171,7 +171,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): def simplify(self): # late import! - from tinygrad.codegen.symbolic import symbolic + from tinygrad.uop.symbolic import symbolic with Context(TRACK_MATCH_STATS=0): return graph_rewrite(self, symbolic) def ssimplify(self) -> Union[UOp, ConstType]: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret diff --git a/tinygrad/codegen/symbolic.py b/tinygrad/uop/symbolic.py similarity index 99% rename from tinygrad/codegen/symbolic.py rename to tinygrad/uop/symbolic.py index f293e66bf4..c93d31d05a 100644 --- a/tinygrad/codegen/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -5,7 +5,7 @@ from collections import defaultdict 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, CORRECT_DIVMOD_FOLDING -from tinygrad.codegen.transcendental import xpow +from tinygrad.uop.transcendental import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** diff --git a/tinygrad/codegen/transcendental.py b/tinygrad/uop/transcendental.py similarity index 100% rename from tinygrad/codegen/transcendental.py rename to tinygrad/uop/transcendental.py