core: optimize convert arrays

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,17 +361,17 @@ 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(
                     f"Argument '{k}' is a CuPy array, which cannot be accessed "
                     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__")

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,35 +953,45 @@ 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
 
     :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")