diff --git a/pykokkos/interface/parallel_dispatch.py b/pykokkos/interface/parallel_dispatch.py index bcb0ef0e..db585f64 100644 --- a/pykokkos/interface/parallel_dispatch.py +++ b/pykokkos/interface/parallel_dispatch.py @@ -30,6 +30,25 @@ import inspect +# check backend availability +cp_available: bool +torch_available: bool + +try: + import cupy as cp + + cp_available = True +except ImportError: + cp_available = False + +try: + import torch + + torch_available = True +except ImportError: + torch_available = False + + workunit_cache: Dict[int, Callable] = {} # Map PyKokkos BuiltinType to numpy dtypes @@ -272,6 +291,37 @@ def check_workunit(workunit: Any) -> None: raise TypeError(f"ERROR: {workunit} is not a valid workunit") +_type_hints_cache: Dict[int, Tuple[Callable, Dict[str, Any]]] = {} + + +def _get_type_hints(workunit: Callable) -> Dict[str, Any]: + """ + Extract and cache a workunit's parameter type hints. + + Cache by id(workunit). Also keeps a strong reference + to the workunit alongside the cached hints to prevent a stale hit if its + id gets reused after garbage collection. + """ + key = id(workunit) + cached = _type_hints_cache.get(key) + if cached is not None and cached[0] is workunit: + return cached[1] + + type_hints: Dict[str, Any] = {} + try: + sig = inspect.signature(workunit) + type_hints = { + name: param.annotation + for name, param in sig.parameters.items() + if param.annotation != inspect.Parameter.empty + } + except (ValueError, TypeError): + pass + + _type_hints_cache[key] = (workunit, type_hints) + return type_hints + + def convert_arrays(kwargs: Dict[str, Any], workunit: Callable, execution_space) -> None: """ Convert all numpy, cupy and pytorch ndarray objects into pk Views @@ -282,39 +332,12 @@ def convert_arrays(kwargs: Dict[str, Any], workunit: Callable, execution_space) (used to convert arrays to the correct memory space) """ - cp_available: bool - torch_available: bool - memory_space = get_default_memory_space(execution_space) - try: - import cupy as cp - - cp_available = True - except ImportError: - cp_available = False - - try: - import torch - - torch_available = True - except ImportError: - torch_available = False - # Get type hints from workunit if available type_hints = {} if workunit is not None and callable(workunit): - import inspect as insp - - try: - sig = insp.signature(workunit) - type_hints = { - name: param.annotation - for name, param in sig.parameters.items() - if param.annotation != insp.Parameter.empty - } - except (ValueError, TypeError): - pass + type_hints = _get_type_hints(workunit) for k, v in kwargs.items(): if isinstance(v, ViewType) or isinstance(v, np.generic): @@ -338,7 +361,7 @@ def convert_arrays(kwargs: Dict[str, Any], workunit: Callable, execution_space) f"from the {execution_space.value} execution space. " f"Use a pk.View (e.g. pk.View([...], dtype)) or a CuPy array instead." ) - kwargs[k] = array(v, space=memory_space) + kwargs[k] = array(v, space=memory_space, is_array_flag=True) elif cp_available and isinstance(v, cp.ndarray): if execution_space not in DeviceExecutionSpace: raise TypeError( @@ -346,9 +369,9 @@ def convert_arrays(kwargs: Dict[str, Any], workunit: Callable, execution_space) f"from the {execution_space.value} (host) execution space. " f"Convert it to a numpy array or pk.View in host memory first." ) - kwargs[k] = array(v, space=memory_space) + kwargs[k] = array(v, space=memory_space, is_array_flag=True) elif torch_available and torch.is_tensor(v): - kwargs[k] = array(v, space=memory_space) + kwargs[k] = array(v, space=memory_space, is_array_flag=True) elif ( hasattr(v, "__array__") or hasattr(v, "__cuda_array_interface__") diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 8859a20b..3885147c 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -406,7 +406,7 @@ def _init_view( # only allow CudaSpace/HIPSpace view for cupy arrays if ( - space in {MemorySpace.CudaSpace, MemorySpace.HIPSpace} + (space is MemorySpace.CudaSpace or space is MemorySpace.HIPSpace) ) and trait is not trait.Unmanaged: space = MemorySpace.HostSpace @@ -417,9 +417,9 @@ def _init_view( is_cpu: bool = self.space is MemorySpace.HostSpace kokkos_lib: ModuleType = km.get_kokkos_module(is_cpu) - if self.dtype in {DataType.float, pk_float}: + if self.dtype is DataType.float or self.dtype is pk_float: self.dtype = float32 - elif self.dtype in {DataType.double, double}: + elif self.dtype is DataType.double or self.dtype is double: self.dtype = float64 if trait is trait.Unmanaged: if array is not None and array.ndim == 0: @@ -953,7 +953,10 @@ def is_array(array) -> bool: def array( - array, space: Optional[MemorySpace] = None, layout: Optional[Layout] = None + array, + space: Optional[MemorySpace] = None, + layout: Optional[Layout] = None, + is_array_flag: Optional[bool] = None, ) -> ViewType: """ Create a PyKokkos View from a generic array @@ -961,27 +964,34 @@ def array( :param array: the data (array?) of unknown type :param space: an optional argument for memory space (used by from_array) :param layout: an optional argument for layout (used by from_array) + :param is_array_flag: an optional flag to determine if array conforms to + the python array API. If not passed, the flag is set by the `is_array` function. :returns: a PyKokkos View wrapping the array """ + # reused type flags + if is_array_flag is None: + is_array_flag: bool = is_array(array) + is_scalar_flag: bool = np.isscalar(array) + is_numpy_instance: bool = isinstance(array, np.ndarray) + # if an array is not a recognized type, try coasting it to a numpy array - if ( - not isinstance(array, np.ndarray) - and not np.isscalar(array) - and not is_array(array) - ): + if not is_numpy_instance and not is_scalar_flag and not is_array_flag: array = np.asarray(array) + is_array_flag = True + is_numpy_instance = True + is_scalar_flag: bool = np.isscalar(array) # check that the array is contiguous if not array.flags["F_CONTIGUOUS"] and not array.flags["C_CONTIGUOUS"]: raise ValueError(f"numpy array is not contiguous") # if numpy array, use from_numpy() - if isinstance(array, np.ndarray) or np.isscalar(array): + if is_numpy_instance or is_scalar_flag: return from_numpy(array, space, layout) # test if the input array can duck-type to a numpy-like array # and run from_array to preprocess the array to numpy - elif is_array(array): + elif is_array_flag: return from_array(array) else: raise TypeError(f"array of type {type(array)} not supported")