From 54924f9969e322b9fa74fb88ed2e7536f4573985 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 19 Jul 2025 14:05:52 -0400 Subject: [PATCH] type remove Union and Optional [pr] (#11283) use `|` for consistency --- test/test_linearizer.py | 5 ++--- test/test_schedule.py | 8 +++---- tinygrad/device.py | 28 ++++++++++++------------ tinygrad/dtype.py | 16 +++++++------- tinygrad/engine/jit.py | 14 ++++++------ tinygrad/engine/realize.py | 18 +++++++-------- tinygrad/helpers.py | 20 ++++++++--------- tinygrad/nn/state.py | 14 ++++++------ tinygrad/opt/kernel.py | 23 ++++++++++--------- tinygrad/opt/search.py | 10 ++++----- tinygrad/renderer/__init__.py | 16 +++++++------- tinygrad/runtime/ops_disk.py | 6 ++--- tinygrad/runtime/ops_gpu.py | 8 ++++--- tinygrad/runtime/ops_metal.py | 4 ++-- tinygrad/runtime/ops_nv.py | 4 ++-- tinygrad/runtime/ops_python.py | 4 ++-- tinygrad/runtime/ops_remote.py | 4 ++-- tinygrad/shape/shapetracker.py | 14 ++++++------ tinygrad/shape/view.py | 20 ++++++++--------- tinygrad/tensor.py | 4 ++-- tinygrad/uop/ops.py | 40 +++++++++++++++++----------------- 21 files changed, 140 insertions(+), 140 deletions(-) diff --git a/test/test_linearizer.py b/test/test_linearizer.py index a1002250bd..39daf7089e 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -1,4 +1,3 @@ -from typing import Union import numpy as np import unittest from dataclasses import replace @@ -15,7 +14,7 @@ from tinygrad.opt.heuristic import hand_coded_optimizations from tinygrad.helpers import prod, Context, getenv, CI, flatten, dedup, AMX, AMD_LLVM from tinygrad.dtype import DType, dtypes -def helper_realized_ast(r:Union[Tensor, list[Tensor]]) -> tuple[UOp, list[Buffer]]: +def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]: if isinstance(r, Tensor): r = [r] s = Tensor.schedule(*r) run_schedule(s[:-1]) # run all kernels except the last one @@ -1079,7 +1078,7 @@ def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs): for out in ast.src] return _helper_linearizer_opt_ast(ast, outbufs+inbufs, *args, **kwargs) -def helper_linearizer_opt(r:Union[Tensor, list[Tensor]], *args, **kwargs): +def helper_linearizer_opt(r:Tensor|list[Tensor], *args, **kwargs): realized_ast, real_bufs = helper_realized_ast(r) return _helper_linearizer_opt_ast(realized_ast, real_bufs, *args, **kwargs) diff --git a/test/test_schedule.py b/test/test_schedule.py index ae04c4690e..140b3d1bd2 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -5,7 +5,7 @@ import unittest import numpy as np import functools -from typing import List, Optional, Union, cast +from typing import cast from hypothesis import assume, given, strategies as strat from tinygrad import nn, dtypes, Device, Tensor @@ -20,11 +20,11 @@ from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule class KernelCountException(Exception): pass -def check_schedule(t:Union[Tensor, List[Tensor], UOp], allowed:int, to_prerealize:Optional[List[Tensor]]=None, filter_sink=True): +def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True): if to_prerealize: with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize) if isinstance(t, Tensor): sched = t.schedule() - elif isinstance(t, List) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t) + elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t) else: assert isinstance(t, UOp), f"can't schedule {t}" sink = UOp.sink(t) if t.op is not Ops.SINK else t @@ -1727,7 +1727,7 @@ class TestSchedule(unittest.TestCase): np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]]) class TestIndexing(unittest.TestCase): - def check_schedule(self, xt:Union[Tensor,List[Tensor]], cnt:int): + def check_schedule(self, xt:Tensor|list[Tensor], cnt:int): with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)): lst = [xt] if isinstance(xt, Tensor) else xt s = Tensor.schedule(*lst) diff --git a/tinygrad/device.py b/tinygrad/device.py index 7c8622a45e..06ef872221 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, replace, field from collections import defaultdict -from typing import Optional, Any, Generic, TypeVar, Iterator +from typing import Any, Generic, TypeVar, Iterator import importlib, inspect, functools, pathlib, os, ctypes, ctypes.util, platform, contextlib, sys, re, atexit, pickle, decimal, time from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, from_mv, PROFILE, temp, mv_address, \ cpu_time_execution, colored, Context, round_up, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, cpu_events, ProfileEvent @@ -18,7 +18,7 @@ class _Device: @functools.cache # this class is a singleton, pylint: disable=method-cache-max-size-none def _canonicalize(self, device:str) -> str: return re.sub(r":0$", "", (d:=device.split(":", 1)[0].upper()) + device[len(d):]) # NOTE: you can't cache canonicalize in case Device.DEFAULT changes - def canonicalize(self, device:Optional[str]) -> str: return self._canonicalize(device if device is not None else Device.DEFAULT) + def canonicalize(self, device:str|None) -> str: return self._canonicalize(device if device is not None else Device.DEFAULT) def __getitem__(self, ix:str) -> Compiled: return self.__get_canonicalized_item(self.canonicalize(ix)) @functools.cache # this class is a singleton, pylint: disable=method-cache-max-size-none def __get_canonicalized_item(self, ix:str) -> Compiled: @@ -72,12 +72,12 @@ class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[l @dataclass(frozen=True, eq=True) class BufferSpec: # TODO: move device, size, dtype here? - image: Optional[ImageDType] = None + image: ImageDType|None = None uncached: bool = False cpu_access: bool = False host: bool = False nolru: bool = False - external_ptr: Optional[int] = None + external_ptr: int|None = None class MultiBuffer: def __init__(self, device:tuple[str, ...], size:int, dtype:DType): @@ -94,8 +94,8 @@ class MultiBuffer: class Buffer: profile_events:list[ProfileEvent] = [] - def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:Optional[BufferSpec]=None, initial_value:Optional[bytes]=None, - uop_refcount=0, base:Optional[Buffer]=None, offset:int=0, preallocate=False): + def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|None=None, + uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False): if isinstance(dtype, ImageDType): options = BufferSpec(image=dtype) # TODO: image hack shouldn't be here. where should it be? else: assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType) self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0 @@ -223,10 +223,10 @@ class Allocator(Generic[DeviceType]): self.dev: DeviceType = dev self.default_buffer_spec: BufferSpec = BufferSpec() # overridden in LRUAllocator - def alloc(self, size:int, options:Optional[BufferSpec]=None): + def alloc(self, size:int, options:BufferSpec|None=None): assert size > 0, f"alloc size must be positive, getting {size}" return self._alloc(size, options if options is not None else self.default_buffer_spec) - def free(self, opaque, size:int, options:Optional[BufferSpec]=None): + def free(self, opaque, size:int, options:BufferSpec|None=None): self._free(opaque, options if options is not None else self.default_buffer_spec) # implemented by the runtime @@ -244,9 +244,9 @@ class LRUAllocator(Allocator, Generic[DeviceType]): It ensures that buffers are not freed until it is absolutely necessary, optimizing performance. """ def __init__(self, dev:DeviceType): - self.cache: dict[tuple[int, Optional[BufferSpec]], Any] = defaultdict(list) + self.cache: dict[tuple[int, BufferSpec|None], Any] = defaultdict(list) super().__init__(dev) - def alloc(self, size:int, options:Optional[BufferSpec]=None): + def alloc(self, size:int, options:BufferSpec|None=None): if len(c := self.cache[(size, options)]): return c.pop() try: return super().alloc(size, options) except (RuntimeError, MemoryError): @@ -256,7 +256,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]): for (sz,options),opaques in self.cache.items(): for opaque in opaques: super().free(opaque, sz, options) opaques.clear() - def free(self, opaque:Any, size:int, options:Optional[BufferSpec]=None): + def free(self, opaque:Any, size:int, options:BufferSpec|None=None): if LRU and (options is None or not options.nolru): self.cache[(size, options)].append(opaque) else: super().free(opaque, size, options) @@ -333,7 +333,7 @@ class CPUProgram: class CompileError(Exception): pass class Compiler: - def __init__(self, cachekey:Optional[str]=None): self.cachekey = None if DISABLE_COMPILER_CACHE else cachekey + def __init__(self, cachekey:str|None=None): self.cachekey = None if DISABLE_COMPILER_CACHE else cachekey def compile(self, src:str) -> bytes: return src.encode() # NOTE: empty compiler is the default def compile_cached(self, src:str) -> bytes: if self.cachekey is None or (lib := diskcache_get(self.cachekey, src)) is None: @@ -346,7 +346,7 @@ class Compiler: class Compiled: profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device. - def __init__(self, device:str, allocator:Allocator, renderer:Optional[Renderer], compiler:Optional[Compiler], runtime, graph=None): + def __init__(self, device:str, allocator:Allocator, renderer:Renderer|None, compiler:Compiler|None, runtime, graph=None): self.device, self.allocator, self.compiler, self.runtime, self.graph = device, allocator, compiler or Compiler(), runtime, graph self.renderer = renderer or Renderer() def synchronize(self): @@ -368,7 +368,7 @@ class Compiled: # override this in your device implementation # TODO: move this to each Device -def is_dtype_supported(dtype:DType, device:Optional[str]=None) -> bool: +def is_dtype_supported(dtype:DType, device:str|None=None) -> bool: if device is None: device = Device.DEFAULT if dtype == dtypes.bfloat16: if device == "METAL": return not CI diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index d7fe21aa6e..bb1d2eb12d 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -1,10 +1,10 @@ from __future__ import annotations -from typing import Final, Optional, ClassVar, Union, Callable, Literal +from typing import Final, ClassVar, Callable, Literal import math, struct, ctypes, functools from dataclasses import dataclass, fields from tinygrad.helpers import getenv, prod -ConstType = Union[float, int, bool] +ConstType = float|int|bool FmtStr = Literal['?', 'b', 'B', 'h', 'H', 'i', 'I', 'q', 'Q', 'e', 'f', 'd'] @@ -21,11 +21,11 @@ class DType(metaclass=DTypeMetaClass): priority: int # this determines when things get upcasted itemsize: int name: str - fmt: Optional[FmtStr] + fmt: FmtStr|None count: int - _scalar: Optional[DType] + _scalar: DType|None @staticmethod - def new(priority:int, itemsize:int, name:str, fmt:Optional[FmtStr]): return DType(priority, itemsize, name, fmt, 1, None) + def new(priority:int, itemsize:int, name:str, fmt:FmtStr|None): return DType(priority, itemsize, name, fmt, 1, None) def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self)) def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.scalar().name]}"+(f".vec({self.count})" if self.count > 1 else "") def __lt__(self, o:DType): return (self.priority, self.itemsize, self.name, self.fmt, self.count) < (o.priority, o.itemsize, o.name, o.fmt, o.count) @@ -167,7 +167,7 @@ if (env_default_float := getenv("DEFAULT_FLOAT", "")): dtypes.default_float = getattr(dtypes, env_default_float.lower()) assert dtypes.is_float(dtypes.default_float), f"{env_default_float} is not a float dtype" -DTypeLike = Union[str, DType] +DTypeLike = str|DType def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower()) # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html @@ -279,7 +279,7 @@ truncate: dict[DType, Callable] = {dtypes.bool: bool, # numpy and torch dtype interop -def _to_np_dtype(dtype:DType) -> Optional[type]: +def _to_np_dtype(dtype:DType) -> type|None: import numpy as np return np.dtype(dtype.fmt).type if dtype.fmt is not None else None def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] # noqa: F821 @@ -287,7 +287,7 @@ def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] # return dtypes.fields()[np.dtype(npdtype).name] @functools.cache -def _to_torch_dtype(dtype:DType) -> Optional['torch.dtype']: # type: ignore [name-defined] # noqa: F821 +def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None: # type: ignore [name-defined] # noqa: F821 import numpy as np, torch # NOTE: torch doesn't expose this mapping with a stable API try: return torch.from_numpy(np.array([], dtype=_to_np_dtype(dtype))).dtype diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 9631718fcf..84bd580f14 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -1,4 +1,4 @@ -from typing import TypeVar, Generic, Callable, Union, cast, Optional, Any +from typing import TypeVar, Generic, Callable, cast, Any import functools, collections from tinygrad.tensor import Tensor from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, colored, JIT, JIT_BATCH_SIZE, dedup, partition, unwrap @@ -21,7 +21,7 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer] # This allows the accelerator to run some batches while subsequent graphs are still being updated. graphed_jit_cache: list[ExecItem] = [] current_batch: list[ExecItem] = [] - current_device: Optional[Compiled] = None + current_device: Compiled|None = None def flush_batch(): nonlocal current_batch, current_device, max_batch_size @@ -31,7 +31,7 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer] graph_runner = current_device.graph(current_batch, input_rawbuffers, var_vals) # clear jit inputs to allow their memory to be freed/reused for (j,i) in graph_runner.input_replace.keys(): graph_runner.jit_cache[j].bufs[i] = None - graphed_jit_cache.append(ExecItem(graph_runner, cast(list[Optional[Buffer]], input_rawbuffers))) + graphed_jit_cache.append(ExecItem(graph_runner, cast(list[Buffer|None], input_rawbuffers))) max_batch_size *= 2 if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels on device {current_device}") except GraphException as e: @@ -76,7 +76,7 @@ class GraphRunner(Runner): self.jit_cache = jit_cache # NOTE: this is not used, but you have to keep these objects alive for the Graph self.input_replace:dict[tuple[int, int], int] = get_input_replace(jit_cache, input_rawbuffers) self.var_vals_replace:dict[int, list[tuple[int, int]]] = {} - self.launch_dims_replace:dict[int, tuple[Optional[int], Optional[int]]] = {} + self.launch_dims_replace:dict[int, tuple[int|None, int|None]] = {} self.launch_dims_base:dict[int, tuple[tuple[int, ...], tuple[int, ...]]] = {} def is_sym_dim(dim) -> bool: return not all(isinstance(d, (int, float)) for d in dim) @@ -149,7 +149,7 @@ class CapturedJit(Generic[ReturnType]): jit_cache: list[ExecItem] input_replace: dict[tuple[int, int], int] extra_view_inputs: list[tuple[int, int, str, int, DType]] - expected_names: list[Union[int, str]] + expected_names: list[int|str] expected_st_vars_dtype_device: list[tuple[ShapeTracker, tuple[Variable, ...], DType, str]] def __reduce__(self): @@ -222,10 +222,10 @@ def _prepare_jit_inputs(args, kwargs): return input_buffers, var_vals, names, st_vars_dtype_device class TinyJit(Generic[ReturnType]): - def __init__(self, fxn:Optional[Callable[..., ReturnType]], captured:Optional[CapturedJit]=None, prune=False, optimize=False): + def __init__(self, fxn:Callable[..., ReturnType]|None, captured:CapturedJit|None=None, prune=False, optimize=False): assert fxn or captured, "need either a function or a CapturedJit" self.fxn = fxn - self.captured: Optional[CapturedJit] = captured + self.captured: CapturedJit|None = captured self.cnt: int = 2 if self.fxn is None else 0 self.prune = prune self.optimize = optimize diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index bfc003cb5c..52da1698e5 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -1,4 +1,4 @@ -from typing import Optional, cast, Generator +from typing import cast, Generator import time, pprint from dataclasses import dataclass, replace, field from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey @@ -54,13 +54,13 @@ class Runner: self.first_run, self.display_name, self.device, self.estimates = True, display_name, device, estimates @property def dev(self): return Device[self.device] - def exec(self, rawbufs:list[Buffer], var_vals:Optional[dict[Variable, int]]=None) -> Optional[float]: + def exec(self, rawbufs:list[Buffer], var_vals:dict[Variable, int]|None=None) -> float|None: return self(rawbufs, {} if var_vals is None else var_vals) - def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False) -> Optional[float]: + def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False) -> float|None: raise NotImplementedError("override this") class CompiledRunner(Runner): - def __init__(self, p:ProgramSpec, precompiled:Optional[bytes]=None, prg=None): + def __init__(self, p:ProgramSpec, precompiled:bytes|None=None, prg=None): if DEBUG >= 4: print(p.src) self.p:ProgramSpec = p self.lib:bytes = precompiled if precompiled is not None else Device[p.device].compiler.compile_cached(p.src) @@ -70,7 +70,7 @@ class CompiledRunner(Runner): def __reduce__(self): return self.__class__, (self.p, self.lib) - def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False) -> Optional[float]: + def __call__(self, rawbufs:list[Buffer], var_vals:dict[Variable, int], wait=False) -> float|None: global_size, local_size = self.p.launch_dims(var_vals) if global_size is not None and local_size is None and all_int(self.p.global_size): # type: ignore[arg-type] # TODO: this is copied from get_program @@ -140,10 +140,10 @@ def get_runner(device:str, ast:UOp) -> CompiledRunner: @dataclass(frozen=True) class ExecItem: prg: Runner - bufs: list[Optional[Buffer]] - metadata: Optional[tuple[Metadata, ...]] = None + bufs: list[Buffer|None] + metadata: tuple[Metadata, ...]|None = None fixedvars: dict[Variable, int] = field(default_factory=dict) - def run(self, _var_vals:Optional[dict[Variable, int]]=None, wait=False, jit=False, do_update_stats=True) -> Optional[float]: + def run(self, _var_vals:dict[Variable, int]|None=None, wait=False, jit=False, do_update_stats=True) -> float|None: var_vals = self.fixedvars if _var_vals is None else (_var_vals|self.fixedvars) bufs = [cast(Buffer, x) for x in self.bufs] if jit else [cast(Buffer, x).ensure_allocated() for x in self.bufs] et = self.prg(bufs, var_vals, wait=wait or DEBUG >= 2) @@ -188,7 +188,7 @@ def lower_schedule(schedule:list[ScheduleItem]) -> Generator[tuple[ScheduleItem, capturing: list = [] # put classes with an add method in here -def run_schedule(schedule:list[ScheduleItem], var_vals:Optional[dict[Variable, int]]=None, do_update_stats=True): +def run_schedule(schedule:list[ScheduleItem], var_vals:dict[Variable, int]|None=None, do_update_stats=True): for si, ei in lower_schedule(schedule): if len(capturing) and CAPTURING: capturing[0].add(ei) if VALIDATE_WITH_CPU and si.ast.op is Ops.SINK: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 5374a12f80..7d47014922 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -2,12 +2,12 @@ from __future__ import annotations import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal from dataclasses import dataclass -from typing import Union, ClassVar, Optional, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator +from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator T = TypeVar("T") U = TypeVar("U") # NOTE: it returns int 1 if x is empty regardless of the type of x -def prod(x:Iterable[T]) -> Union[T,int]: return functools.reduce(operator.mul, x, 1) +def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1) # NOTE: helpers is not allowed to import from anything else in tinygrad OSX = platform.system() == "Darwin" @@ -23,14 +23,14 @@ def argfix(*x): return tuple(x[0]) return x def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) # https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python -def all_same(items:Union[tuple[T, ...], list[T]]): return all(x == items[0] for x in items) +def all_same(items:tuple[T, ...]|list[T]): return all(x == items[0] for x in items) def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t) -def colored(st, color:Optional[str], background=False): return f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m" if color is not None else st # replace the termcolor library with one line # noqa: E501 +def colored(st, color:str|None, background=False): return f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m" if color is not None else st # replace the termcolor library with one line # noqa: E501 def colorize_float(x: float): return colored(f"{x:7.2f}x", 'green' if x < 0.75 else 'red' if x > 1.15 else 'yellow') def time_to_str(t:float, w=8) -> str: return next((f"{t * d:{w}.2f}{pr}" for d,pr in [(1, "s "),(1e3, "ms")] if t > 10/d), f"{t * 1e6:{w}.2f}us") def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s) def ansilen(s:str): return len(ansistrip(s)) -def make_tuple(x:Union[int, Sequence[int]], cnt:int) -> tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else tuple(x) +def make_tuple(x:int|Sequence[int], cnt:int) -> tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else tuple(x) def flatten(l:Iterable[Iterable[T]]): return [item for sublist in l for item in sublist] def fully_flatten(l): if hasattr(l, "__len__") and hasattr(l, "__getitem__") and not isinstance(l, str): @@ -62,7 +62,7 @@ def partition(itr:Iterable[T], fxn:Callable[[T],bool]) -> tuple[list[T], list[T] ret:tuple[list[T], list[T]] = ([], []) for s in itr: (ret[0] if fxn(s) else ret[1]).append(s) return ret -def unwrap(x:Optional[T]) -> T: +def unwrap(x:T|None) -> T: assert x is not None return x def get_single_element(x:Sequence[T]) -> T: @@ -227,7 +227,7 @@ def diskcache_clear(): drop_tables = cur.execute("SELECT 'DROP TABLE IF EXISTS ' || quote(name) || ';' FROM sqlite_master WHERE type = 'table';").fetchall() cur.executescript("\n".join([s[0] for s in drop_tables] + ["VACUUM;"])) -def diskcache_get(table:str, key:Union[dict, str, int]) -> Any: +def diskcache_get(table:str, key:dict|str|int) -> Any: if CACHELEVEL < 1: return None if isinstance(key, (str,int)): key = {"key": key} cur = db_connection().cursor() @@ -239,7 +239,7 @@ def diskcache_get(table:str, key:Union[dict, str, int]) -> Any: return None _db_tables = set() -def diskcache_put(table:str, key:Union[dict, str, int], val:Any, prepickled=False): +def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False): if CACHELEVEL < 1: return val if isinstance(key, (str,int)): key = {"key": key} conn = db_connection() @@ -274,7 +274,7 @@ def _ensure_downloads_dir() -> pathlib.Path: return downloads_dir return pathlib.Path(cache_dir) / "downloads" -def fetch(url:str, name:Optional[Union[pathlib.Path, str]]=None, subdir:Optional[str]=None, gunzip:bool=False, +def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False, allow_caching=not getenv("DISABLE_HTTP_CACHE")) -> pathlib.Path: if url.startswith(("/", ".")): return pathlib.Path(url) if name is not None and (isinstance(name, pathlib.Path) or '/' in name): fp = pathlib.Path(name) @@ -344,7 +344,7 @@ def flat_mv(mv:memoryview): return mv if len(mv) == 0 else mv.cast("B", shape=(m class tqdm(Generic[T]): def __init__(self, iterable:Iterable[T]|None=None, desc:str='', disable:bool=False, - unit:str='it', unit_scale=False, total:Optional[int]=None, rate:int=100): + unit:str='it', unit_scale=False, total:int|None=None, rate:int=100): self.iterable, self.disable, self.unit, self.unit_scale, self.rate = iterable, disable, unit, unit_scale, rate self.st, self.i, self.n, self.skip, self.t = time.perf_counter(), -1, 0, 1, getattr(iterable, "__len__", lambda:0)() if total is None else total self.set_description(desc) diff --git a/tinygrad/nn/state.py b/tinygrad/nn/state.py index 6b56a45b6c..aa8c1b1a31 100644 --- a/tinygrad/nn/state.py +++ b/tinygrad/nn/state.py @@ -1,6 +1,6 @@ import json, pathlib, zipfile, pickle, tarfile, struct, functools, io from collections import OrderedDict -from typing import Union, Optional, Any, Callable, BinaryIO, Iterable +from typing import Any, Callable, BinaryIO, Iterable from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes from tinygrad.helpers import prod, argsort, DEBUG, Timing, CI, unwrap, GlobalCounters, tqdm, round_up, T @@ -35,9 +35,9 @@ safe_dtypes = {"BOOL":dtypes.bool, "I8":dtypes.int8, "U8":dtypes.uint8, "I16":dt "I64":dtypes.int64, "U64":dtypes.uint64, "F16":dtypes.float16, "BF16":dtypes.bfloat16, "F32":dtypes.float32, "F64":dtypes.float64} inverse_safe_dtypes = {v:k for k,v in safe_dtypes.items()} -def accept_filename(func: Callable[[Tensor], T]) -> Callable[[Union[Tensor, str, pathlib.Path]], T]: +def accept_filename(func: Callable[[Tensor], T]) -> Callable[[Tensor|str|pathlib.Path], T]: @functools.wraps(func) - def wrapper(fn: Union[Tensor, str, pathlib.Path]) -> T: return func(Tensor(pathlib.Path(fn)) if not isinstance(fn, Tensor) else fn) + def wrapper(fn: Tensor|str|pathlib.Path) -> T: return func(Tensor(pathlib.Path(fn)) if not isinstance(fn, Tensor) else fn) return wrapper @accept_filename @@ -48,7 +48,7 @@ def safe_load_metadata(t:Tensor) -> tuple[Tensor, int, dict[str, Any]]: data_start = int.from_bytes(t[0:8].data(), "little") + 8 return t, data_start, json.loads(t[8:data_start].data().tobytes()) -def safe_load(fn:Union[Tensor, str, pathlib.Path]) -> dict[str, Tensor]: +def safe_load(fn:Tensor|str|pathlib.Path) -> dict[str, Tensor]: """ Loads a .safetensor file, returning the `state_dict`. @@ -61,7 +61,7 @@ def safe_load(fn:Union[Tensor, str, pathlib.Path]) -> dict[str, Tensor]: return { k: data[v['data_offsets'][0]:v['data_offsets'][1]].bitcast(safe_dtypes[v['dtype']]).reshape(v['shape']) for k, v in metadata.items() if k != "__metadata__" } -def safe_save(tensors:dict[str, Tensor], fn:str, metadata:Optional[dict[str, Any]]=None): +def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=None): """ Saves a `state_dict` to disk in a .safetensor file with optional metadata. @@ -193,8 +193,8 @@ def torch_load(t:Tensor) -> dict[str, Tensor]: state_dict = nn.state.torch_load("test.pth") ``` """ - offsets: dict[Union[str, int], int] = {} - lens: dict[Union[str, int], int] = {} + offsets: dict[str|int, int] = {} + lens: dict[str|int, int] = {} def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad=None, backward_hooks=None, metadata=None): #print(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata) lens[storage[2]] = storage[4] * storage[1].itemsize diff --git a/tinygrad/opt/kernel.py b/tinygrad/opt/kernel.py index cde589afca..cacc8a2824 100644 --- a/tinygrad/opt/kernel.py +++ b/tinygrad/opt/kernel.py @@ -2,7 +2,7 @@ from __future__ import annotations import itertools, functools, math from dataclasses import dataclass from collections import defaultdict -from typing import Optional, cast, Final, Callable, Sequence +from typing import cast, Final, Callable, Sequence from enum import Enum, auto from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, smax, AxisType @@ -24,8 +24,8 @@ class OptOps(Enum): @dataclass(frozen=True, order=True) class Opt: op: OptOps - axis: Optional[int] = None - arg: Optional[int | tuple] = None + axis: int|None = None + arg: int|tuple|None = None def __repr__(self): return f"Opt(op={self.op}, axis={self.axis}, arg={self.arg})" axis_letters = {AxisType.GLOBAL: "g", AxisType.LOCAL: "l", AxisType.LOOP: "L", AxisType.UPCAST: "u", @@ -50,7 +50,7 @@ class TensorCoreOptions: self.axes, self.axes_exist = tuple(axes), tuple(axes_exist) class Kernel: - def __init__(self, ast:UOp, opts:Optional[Renderer]=None): + def __init__(self, ast:UOp, opts:Renderer|None=None): assert ast.op is Ops.SINK, ast.op self.ast = ast @@ -76,8 +76,8 @@ class Kernel: self.sts.append(ShapeTracker.from_shape(tuple([smax(*s) for s in zip(*[x.shape for x in self.sts])]), (0,)*len(self.sts[0].shape))) # parameters for optimization - self.tensor_core: Optional[TensorCore] = None - self.tensor_core_opts: Optional[TensorCoreOptions] = None + self.tensor_core: TensorCore|None = None + self.tensor_core_opts: TensorCoreOptions|None = None self.use_tensor_cores: int = 0 self.applied_opts: list[Opt] = [] self.dont_use_locals = False @@ -144,7 +144,7 @@ class Kernel: assert len(self.axis_types) == self.shape_len, "colors size mismatch" return [axis_colors[x] if not self.dont_use_locals or not x == AxisType.GLOBAL else "BLUE" for x in self.axis_types] - def colored_shape(self, pad:Optional[int]=None, dense=False) -> str: + def colored_shape(self, pad:int|None=None, dense=False) -> str: shape_strs = [(s if dense else f"{s:4d}") if isinstance(s, int) else s.render() for s in self.full_shape] ret = ' '.join(colored(s, color) for s,color in zip(shape_strs, self.colors())) if pad: ret += ' '*(pad-ansilen(ret)) @@ -340,14 +340,14 @@ class Kernel: # **** kernel outputs, mostly tensor cores **** - def _create_tc_opts(self, reduceop:UOp, tc:TensorCore, axis:int, opt_level:int) -> Optional[TensorCoreOptions]: + def _create_tc_opts(self, reduceop:UOp, tc:TensorCore, axis:int, opt_level:int) -> TensorCoreOptions|None: has_cast = tc.dtype_in != tc.dtype_out if has_cast and not (reduceop.src[0].op is Ops.CAST and reduceop.src[0].dtype == tc.dtype_out): return None mul_op = reduceop.src[0].src[0] if has_cast else reduceop.src[0] if mul_op.op is not Ops.MUL: return None - def buf_index(src:UOp) -> Optional[int]: + def buf_index(src:UOp) -> int|None: # TODO: apply tc even if the sources are not from LOAD if src.op is Ops.LOAD and src.dtype == tc.dtype_in: return self.bufs.index(src) try: @@ -392,8 +392,7 @@ class Kernel: return True return False - def apply_tensor_cores(self, use_tensor_cores=1, extra_opts:Optional[list[Opt]]=None, axis:int=0, tc_select:Optional[int]=None, - tc_opt:Optional[int]=None) -> bool: + def apply_tensor_cores(self, use_tensor_cores=1, extra_opts:list[Opt]|None=None, axis:int=0, tc_select:int|None=None, tc_opt:int|None=None) -> bool: """ Attempts to apply a tensor core optimization to the kernel. If one exists and applies properly, return true, otherwise return false. Tensor cores are optimized instructions that matrix multiply-accumulate across a wave of threads: D(M, N) = A(M, K) * B(K, N) + C(M, N). @@ -442,7 +441,7 @@ class Kernel: return ret def shape_str_to_axis(self, nms:list[str]) -> tuple[int, ...]: return tuple([self.shape_str().index(x) for x in nms]) - def get_optimized_ast(self, name_override:Optional[str]=None) -> UOp: + def get_optimized_ast(self, name_override:str|None=None) -> UOp: @functools.cache def fixup_ast(op:UOp) -> UOp: ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821 diff --git a/tinygrad/opt/search.py b/tinygrad/opt/search.py index 6d854d734e..2a798f16c1 100644 --- a/tinygrad/opt/search.py +++ b/tinygrad/opt/search.py @@ -1,4 +1,4 @@ -from typing import cast, Optional, Callable +from typing import cast, Callable import itertools, functools, random, math, time, multiprocessing, traceback, signal, atexit from collections import defaultdict from dataclasses import replace @@ -35,8 +35,8 @@ def get_test_global_size(global_size, max_global_size, var_vals): break return test_global_size, input_size / prod(test_global_size) -def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[Variable, int], rawbufs:list[Buffer], early_stop:Optional[float]=None, - allow_test_size:int=True, max_global_size:Optional[int]=65536, clear_l2=False, cnt=3, name="test") -> list[float]: +def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[Variable, int], rawbufs:list[Buffer], early_stop:float|None=None, + allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]: factor = 1 if allow_test_size and p.global_size is not None and max_global_size is not None: global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals) @@ -57,7 +57,7 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[Variable, int], rawbuf class TimeoutException(Exception): pass def timeout_handler(signum, frame): raise TimeoutException() -def _try_compile_linearized_w_idx(x:tuple[int,Kernel], compiler:Compiler) -> tuple[int, Optional[tuple[ProgramSpec, bytes, float]]]: +def _try_compile_linearized_w_idx(x:tuple[int,Kernel], compiler:Compiler) -> tuple[int, tuple[ProgramSpec, bytes, float]|None]: if hasattr(signal, "alarm"): signal.signal(getattr(signal, 'SIGALRM'), timeout_handler) # set timeout @@ -97,7 +97,7 @@ def bufs_from_lin(lin:Kernel, allocate:bool=True) -> list[Buffer]: if x.src[0].base.op is Ops.DEFINE_GLOBAL: bufsts[x.src[0].base.arg].append(x) # TODO: Nones are staying in here if buffers are optimized out! # TODO: add a test for this - rawbufs: list[Optional[Buffer]] = [None]*(max(bufsts)+1) + rawbufs: list[Buffer|None] = [None]*(max(bufsts)+1) for k,lx in bufsts.items(): buf_size = prod(dtype.shape) if isinstance(dtype:=lx[0].src[0].dtype, ImageDType) else max(y.st_arg.real_size() for y in lx) assert isinstance(dtype, (PtrDType, ImageDType)) diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index d58de2da79..c7bc893414 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -1,5 +1,5 @@ from __future__ import annotations -from typing import Optional, Callable, cast, TYPE_CHECKING +from typing import Callable, cast, TYPE_CHECKING import functools, itertools from dataclasses import dataclass, field, replace from tinygrad.helpers import to_function_name, dedup, prod @@ -52,11 +52,11 @@ class ProgramSpec: src:str device:str ast:UOp # save the base ast (this is method cache key) - uops:Optional[list[UOp]]=None + uops:list[UOp]|None=None # filled in from uops (if we have uops) - global_size:Optional[list[int]]=None - local_size:Optional[list[int]]=None + global_size:list[int]|None=None + local_size:list[int]|None=None vars:list[Variable]=field(default_factory=list) globals:list[int]=field(default_factory=list) outs:list[int]=field(default_factory=list) @@ -113,12 +113,12 @@ class Renderer: has_local: bool = True has_shared: bool = True # NOTE: these two should be in (x,y,z) order to match the max_sizes argument in get_grouped_dims - global_max: Optional[tuple[int, ...]] = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now - local_max: Optional[tuple[int, ...]] = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now + global_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now + local_max: tuple[int, ...]|None = (0x8FFFFFFF,) * (3) # TODO: Ops.SPECIAL int32 indexes right now shared_max: int = 32768 tensor_cores: list[TensorCore] = [] - pre_matcher: Optional[PatternMatcher] = None - extra_matcher: Optional[PatternMatcher] = None + pre_matcher: PatternMatcher|None = None + extra_matcher: PatternMatcher|None = None code_for_op: dict[Ops, Callable] = {} def __reduce__(self): return self.__class__, () diff --git a/tinygrad/runtime/ops_disk.py b/tinygrad/runtime/ops_disk.py index 1cfd5d4e6c..d71bbe3609 100644 --- a/tinygrad/runtime/ops_disk.py +++ b/tinygrad/runtime/ops_disk.py @@ -1,5 +1,5 @@ import os, sys, mmap, io, ctypes, ctypes.util, contextlib -from typing import Optional, Generator, Callable +from typing import Generator, Callable from tinygrad.helpers import OSX, round_up from tinygrad.device import Compiled, Allocator with contextlib.suppress(ImportError): @@ -12,8 +12,8 @@ class DiskDevice(Compiled): def __init__(self, device:str): if not DiskDevice._tried_io_uring_init: self._iouring_setup() - self.size: Optional[int] = None - self.fd: Optional[int] = None + self.size: int|None = None + self.fd: int|None = None self.count = 0 super().__init__(device, DiskAllocator(self), None, None, None) def _might_open(self, size:int): diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py index 0a8304d891..3b9b1d6965 100644 --- a/tinygrad/runtime/ops_gpu.py +++ b/tinygrad/runtime/ops_gpu.py @@ -1,5 +1,5 @@ from __future__ import annotations -from typing import Optional, cast +from typing import cast import ctypes, functools, hashlib from tinygrad.runtime.autogen import opencl as cl from tinygrad.helpers import init_c_var, to_char_p_p, from_mv, OSX, DEBUG, getenv, mv_address @@ -46,12 +46,14 @@ class CLProgram: try: check(cl.clReleaseProgram(self.program)) except (TypeError, AttributeError): pass - def __call__(self, *bufs:tuple[ctypes._CData, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:Optional[tuple[int,int,int]]=None, vals:tuple[int, ...]=(), wait=False) -> Optional[float]: # noqa: E501 + def __call__(self, *bufs:tuple[ctypes._CData, BufferSpec], global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None, + vals:tuple[int, ...]=(), wait=False) -> float|None: for i,(b,_) in enumerate(bufs): cl.clSetKernelArg(self.kernel, i, ctypes.sizeof(b), ctypes.byref(b)) for i,v in enumerate(vals,start=len(bufs)): cl.clSetKernelArg(self.kernel, i, 4, ctypes.byref(ctypes.c_int32(v))) if local_size is not None: global_size = cast(tuple[int,int,int], tuple(int(g*l) for g,l in zip(global_size, local_size))) event = cl.cl_event() if wait else None - check(cl.clEnqueueNDRangeKernel(self.dev.queue, self.kernel, len(global_size), None, (ctypes.c_size_t * len(global_size))(*global_size), (ctypes.c_size_t * len(local_size))(*local_size) if local_size else None, 0, None, event)) # noqa: E501 + check(cl.clEnqueueNDRangeKernel(self.dev.queue, self.kernel, len(global_size), None, (ctypes.c_size_t * len(global_size))(*global_size), + (ctypes.c_size_t * len(local_size))(*local_size) if local_size else None, 0, None, event)) if wait: assert event is not None check(cl.clWaitForEvents(1, event)) diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index c9839208f1..abf123b03c 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -1,5 +1,5 @@ import os, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform -from typing import Any, Union, cast +from typing import Any, cast from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE, ProfileRangeEvent, cpu_profile from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, ProfileDeviceEvent from tinygrad.renderer.cstyle import MetalRenderer @@ -111,7 +111,7 @@ class MetalCompiler(Compiler): super().__init__("compile_metal_direct") def __reduce__(self): return (MetalCompiler,()) # force pickle to create new instance for each multiprocessing fork def compile(self, src:str) -> bytes: - ret: Union[Exception, bytes] = CompileError("MTLCodeGenServiceBuildRequest returned without calling the callback") + ret: Exception|bytes = CompileError("MTLCodeGenServiceBuildRequest returned without calling the callback") @ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_int32, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_char_p) def callback(blockptr, error, dataPtr, dataLen, errorMessage): nonlocal ret diff --git a/tinygrad/runtime/ops_nv.py b/tinygrad/runtime/ops_nv.py index fdd911ebd9..d477724ea2 100644 --- a/tinygrad/runtime/ops_nv.py +++ b/tinygrad/runtime/ops_nv.py @@ -1,7 +1,7 @@ from __future__ import annotations import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, weakref assert sys.platform != 'win32' -from typing import cast, Union, ClassVar +from typing import cast, 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 @@ -299,7 +299,7 @@ class NVKIface: root = None fd_ctl: FileIOInterface fd_uvm: FileIOInterface - gpus_info: Union[list, ctypes.Array] = [] + gpus_info: list|ctypes.Array = [] # TODO: Need a proper allocator for va addresses # 0x1000000000 - 0x2000000000, reserved for system/cpu mappings diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 229697ec3f..99521edcae 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -2,7 +2,7 @@ # a python uops emulator # works to test the tensor cores, and all the uops in general # this is the (living) definition of uops -from typing import Optional, Any, TYPE_CHECKING +from typing import Any, TYPE_CHECKING 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 @@ -26,7 +26,7 @@ def _store(m, i, v): class PythonProgram: def __init__(self, name:str, lib:bytes): - self.uops: list[tuple[Ops, Optional[DType], list[int], Any]] = pickle.loads(lib) + self.uops: list[tuple[Ops, DType|None, list[int], Any]] = pickle.loads(lib) def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False): st = time.perf_counter() warp = list(itertools.product(*[range(x) for x in local_size[::-1]])) diff --git a/tinygrad/runtime/ops_remote.py b/tinygrad/runtime/ops_remote.py index 71f7dc22f6..ad4b286e4e 100644 --- a/tinygrad/runtime/ops_remote.py +++ b/tinygrad/runtime/ops_remote.py @@ -5,7 +5,7 @@ # it should be a secure (example: no use of pickle) boundary. HTTP is used for RPC from __future__ import annotations -from typing import Callable, Iterator, Optional, Any, cast +from typing import Callable, Iterator, Any, cast from collections import defaultdict from dataclasses import dataclass, field, replace import multiprocessing, threading, functools, itertools, asyncio, http, http.client, hashlib, time, os, binascii, struct, ast, contextlib, weakref @@ -77,7 +77,7 @@ class ProgramFree(RemoteRequest): name: str; datahash: str # noqa: E702 @dataclass(frozen=True) class ProgramExec(RemoteRequest): name: str; datahash: str; bufs: tuple[int, ...]; vals: tuple[int, ...] # noqa: E702 - global_size: Optional[tuple[int, ...]]; local_size: Optional[tuple[int, ...]]; wait: bool # noqa: E702 + global_size: tuple[int, ...]|None; local_size: tuple[int, ...]|None; wait: bool # noqa: E702 @dataclass(frozen=True) class GraphComputeItem: diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 5b5ca7a96b..bdc83b802f 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass import functools -from typing import Optional, Callable +from typing import Callable from tinygrad.helpers import merge_dicts, getenv from tinygrad.shape.view import View, strides_for_shape, unravel from tinygrad.dtype import dtypes @@ -23,7 +23,7 @@ def handle_upcast(u: UOp) -> UOp|None: pm_upcast = PatternMatcher([(UPat(GroupOp.ALU, dtype=dtypes.int, name="u"), handle_upcast),]) @functools.cache -def views_to_indexed_uops(views: tuple[View, ...], _idxs:Optional[tuple[UOp, ...]]=None) -> tuple[UOp, UOp]: +def views_to_indexed_uops(views: tuple[View, ...], _idxs:tuple[UOp, ...]|None=None) -> tuple[UOp, UOp]: idx, valid = views[-1].to_indexed_uops(_idxs) for view in reversed(views[0:-1]): view = view.minify() @@ -38,10 +38,10 @@ def views_to_indexed_uops(views: tuple[View, ...], _idxs:Optional[tuple[UOp, ... return graph_rewrite(UOp.sink(idx, valid), symbolic_flat+pm_upcast, name="indexing sym @ 2").src @functools.cache -def views_to_real_strides(views: tuple[View, ...], ignore_valid=False) -> tuple[Optional[sint], ...]: +def views_to_real_strides(views: tuple[View, ...], ignore_valid=False) -> tuple[sint|None, ...]: # NOTE: if a stride is not always valid, it will be None if len(views) == 1 and views[-1].mask is None: return views[-1].strides - ret: list[Optional[sint]] = [None] * len(views[-1].shape) + ret: list[sint|None] = [None] * len(views[-1].shape) idx, valid = views_to_indexed_uops(views) for c in split_uop(idx, Ops.ADD): if c.op is Ops.RANGE: ret[c.arg] = 1 @@ -62,7 +62,7 @@ class ShapeTracker: for v in st.views: ret = ShapeTracker(ret.views + (v,)).simplify() # one view at a time = better simplification return ret - def invert(self, out_shape:tuple[sint, ...]) -> Optional[ShapeTracker]: + def invert(self, out_shape:tuple[sint, ...]) -> ShapeTracker|None: inverted_views:list[View] = [] for v,s in zip(self.views[::-1], [x.shape for x in self.views[::-1][1:]]+[out_shape]): if (inverted:= v.invert(s)) is None: return None @@ -87,7 +87,7 @@ class ShapeTracker: def reduce(self, axis:tuple[int, ...]) -> tuple[sint, ...]: return tuple(1 if i in axis else s for i,s in enumerate(self.shape)) def to_uop(self) -> UOp: return UOp(Ops.VIEW, dtypes.void, (), self) - def to_indexed_uops(self, _idxs:Optional[list[UOp]|tuple[UOp, ...]]=None) -> tuple[UOp, UOp]: + def to_indexed_uops(self, _idxs:list[UOp]|tuple[UOp, ...]|None=None) -> tuple[UOp, UOp]: return views_to_indexed_uops(self.views, tuple(_idxs) if _idxs is not None else None) # upper bound on buffer size required to fit this shapetracker @@ -109,7 +109,7 @@ class ShapeTracker: return ShapeTracker(tuple(unbound_views)), merge_dicts(var_vals) def substitute(self, dvars:dict[UOp, UOp]): return ShapeTracker(tuple(x.substitute(dvars) for x in self.views)) - def real_strides(self, ignore_valid=False) -> tuple[Optional[sint], ...]: + def real_strides(self, ignore_valid=False) -> tuple[sint|None, ...]: with Context(TRACK_MATCH_STATS=0): return views_to_real_strides(self.views, ignore_valid) def unit_stride_axes(self, ignore_valid=False) -> list[int]: return [i for i,st in enumerate(self.real_strides(ignore_valid)) if st == 1] diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index df80eae0aa..0475fc6506 100644 --- a/tinygrad/shape/view.py +++ b/tinygrad/shape/view.py @@ -1,7 +1,7 @@ from __future__ import annotations import functools, operator, itertools from dataclasses import dataclass -from typing import Optional, cast, Sequence +from typing import cast, Sequence from tinygrad.dtype import dtypes from tinygrad.uop.ops import resolve, UOp, Variable, sint, sym_infer, smax, smin, sint_to_uop, Ops, ssimplify from tinygrad.helpers import prod, all_int, argsort, flatten, ceildiv @@ -42,7 +42,7 @@ def strides_for_shape(shape:tuple[sint, ...]) -> tuple[sint, ...]: return canonicalize_strides(shape, strides) @functools.cache -def merge_dims(shape:tuple[int, ...], strides:tuple[int, ...], mask:Optional[tuple[tuple[int, int], ...]]=None) -> tuple[tuple[int, int, int], ...]: +def merge_dims(shape:tuple[int, ...], strides:tuple[int, ...], mask:tuple[tuple[int, int], ...]|None=None) -> tuple[tuple[int, int, int], ...]: # merge contiguous sub-parts or zero strided dims # any stride 0, masked from dim=1, or contiguous part is merged into next dim. # stride != 0 to stride == 0 starts a new merging block @@ -64,8 +64,8 @@ def merge_dims(shape:tuple[int, ...], strides:tuple[int, ...], mask:Optional[tup return tuple(ret) @functools.cache -def _reshape_mask(_mask:Optional[tuple[tuple[sint, sint], ...]], old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) \ - -> Optional[tuple[tuple[sint, sint], ...]]: +def _reshape_mask(_mask:tuple[tuple[sint, sint], ...]|None, old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) \ + -> tuple[tuple[sint, sint], ...]|None: """Returns the new mask if reshape is possible, and None if not possible.""" if _mask is None: return tuple((0, s) for s in new_shape) if not all_int(flatten(_mask)): return None @@ -109,10 +109,10 @@ class View: shape:tuple[sint, ...] strides:tuple[sint, ...] offset:sint - mask:Optional[tuple[tuple[sint, sint], ...]] + mask:tuple[tuple[sint, sint], ...]|None contiguous:bool - def to_indexed_uops(self:View, idxs:Optional[Sequence[UOp]]=None, vexpr:UOp=UOp.const(dtypes.bool, True)) -> tuple[UOp, UOp]: + def to_indexed_uops(self:View, idxs:Sequence[UOp]|None=None, vexpr:UOp=UOp.const(dtypes.bool, True)) -> tuple[UOp, UOp]: """(idx, valid)""" if idxs is None: idxs = [UOp.range(dtypes.int, s, i) for i,s in enumerate(self.shape)] iexpr = sint_to_uop(self.offset) @@ -131,7 +131,7 @@ class View: @staticmethod @functools.cache - def create(shape:tuple[sint, ...], strides:Optional[tuple[sint, ...]]=None, offset:sint=0, mask:Optional[tuple[tuple[sint, sint], ...]]=None): + def create(shape:tuple[sint, ...], strides:tuple[sint, ...]|None=None, offset:sint=0, mask:tuple[tuple[sint, sint], ...]|None=None): # TODO: resolve shouldn't be needed here if not all(resolve(s >= 0) for s in shape): raise ValueError(f"Trying to create View with negative dimension: {shape=}") strides = canonicalize_strides(shape, strides) if strides else strides_for_shape(shape) @@ -177,7 +177,7 @@ class View: return View.create(new_shape, new_strides, new_offset, new_mask) @functools.cache # pylint: disable=method-cache-max-size-none - def __add__(self, vm1:View) -> Optional[View]: + def __add__(self, vm1:View) -> View|None: vm2 = self if vm2.contiguous or vm1.size() == 0: return vm1 if vm1.contiguous and vm1.shape == vm2.shape: return vm2 @@ -244,7 +244,7 @@ class View: return View.create(vm1.shape, tuple(strides), ssimplify(sum(o * s for o, s in zip(origin, vm2.strides)) + vm2.offset)) @functools.cache # pylint: disable=method-cache-max-size-none - def invert(self, out_shape:tuple[sint, ...]) -> Optional[View]: + def invert(self, out_shape:tuple[sint, ...]) -> View|None: ret = View.create(self.shape) if self.mask: ret = ret.shrink(self.mask) ret = ret.flip(tuple(x < 0 for x in self.strides)).permute(argsort(tuple(-x if x > 0 else x for x in self.strides))) @@ -306,7 +306,7 @@ class View: return View.create(self.shape, tuple(-z if f else z for z,f in zip(self.strides, arg)), self.offset+offset, mask) @functools.cache # pylint: disable=method-cache-max-size-none - def reshape(self, new_shape: tuple[sint, ...]) -> Optional[View]: + def reshape(self, new_shape: tuple[sint, ...]) -> View|None: if self.shape == new_shape: return self if not all(x >= 0 for x in new_shape): raise ValueError(f"shape can't contain negative numbers {new_shape}") diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index f82436a8fd..6361a2775e 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -2,7 +2,7 @@ from __future__ import annotations import time, math, itertools, functools, struct, sys, inspect, pathlib, string, hashlib, weakref, contextvars from contextlib import ContextDecorator -from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, SupportsIndex, ParamSpec, TypeVar, Optional +from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, SupportsIndex, ParamSpec, TypeVar from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate from tinygrad.dtype import _from_np_dtype, _to_np_dtype from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup @@ -47,7 +47,7 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str|None=None) -> Non # **** Tensor helper functions **** # this tracks the tensor.py METADATA -_METADATA: contextvars.ContextVar[Optional[Metadata]] = contextvars.ContextVar("_METADATA", default=None) +_METADATA: contextvars.ContextVar[Metadata|None] = contextvars.ContextVar("_METADATA", default=None) def _fromnp(x: 'np.ndarray') -> UOp: # type: ignore [name-defined] # noqa: F821 ret = UOp.new_buffer("NPY", x.size, _from_np_dtype(x.dtype)) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 7d99f69962..1119a8c802 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1,5 +1,5 @@ from __future__ import annotations -from typing import Any, Optional, Union, Callable, cast, TYPE_CHECKING, Type, Sequence +from typing import Any, Callable, cast, TYPE_CHECKING, Type, Sequence import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref from dataclasses import dataclass, field from enum import Enum, auto @@ -34,7 +34,7 @@ def smin(*lst): return _suop(argfix(*lst), UOp.minimum, min) def srender(x) -> str: return x.render() if isinstance(x, UOp) else str(x) def ssimplify(uop): return uop.ssimplify() if isinstance(uop, UOp) else uop -def sym_infer(uop: Union[UOp, int], var_vals: dict[UOp, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop +def sym_infer(uop: UOp|int, var_vals: dict[UOp, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop # used for UOp and UPat def pretty_print(x:Any, rep:Callable, srcfn=lambda x: x.src, cache=None, d=0)->str: @@ -180,7 +180,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): 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 + def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret def _eval(self, dtype, expected_type:Type[T]) -> T: assert self.dtype in dtype, f"eval with wrong dtype {self}" vmin, vmax = (simple_self:=self.simplify())._min_max @@ -223,7 +223,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return UOp(Ops.CAST, dtype, (self,)) def cast_vec(self, dtype:DType): return UOp(Ops.CAST, dtype.vec(self.dtype.count), (self,)) def bitcast(self, dtype:DType): return UOp(Ops.BITCAST, dtype, (self,)) - def gep(self, i:Union[tuple[int, ...], int]): + def gep(self, i:tuple[int, ...]|int): if isinstance(i, tuple) and len(i) == 1: return self.gep(i[0]) if isinstance(i, int): # NOTE: these are just shortcuts to not have to create and fold later @@ -290,7 +290,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return tuple(itertools.pairwise(itertools.accumulate([self.src[0].shape[self.axis] for _ in self.device], initial=0))) @functools.cached_property - def axis(self) -> Optional[int]: + def axis(self) -> int|None: if self.op is Ops.MULTI: return self.arg # NOTE: they all have to share an axis, we always choose [-1] if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None @@ -364,7 +364,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): @property def device(self) -> str|tuple[str, ...]: return cast(str|tuple[str, ...], unwrap(self._device)) @functools.cached_property - def _device(self) -> Optional[str|tuple[str, ...]]: + def _device(self) -> str|tuple[str, ...]|None: if self.op is Ops.DEVICE: return self.arg if self.op is Ops.MSELECT: assert isinstance(self.src[0].device, tuple), "mselect must be on tuple device" @@ -402,7 +402,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): buffers[self] = ret return ret @property - def realized(self) -> Optional[Buffer|MultiBuffer]: + def realized(self) -> Buffer|MultiBuffer|None: # NOTE: this is used by the JIT to determine which inputs we capture return self.buffer if self.op in {Ops.BUFFER, Ops.MSTACK} and self.buffer.is_allocated() else None @property @@ -523,7 +523,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): fxn, varnames = self._sym_fxn return fxn(**{k.arg[0]:v for k,v in var_vals.items() if k.arg[0] in varnames}) - def render(self, simplify=True, pm:Optional[PatternMatcher]=None) -> str: + def render(self, simplify=True, pm:PatternMatcher|None=None) -> str: ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm) return ret.arg if ret.op is Ops.NOOP else str(ret) @@ -596,12 +596,12 @@ def printable(loc:tuple[str, int]) -> str: class UPat(MathTrait): __slots__ = ("op", "dtype", "arg", "name", "src") - def __init__(self, op:Optional[Union[Ops, tuple[Ops, ...], set[Ops]]]=None, dtype:Optional[Union[DType, tuple[DType, ...]]]=None, - src:Optional[Union[tuple[UPat, ...], list[UPat], UPat]]=None, arg:Any=None, - name:Optional[str]=None, allow_any_len:bool=False, custom_early_reject:Optional[set[Ops]]=None, location=None): + def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|None=None, + src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None, + name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None): assert op is None or isinstance(op, (Ops, tuple, set)), "op must be Ops or tuple of Ops" - self.op: Optional[tuple[Ops, ...]] = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op) - self.dtype: Optional[tuple[DType, ...]] = (dtype,) if isinstance(dtype, DType) else dtype + self.op: tuple[Ops, ...]|None = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op) + self.dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else dtype self.arg, self.name, self._in_src, self.custom_early_reject = arg, name, src, custom_early_reject self.src: Any = None assert self.name != "ctx", "UPat can't be named ctx" @@ -633,15 +633,15 @@ class UPat(MathTrait): @staticmethod @functools.cache - def var(name:Optional[str]=None, dtype:Optional[Union[DType, tuple[DType, ...]]]=None): return UPat(dtype=dtype, name=name) + def var(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None): return UPat(dtype=dtype, name=name) @staticmethod @functools.cache - def cvar(name:Optional[str]=None, dtype:Optional[DType]=None, vec=True): return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name) + def cvar(name:str|None=None, dtype:DType|None=None, vec=True): return UPat((Ops.CONST,Ops.VCONST) if vec else Ops.CONST, dtype, name=name) @staticmethod - def const(dtype:Optional[Union[DType, tuple[DType, ...]]], b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b) + def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b) # copied from UOp - def index(self, idx:UPat, valid:Optional[UPat]=None): return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx)) + def index(self, idx:UPat, valid:UPat|None=None): return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx)) def view(self, st=None, **kwargs): return UPat(Ops.VIEW, self.dtype, (self,), st, **kwargs) def cast(self, dtype=None, **kwargs): return UPat(Ops.CAST, dtype, (self,), **kwargs) def bitcast(self, dtype=None): return UPat(Ops.BITCAST, dtype, (self,)) @@ -764,7 +764,7 @@ def track_uop(u:UOp): VIZ = ContextVar("VIZ", 0) TRACK_MATCH_STATS = ContextVar("TRACK_MATCH_STATS", 2 if VIZ else 0) -match_stats:dict[UPat, list[Union[int, float]]] = dict() +match_stats:dict[UPat, list[int|float]] = dict() @dataclass(frozen=True) class TrackedGraphRewrite: @@ -959,7 +959,7 @@ renderer_infer = PatternMatcher([ # *** what was symbolic.py *** -sint = Union[int, UOp] +sint = int|UOp Variable = UOp -ConstLike = Union[ConstType, Variable, tuple[ConstType, ...]] +ConstLike = ConstType|Variable|tuple[ConstType, ...]