From 9a2ac00f3cf51bd041959292d93410c5dafda280 Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:54:09 -0500 Subject: [PATCH 1/7] Backend availability at module initialization Speeds up convert array on 6 arrays from 2.8 ms --> 1.5 ms --- pykokkos/interface/parallel_dispatch.py | 36 +++++++++++++------------ 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/pykokkos/interface/parallel_dispatch.py b/pykokkos/interface/parallel_dispatch.py index bcb0ef0e..4dcf2d43 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 @@ -282,25 +301,8 @@ 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): From b29d3bb58dd9232677fa970c5130437c37325342 Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:01:54 -0500 Subject: [PATCH 2/7] Remove extra call to is_array Each call to is_array cost ~40\mu s --- pykokkos/interface/views.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 8859a20b..2c4b482d 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -965,10 +965,11 @@ def array( """ # if an array is not a recognized type, try coasting it to a numpy array + is_array_flag: bool = is_array(array) if ( not isinstance(array, np.ndarray) and not np.isscalar(array) - and not is_array(array) + and not is_array_flag ): array = np.asarray(array) @@ -981,7 +982,7 @@ def array( 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") From 4dec3f9b7317518c7b1111e83befc44df1764dbe Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:15:35 -0500 Subject: [PATCH 3/7] Remove redundant flag checking in convert array --- pykokkos/interface/views.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 2c4b482d..0d2bfac2 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -964,13 +964,13 @@ def array( :returns: a PyKokkos View wrapping the array """ - # if an array is not a recognized type, try coasting it to a numpy array + # reused type flags is_array_flag: bool = is_array(array) - if ( - not isinstance(array, np.ndarray) - and not np.isscalar(array) - and not is_array_flag - ): + 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 is_numpy_instance and not is_scalar_flag and not is_array_flag: array = np.asarray(array) # check that the array is contiguous @@ -978,7 +978,7 @@ def array( 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 From 2ef9f25e254dd7dac59bab682b791b5bf804eedc Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Wed, 1 Jul 2026 10:39:33 -0500 Subject: [PATCH 4/7] Replace slow set comparisons with fast is comps --- pykokkos/interface/views.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 0d2bfac2..3d3619e9 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: From bfef3043ba0787063e60004836908e6ad6106965 Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Wed, 1 Jul 2026 14:53:37 -0500 Subject: [PATCH 5/7] Avoid call to is_array when array status known --- pykokkos/interface/parallel_dispatch.py | 6 +++--- pykokkos/interface/views.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/pykokkos/interface/parallel_dispatch.py b/pykokkos/interface/parallel_dispatch.py index 4dcf2d43..5501db5f 100644 --- a/pykokkos/interface/parallel_dispatch.py +++ b/pykokkos/interface/parallel_dispatch.py @@ -340,7 +340,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( @@ -348,9 +348,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 3d3619e9..3d764551 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -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,11 +964,14 @@ 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 - is_array_flag: bool = is_array(array) + 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) From b4536fbf9232b022a6f469b5218e733dde1b1c57 Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Wed, 1 Jul 2026 15:21:38 -0500 Subject: [PATCH 6/7] Cache type hints from in convert arrays The call to inspect signature alone cost ~300 mu s --- pykokkos/interface/parallel_dispatch.py | 43 ++++++++++++++++++------- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/pykokkos/interface/parallel_dispatch.py b/pykokkos/interface/parallel_dispatch.py index 5501db5f..db585f64 100644 --- a/pykokkos/interface/parallel_dispatch.py +++ b/pykokkos/interface/parallel_dispatch.py @@ -291,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 @@ -306,17 +337,7 @@ def convert_arrays(kwargs: Dict[str, Any], workunit: Callable, execution_space) # 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): From 5ff0906080f2ae2929870fba216e857de84c4923 Mon Sep 17 00:00:00 2001 From: Gabriel Kosmacher <73120774+kennykos@users.noreply.github.com> Date: Thu, 2 Jul 2026 08:01:55 -0500 Subject: [PATCH 7/7] Update flags after cast to numpy --- pykokkos/interface/views.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 3d764551..3885147c 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -978,6 +978,9 @@ def array( # if an array is not a recognized type, try coasting it to a numpy 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"]: