ENH: functional support for Array API (#1861)

* updates on `_device_offload` modules both of `onedal4py` and `sklearne`x.

* extensions for `sklearnex._device_offload.wrap_output_data` decorator and `sklearnex._device_offload.dispatch` function.

* `onedal._device_offload._transfer_to_host` extended for the Array API inputs.

* `onedal._device_offload.support_usm_ndarray` renamed to `onedal._device_offload.support_input_format` extend for Array API inputs handling. (several files changed for this this renaming)

* onedal4py's testing utility `get_dataframes_and_queues` extended for the array_api inputs
_array_api modules added in the both onedal4py and sklearnex.\

* `onedal.utils._array_api` module includes several array api utilities and `_get_sycl_namespace` that is used in the `get_namespace` (from sklearnex.utils._arrray_api), that returns supported array api namespace.

Author: samir-nasibli

onedal/_device_offload.py

@@ -19,8 +19,10 @@
 from functools import wraps
 
 import numpy as np
+from sklearn import get_config
 
 from ._config import _get_config
+from .utils._array_api import _asarray, _is_numpy_namespace
 
 try:
     from dpctl import SyclQueue
@@ -34,6 +36,8 @@
 try:
     import dpnp
 
+    from .utils._array_api import _convert_to_dpnp
+
     dpnp_available = True
 except ImportError:
     dpnp_available = False
@@ -94,6 +98,7 @@ def _transfer_to_host(queue, *data):
     host_data = []
     for item in data:
         usm_iface = getattr(item, "__sycl_usm_array_interface__", None)
+        array_api = getattr(item, "__array_namespace__", lambda: None)()
         if usm_iface is not None:
             if not dpctl_available:
                 raise RuntimeError(
@@ -120,6 +125,11 @@ def _transfer_to_host(queue, *data):
                 order=order,
             )
             has_usm_data = True
+        elif array_api and not _is_numpy_namespace(array_api):
+            # `copy`` param for the `asarray`` is not setted.
+            # The object is copied only if needed.
+            item = np.asarray(item)
+            has_host_data = True
         else:
             has_host_data = True
 
@@ -153,34 +163,17 @@ def _get_host_inputs(*args, **kwargs):
     return q, hostargs, hostkwargs
 
 
-def _extract_usm_iface(*args, **kwargs):
-    allargs = (*args, *kwargs.values())
-    if len(allargs) == 0:
-        return None
-    return getattr(allargs[0], "__sycl_usm_array_interface__", None)
-
-
 def _run_on_device(func, obj=None, *args, **kwargs):
     if obj is not None:
         return func(obj, *args, **kwargs)
     return func(*args, **kwargs)
 
 
-if dpnp_available:
-
-    def _convert_to_dpnp(array):
-        if isinstance(array, usm_ndarray):
-            return dpnp.array(array, copy=False)
-        elif isinstance(array, Iterable):
-            for i in range(len(array)):
-                array[i] = _convert_to_dpnp(array[i])
-        return array
-
-
-def support_usm_ndarray(freefunc=False, queue_param=True):
+def support_input_format(freefunc=False, queue_param=True):
     """
-    Handles USMArray input. Puts SYCLQueue from data to decorated function arguments.
-    Converts output of decorated function to dpctl.tensor/dpnp.ndarray if input was of this type.
+    Converts and moves the output arrays of the decorated function
+    to match the input array type and device.
+    Puts SYCLQueue from data to decorated function arguments.
 
     Parameters
     ----------
@@ -194,17 +187,29 @@ def support_usm_ndarray(freefunc=False, queue_param=True):
 
     def decorator(func):
         def wrapper_impl(obj, *args, **kwargs):
-            usm_iface = _extract_usm_iface(*args, **kwargs)
+            if len(args) == 0 and len(kwargs) == 0:
+                return _run_on_device(func, obj, *args, **kwargs)
+            data = (*args, *kwargs.values())
             data_queue, hostargs, hostkwargs = _get_host_inputs(*args, **kwargs)
             if queue_param and not (
                 "queue" in hostkwargs and hostkwargs["queue"] is not None
             ):
                 hostkwargs["queue"] = data_queue
             result = _run_on_device(func, obj, *hostargs, **hostkwargs)
-            if usm_iface is not None and hasattr(result, "__array_interface__"):
+            usm_iface = getattr(data[0], "__sycl_usm_array_interface__", None)
+            if usm_iface is not None:
                 result = _copy_to_usm(data_queue, result)
-                if dpnp_available and len(args) > 0 and isinstance(args[0], dpnp.ndarray):
+                if dpnp_available and isinstance(data[0], dpnp.ndarray):
                     result = _convert_to_dpnp(result)
+                return result
+            config = get_config()
+            if not ("transform_output" in config and config["transform_output"]):
+                input_array_api = getattr(data[0], "__array_namespace__", lambda: None)()
+                if input_array_api:
+                    input_array_api_device = data[0].device
+                    result = _asarray(
+                        result, input_array_api, device=input_array_api_device
+                    )
             return result
 
         if freefunc:

onedal/spmd/basic_statistics/basic_statistics.py

@@ -16,15 +16,15 @@
 
 from onedal.basic_statistics import BasicStatistics as BasicStatistics_Batch
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
 class BasicStatistics(BaseEstimatorSPMD, BasicStatistics_Batch):
-    @support_usm_ndarray()
+    @support_input_format()
     def compute(self, data, weights=None, queue=None):
         return super().compute(data, weights=weights, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, data, sample_weight=None, queue=None):
         return super().fit(data, sample_weight=sample_weight, queue=queue)

onedal/spmd/cluster/kmeans.py

@@ -18,7 +18,7 @@
 from onedal.cluster import KMeansInit as KMeansInit_Batch
 from onedal.spmd.basic_statistics import BasicStatistics
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
@@ -37,15 +37,15 @@ def _get_basic_statistics_backend(self, result_options):
     def _get_kmeans_init(self, cluster_count, seed, algorithm):
         return KMeansInit(cluster_count=cluster_count, seed=seed, algorithm=algorithm)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y=None, queue=None):
         return super().fit(X, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict(self, X, queue=None):
         return super().predict(X, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def fit_predict(self, X, y=None, queue=None):
         return super().fit_predict(X, queue=queue)
 

onedal/spmd/covariance/covariance.py

@@ -16,11 +16,11 @@
 
 from onedal.covariance import EmpiricalCovariance as EmpiricalCovariance_Batch
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
 class EmpiricalCovariance(BaseEstimatorSPMD, EmpiricalCovariance_Batch):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y=None, queue=None):
         return super().fit(X, queue=queue)

onedal/spmd/decomposition/pca.py

@@ -16,11 +16,11 @@
 
 from onedal.decomposition.pca import PCA as PCABatch
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
 class PCA(BaseEstimatorSPMD, PCABatch):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y=None, queue=None):
         return super().fit(X, queue=queue)

onedal/spmd/linear_model/linear_model.py

@@ -16,15 +16,15 @@
 
 from onedal.linear_model import LinearRegression as LinearRegression_Batch
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
 class LinearRegression(BaseEstimatorSPMD, LinearRegression_Batch):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict(self, X, queue=None):
         return super().predict(X, queue=queue)

onedal/spmd/linear_model/logistic_regression.py

@@ -16,23 +16,23 @@
 
 from onedal.linear_model import LogisticRegression as LogisticRegression_Batch
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
 class LogisticRegression(BaseEstimatorSPMD, LogisticRegression_Batch):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict(self, X, queue=None):
         return super().predict(X, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict_proba(self, X, queue=None):
         return super().predict_proba(X, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict_log_proba(self, X, queue=None):
         return super().predict_log_proba(X, queue=queue)

onedal/spmd/neighbors/neighbors.py

@@ -17,30 +17,30 @@
 from onedal.neighbors import KNeighborsClassifier as KNeighborsClassifier_Batch
 from onedal.neighbors import KNeighborsRegressor as KNeighborsRegressor_Batch
 
-from ..._device_offload import support_usm_ndarray
+from ..._device_offload import support_input_format
 from .._base import BaseEstimatorSPMD
 
 
 class KNeighborsClassifier(BaseEstimatorSPMD, KNeighborsClassifier_Batch):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict(self, X, queue=None):
         return super().predict(X, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict_proba(self, X, queue=None):
         raise NotImplementedError("predict_proba not supported in distributed mode.")
 
-    @support_usm_ndarray()
+    @support_input_format()
     def kneighbors(self, X=None, n_neighbors=None, return_distance=True, queue=None):
         return super().kneighbors(X, n_neighbors, return_distance, queue=queue)
 
 
 class KNeighborsRegressor(BaseEstimatorSPMD, KNeighborsRegressor_Batch):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y, queue=None):
         if queue is not None and queue.sycl_device.is_gpu:
             return super()._fit(X, y, queue=queue)
@@ -50,11 +50,11 @@ def fit(self, X, y, queue=None):
                 "CPU. Consider running on it on GPU."
             )
 
-    @support_usm_ndarray()
+    @support_input_format()
     def kneighbors(self, X=None, n_neighbors=None, return_distance=True, queue=None):
         return super().kneighbors(X, n_neighbors, return_distance, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def predict(self, X, queue=None):
         return self._predict_gpu(X, queue=queue)
 
@@ -66,10 +66,10 @@ def _get_onedal_params(self, X, y=None):
 
 
 class NearestNeighbors(BaseEstimatorSPMD):
-    @support_usm_ndarray()
+    @support_input_format()
     def fit(self, X, y, queue=None):
         return super().fit(X, y, queue=queue)
 
-    @support_usm_ndarray()
+    @support_input_format()
     def kneighbors(self, X=None, n_neighbors=None, return_distance=True, queue=None):
         return super().kneighbors(X, n_neighbors, return_distance, queue=queue)

onedal/tests/utils/_dataframes_support.py

@@ -16,10 +16,10 @@
 
 import pytest
 import scipy.sparse as sp
-from sklearn import get_config
+
+from sklearnex import get_config
 
 try:
-    import dpctl
     import dpctl.tensor as dpt
 
     dpctl_available = True
@@ -40,7 +40,6 @@
     # GPU-no-copy.
     import array_api_strict
 
-    # Run check if "array_api_dispatch" is configurable
     array_api_enabled = lambda: get_config()["array_api_dispatch"]
     array_api_enabled()
     array_api_modules = {"array_api": array_api_strict}
@@ -58,7 +57,7 @@
 
 
 def get_dataframes_and_queues(
-    dataframe_filter_="numpy,pandas,dpnp,dpctl", device_filter_="cpu,gpu"
+    dataframe_filter_="numpy,pandas,dpnp,dpctl,array_api", device_filter_="cpu,gpu"
 ):
     """Get supported dataframes for testing.
 
@@ -107,13 +106,18 @@ def get_df_and_q(dataframe: str):
         dataframes_and_queues.extend(get_df_and_q("dpctl"))
     if dpnp_available and "dpnp" in dataframe_filter_:
         dataframes_and_queues.extend(get_df_and_q("dpnp"))
-    if "array_api" in dataframe_filter_ or array_api_enabled():
+    if (
+        "array_api" in dataframe_filter_
+        and "array_api" in array_api_modules
+        or array_api_enabled()
+    ):
         dataframes_and_queues.append(pytest.param("array_api", None, id="array_api"))
 
     return dataframes_and_queues
 
 
 def _as_numpy(obj, *args, **kwargs):
+    """Converted input object to numpy.ndarray format."""
     if dpnp_available and isinstance(obj, dpnp.ndarray):
         return obj.asnumpy(*args, **kwargs)
     if dpctl_available and isinstance(obj, dpt.usm_ndarray):
@@ -155,17 +159,10 @@ def _convert_to_dataframe(obj, sycl_queue=None, target_df=None, *args, **kwargs)
         # DPCtl tensor.
         return dpt.asarray(obj, usm_type="device", sycl_queue=sycl_queue, *args, **kwargs)
     elif target_df in array_api_modules:
-        # use dpctl to define gpu devices via queues and
-        # move data to the device. This is necessary as
-        # the standard for defining devices is
-        # purposefully not defined in the array_api
-        # standard, but maintaining data on a device
-        # using the method `from_dlpack` is.
+        # Array API input other than DPNP ndarray, DPCtl tensor or
+        # Numpy ndarray.
+
         xp = array_api_modules[target_df]
-        return xp.from_dlpack(
-            _convert_to_dataframe(
-                obj, sycl_queue=sycl_queue, target_df="dpctl", *args, **kwargs
-            )
-        )
+        return xp.asarray(obj)
 
     raise RuntimeError("Unsupported dataframe conversion")