[maintenance] lazy load dpnp.tensor/dpnp and prepare for array_api lazy importing (#2509)

* starting point

* first cut

* rename

* fix various testing imports

* don't get ahead of my skis

* attempt to further move things apart

* remove get_unique_values_with_dpep

* remove actually

* Update _array_api.py

* try to fix

* Update _device_offload.py

* Update _device_offload.py

* Update _device_offload.py

* Update _device_offload.py

* Update _sycl_usm.py

* Update _third_party.py

* Update _device_offload.py

* Update _device_offload.py

* Update _device_offload.py

* Update _device_offload.py

* Update _sycl_usm.py

* Update _sycl_usm.py

* Update _third_party.py

* Update _third_party.py

* Update _sycl_usm.py

* Update _third_party.py

* Update _third_party.py

* Update _array_api.py

* Update _array_api.py

* formatting

* Update setup.py

* Update _third_party.py

* Update _third_party.py

* Update _third_party.py

* Update _data_conversion.py

* Update __init__.py

* Update __init__.py

* Update _data_conversion.py

* Update _device_offload.py

* Update _third_party.py

* add requested comments to code

* add requested comments to code

* fix codespell hits

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update test_common.py

* Update _data_conversion.py

* Update _data_conversion.py

* Update test_common.py

* Update _data_conversion.py

* Update _data_conversion.py

Author: icfaust

onedal/_device_offload.py

@@ -16,7 +16,6 @@
 
 import inspect
 import logging
-from collections.abc import Iterable
 from functools import wraps
 
 import numpy as np
@@ -25,14 +24,10 @@
 from onedal import _default_backend as backend
 
 from ._config import _get_config
+from .datatypes import copy_to_dpnp, copy_to_usm, usm_to_numpy
 from .utils import _sycl_queue_manager as QM
 from .utils._array_api import _asarray, _is_numpy_namespace
-from .utils._dpep_helpers import dpctl_available, dpnp_available
-
-if dpctl_available:
-    from dpctl.memory import MemoryUSMDevice, as_usm_memory
-    from dpctl.tensor import usm_ndarray
-
+from .utils._third_party import is_dpnp_ndarray
 
 logger = logging.getLogger("sklearnex")
 cpu_dlpack_device = (backend.kDLCPU, 0)
@@ -69,61 +64,13 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
-if dpnp_available:
-    import dpnp
-
-    from .utils._array_api import _convert_to_dpnp
-
-
-def _copy_to_usm(queue, array):
-    if not dpctl_available:
-        raise RuntimeError(
-            "dpctl need to be installed to work " "with __sycl_usm_array_interface__"
-        )
-
-    if hasattr(array, "__array__"):
-
-        try:
-            mem = MemoryUSMDevice(array.nbytes, queue=queue)
-            mem.copy_from_host(array.tobytes())
-            return usm_ndarray(array.shape, array.dtype, buffer=mem)
-        except ValueError as e:
-            # ValueError will raise if device does not support the dtype
-            # retry with float32 (needed for fp16 and fp64 support issues)
-            # try again as float32, if it is a float32 just raise the error.
-            if array.dtype == np.float32:
-                raise e
-            return _copy_to_usm(queue, array.astype(np.float32))
-    else:
-        if isinstance(array, Iterable):
-            array = [_copy_to_usm(queue, i) for i in array]
-        return array
-
-
 def _transfer_to_host(*data):
     has_usm_data, has_host_data = False, False
 
     host_data = []
     for item in data:
-        usm_iface = getattr(item, "__sycl_usm_array_interface__", None)
-        if usm_iface is not None:
-            if not dpctl_available:
-                raise RuntimeError(
-                    "dpctl need to be installed to work "
-                    "with __sycl_usm_array_interface__"
-                )
-
-            buffer = as_usm_memory(item).copy_to_host()
-            order = "C"
-            if usm_iface["strides"] is not None and len(usm_iface["strides"]) > 1:
-                if usm_iface["strides"][0] < usm_iface["strides"][1]:
-                    order = "F"
-            item = np.ndarray(
-                shape=usm_iface["shape"],
-                dtype=usm_iface["typestr"],
-                buffer=buffer,
-                order=order,
-            )
+        if usm_iface := getattr(item, "__sycl_usm_array_interface__", None):
+            item = usm_to_numpy(item, usm_iface)
             has_usm_data = True
         elif not isinstance(item, np.ndarray) and (
             device := getattr(item, "__dlpack_device__", None)
@@ -215,15 +162,16 @@ def wrapper_impl(*args, **kwargs):
             )
         if _get_config()["use_raw_input"] is True and not override_raw_input:
             if "queue" not in kwargs:
-                usm_iface = getattr(args[0], "__sycl_usm_array_interface__", None)
-                data_queue = usm_iface["syclobj"] if usm_iface is not None else None
-                kwargs["queue"] = data_queue
+                if usm_iface := getattr(args[0], "__sycl_usm_array_interface__", None):
+                    kwargs["queue"] = usm_iface["syclobj"]
+                else:
+                    kwargs["queue"] = None
             return invoke_func(self, *args, **kwargs)
         elif len(args) == 0 and len(kwargs) == 0:
             # no arguments, there's nothing we can deduce from them -> just call the function
             return invoke_func(self, *args, **kwargs)
 
-        data = (*args, *kwargs.values())
+        data = (*args, *kwargs.values())[0]
         # get and set the global queue from the kwarg or data
         with QM.manage_global_queue(kwargs.get("queue"), *args) as queue:
             hostargs, hostkwargs = _get_host_inputs(*args, **kwargs)
@@ -232,17 +180,17 @@ def wrapper_impl(*args, **kwargs):
                 hostkwargs["queue"] = queue
             result = invoke_func(self, *hostargs, **hostkwargs)
 
-            usm_iface = getattr(data[0], "__sycl_usm_array_interface__", None)
-            if queue is not None and usm_iface is not None:
-                result = _copy_to_usm(queue, result)
-                if dpnp_available and isinstance(data[0], dpnp.ndarray):
-                    result = _convert_to_dpnp(result)
-                return result
+            if queue and hasattr(data, "__sycl_usm_array_interface__"):
+                return (
+                    copy_to_dpnp(queue, result)
+                    if is_dpnp_ndarray(data)
+                    else copy_to_usm(queue, result)
+                )
 
         if get_config().get("transform_output") in ("default", None):
-            input_array_api = getattr(data[0], "__array_namespace__", lambda: None)()
+            input_array_api = getattr(data, "__array_namespace__", lambda: None)()
             if input_array_api and not _is_numpy_namespace(input_array_api):
-                input_array_api_device = data[0].device
+                input_array_api_device = data.device
                 result = _asarray(result, input_array_api, device=input_array_api_device)
         return result
 

onedal/common/tests/test_sycl.py

@@ -18,7 +18,7 @@
 
 from onedal import _default_backend as backend
 from onedal.tests.utils._device_selection import get_queues
-from onedal.utils._dpep_helpers import dpctl_available
+from onedal.utils._third_party import dpctl_available
 
 
 @pytest.mark.skipif(

onedal/datatypes/__init__.py

@@ -15,5 +15,12 @@
 # ==============================================================================
 
 from ._data_conversion import from_table, to_table
+from ._sycl_usm import copy_to_dpnp, copy_to_usm, usm_to_numpy
 
-__all__ = ["from_table", "to_table"]
+__all__ = [
+    "copy_to_dpnp",
+    "copy_to_usm",
+    "from_table",
+    "to_table",
+    "usm_to_numpy",
+]

onedal/datatypes/_data_conversion.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from types import ModuleType
+
 import numpy as np
 
 from onedal import _default_backend as backend
@@ -61,29 +63,6 @@ def to_table(*args, queue=None):
 
 if backend.is_dpc:
 
-    try:
-        # try/catch is used here instead of dpep_helpers because
-        # of circular import issues of _data_conversion.py and
-        # utils/validation.py. This is a temporary fix until the
-        # issue with dpnp is addressed, at which point this can
-        # be removed entirely.
-        import dpnp
-
-        def _table_to_array(table, xp=None):
-            # By default DPNP ndarray created with a copy.
-            # TODO:
-            # investigate why dpnp.array(table, copy=False) doesn't work.
-            # Work around with using dpctl.tensor.asarray.
-            if xp == dpnp:
-                return dpnp.array(dpnp.dpctl.tensor.asarray(table), copy=False)
-            else:
-                return xp.asarray(table)
-
-    except ImportError:
-
-        def _table_to_array(table, xp=None):
-            return xp.asarray(table)
-
     def convert_one_from_table(table, sycl_queue=None, sua_iface=None, xp=None):
         # Currently only `__sycl_usm_array_interface__` protocol used to
         # convert into dpnp/dpctl tensors.
@@ -102,7 +81,14 @@ def convert_one_from_table(table, sycl_queue=None, sua_iface=None, xp=None):
                     backend.from_table(table), usm_type="device", sycl_queue=sycl_queue
                 )
             else:
-                return _table_to_array(table, xp=xp)
+                # By default DPNP ndarray created with a copy.
+                # TODO:
+                # investigate why dpnp.array(table, copy=False) doesn't work.
+                # Work around with using dpctl.tensor.asarray.
+                if isinstance(xp, ModuleType) and xp.__name__ == "dpnp":
+                    return xp.array(xp.dpctl.tensor.asarray(table), copy=False)
+                else:
+                    return xp.asarray(table)
 
         return backend.from_table(table)
 

onedal/datatypes/_sycl_usm.py

@@ -0,0 +1,78 @@
+# ==============================================================================
+# Copyright Contributors to the oneDAL Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from collections.abc import Iterable
+
+import numpy as np
+
+from ..utils._third_party import lazy_import
+
+
+@lazy_import("dpctl.memory", "dpctl.tensor")
+def array_to_usm(memory, tensor, queue, array):
+    try:
+        mem = memory.MemoryUSMDevice(array.nbytes, queue=queue)
+        mem.copy_from_host(array.tobytes())
+        return tensor.usm_ndarray(array.shape, array.dtype, buffer=mem)
+    except ValueError as e:
+        # ValueError will raise if device does not support the dtype
+        # retry with float32 (needed for fp16 and fp64 support issues)
+        # try again as float32, if it is a float32 just raise the error.
+        if array.dtype == np.float32:
+            raise e
+        return _array_to_usm(queue, array.astype(np.float32))
+
+
+@lazy_import("dpnp", "dpctl.tensor")
+def to_dpnp(dpnp, tensor, array):
+    if isinstance(array, tensor.usm_ndarray):
+        return dpnp.array(array, copy=False)
+    else:
+        return array
+
+
+def copy_to_usm(queue, array):
+    if hasattr(array, "__array__"):
+        return array_to_usm(queue, array)
+    else:
+        if isinstance(array, Iterable):
+            array = [copy_to_usm(queue, i) for i in array]
+        return array
+
+
+def copy_to_dpnp(queue, array):
+    if hasattr(array, "__array__"):
+        return to_dpnp(array_to_usm(queue, array))
+    else:
+        if isinstance(array, Iterable):
+            array = [copy_to_dpnp(queue, i) for i in array]
+        return array
+
+
+@lazy_import("dpctl.memory")
+def usm_to_numpy(memorymod, item, usm_iface):
+    buffer = memorymod.as_usm_memory(item).copy_to_host()
+    order = "C"
+    if usm_iface["strides"] is not None and len(usm_iface["strides"]) > 1:
+        if usm_iface["strides"][0] < usm_iface["strides"][1]:
+            order = "F"
+    item = np.ndarray(
+        shape=usm_iface["shape"],
+        dtype=usm_iface["typestr"],
+        buffer=buffer,
+        order=order,
+    )
+    return item

onedal/datatypes/tests/common.py

@@ -14,11 +14,16 @@
 # limitations under the License.
 # ===============================================================================
 
-from onedal.utils._dpep_helpers import dpctl_available, dpnp_available
+from onedal.utils._third_party import dpctl_available
 
-if dpnp_available:
+try:
     import dpnp
 
+    dpnp_available = True
+except ImportError:
+    dpnp_available = False
+
+
 if dpctl_available:
     import dpctl
     from dpctl.tensor import usm_ndarray

onedal/datatypes/tests/test_data.py

@@ -24,7 +24,7 @@
 from onedal import _default_backend, _dpc_backend
 from onedal._device_offload import supports_queue
 from onedal.datatypes import from_table, to_table
-from onedal.utils._dpep_helpers import dpctl_available
+from onedal.utils._third_party import dpctl_available
 
 backend = _dpc_backend or _default_backend
 

onedal/ensemble/forest.py

@@ -34,7 +34,6 @@
 from ..common._mixin import ClassifierMixin, RegressorMixin
 from ..datatypes import from_table, to_table
 from ..utils._array_api import _get_sycl_namespace
-from ..utils._dpep_helpers import get_unique_values_with_dpep
 from ..utils.validation import (
     _check_array,
     _check_n_features,
@@ -315,7 +314,13 @@ def _fit(self, X, y, sample_weight):
         else:
             if sua_iface is not None:
                 queue = X.sycl_queue
-            self.classes_ = get_unique_values_with_dpep(y)
+            # try catch needed for raw_inputs + array_api data where unlike
+            # numpy the way to yield unique values is via `unique_values`
+            # This should be removed when refactored for gpu zero-copy
+            try:
+                self.classes_ = xp.unique(y)
+            except AttributeError:
+                self.classes_ = xp.unique_values(y)
 
         self.n_features_in_ = X.shape[1]
 

onedal/linear_model/logistic_regression.py

@@ -29,7 +29,6 @@
 from ..common._mixin import ClassifierMixin
 from ..datatypes import from_table, to_table
 from ..utils._array_api import _get_sycl_namespace
-from ..utils._dpep_helpers import get_unique_values_with_dpep
 from ..utils.validation import (
     _check_array,
     _check_n_features,
@@ -96,7 +95,15 @@ def _fit(self, X, y):
             self.classes_, y = np.unique(y, return_inverse=True)
             y = y.astype(dtype=np.int32)
         else:
-            self.classes_ = get_unique_values_with_dpep(y)
+            _, xp, _ = _get_sycl_namespace(X)
+            # try catch needed for raw_inputs + array_api data where unlike
+            # numpy the way to yield unique values is via `unique_values`
+            # This should be removed when refactored for gpu zero-copy
+            try:
+                self.classes_ = xp.unique(y)
+            except AttributeError:
+                self.classes_ = xp.unique_values(y)
+
             n_classes = len(self.classes_)
             if n_classes != 2:
                 raise ValueError("Only binary classification is supported")

onedal/tests/utils/_dataframes_support.py

@@ -20,14 +20,18 @@
 
 from sklearnex import get_config
 
-from ...utils._dpep_helpers import dpctl_available, dpnp_available
+from ...utils._third_party import dpctl_available
 
 if dpctl_available:
     import dpctl.tensor as dpt
 
-if dpnp_available:
+try:
     import dpnp
 
+    dpnp_available = True
+except ImportError:
+    dpnp_available = False
+
 try:
     # This should be lazy imported in the
     # future along with other popular