diff --git a/test/unit/test_device.py b/test/unit/test_device.py index e1eaaa1314..9dd9f1e5c0 100644 --- a/test/unit/test_device.py +++ b/test/unit/test_device.py @@ -30,7 +30,7 @@ class TestDevice(unittest.TestCase): @unittest.skipIf(WIN and CI, "skipping windows test") # TODO: subproccess causes memory violation? def test_env_overwrite_default_compiler(self): - expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept RuntimeError: pass" + expect_failure = "\ntry: assert Device[Device.DEFAULT].compiler is None;\nexcept Exception: pass" if Device.DEFAULT == "CPU": from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler, ClangJITCompiler diff --git a/tinygrad/device.py b/tinygrad/device.py index a8186e0b1a..8d958531d0 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Sequence, cast, Generator import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored, CPU_LLVM from tinygrad.helpers import Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup -from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM +from tinygrad.helpers import unwrap_class_type, suppress_finalizing, AMD_LLVM, select_first_inited from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype from tinygrad.renderer import Renderer @@ -291,8 +291,8 @@ class Compiled: if len(enable_comps) > 1: raise RuntimeError(f"{self.device}: multiple compilers set in env {enable_comps}") for _, comp_pair in disable_comps: self.compilers.remove(comp_pair) - try: self.renderer, self.compiler = next(self._get_available_compilers([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers)) - except StopIteration as exc: raise RuntimeError(f"no usable compilers for {self.device}") from exc + self.renderer, self.compiler = select_first_inited([list(enable_comps)[0][1]] if len(enable_comps) == 1 else self.compilers, + f"No compiler for {self.device} is available") if DEBUG >= 1: print(f"{self.device}: using {self.compiler.__class__.__name__}") @@ -300,10 +300,6 @@ class Compiled: compiler_name = f"{unwrap_class_type(c).__name__.upper().removesuffix('COMPILER').removeprefix(devname:=self.device.split(':')[0].upper())}" return f"{devname}_{compiler_name if len(compiler_name) > 0 else unwrap_class_type(c).__name__.upper()}" - def _get_available_compilers(self, compilers) -> Iterator[tuple[Renderer, Compiler]]: - for renderer, compiler in compilers: - with contextlib.suppress(Exception): yield renderer(), compiler() - def synchronize(self): """ Synchronize all pending operations on the device. diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 013d6d53e1..1377bcfb65 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -114,6 +114,13 @@ def suppress_finalizing(func): if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing return wrapper +def select_first_inited(candidates:Sequence[Callable[...,T]|Sequence[Callable[...,T]]], err_msg: str) -> tuple[T,...]|T: + excs = [] + for typ in candidates: + try: return tuple([cast(Callable, t)() for t in typ]) if isinstance(typ, Sequence) else cast(Callable, typ)() + except Exception as e: excs.append(e) + raise ExceptionGroup(err_msg, excs) + def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's') diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index ace5c77955..4341625d57 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -1,9 +1,9 @@ from __future__ import annotations from typing import cast, Callable, Type, TypeVar, Generic, Any, Sequence -import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections +import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools try: import fcntl # windows misses that except ImportError: fcntl = None #type:ignore[assignment] -from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent +from tinygrad.helpers import PROFILE, getenv, to_mv, ProfileRangeEvent, select_first_inited from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, CompilerPairT from tinygrad.uop.ops import sym_infer, sint, UOp from tinygrad.runtime.autogen import libc @@ -438,12 +438,9 @@ class HCQCompiled(Compiled, Generic[SignalType]): return buf, realloced def _select_iface(self, *ifaces:Type): - excs = [] if val:=getenv(f'{type(self).__name__[:-6].upper()}_IFACE', ""): ifaces = tuple(x for x in ifaces if x.__name__.startswith(val.upper())) - for iface_t in ifaces: - try: return iface_t(self, self.device_id) - except Exception as e: excs.append(e) - raise ExceptionGroup(f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available", excs) + return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in ifaces], + f"No interface for {type(self).__name__[:-6]}:{self.device_id} is available") def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"