[experiment] ENH: using only raw inputs for onedal backend

onedal/_config.py

@@ -22,6 +22,7 @@
     "target_offload": "auto",
     "allow_fallback_to_host": False,
     "allow_sklearn_after_onedal": True,
+    "use_raw_input": False,
 }
 
 _threadlocal = threading.local()

onedal/_device_offload.py

@@ -162,30 +162,45 @@ def support_input_format(freefunc=False, queue_param=True):
 
     def decorator(func):
         def wrapper_impl(obj, *args, **kwargs):
-            if len(args) == 0 and len(kwargs) == 0:
+            # Check if the function is KNeighborsClassifier.fit
+            override_raw_input = (
+                obj.__class__.__name__ == "KNeighborsClassifier"
+                and func.__name__ == "fit"
+            )
+            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
+                return _run_on_device(func, obj, *args, **kwargs)
+
+            elif 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)
-            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 isinstance(data[0], dpnp.ndarray):
-                    result = _convert_to_dpnp(result)
+
+            else:
+                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)
+                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 isinstance(data[0], dpnp.ndarray):
+                        result = _convert_to_dpnp(result)
+                    return result
+                if not get_config().get("transform_output", False):
+                    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
-            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/basic_statistics/basic_statistics.py

@@ -14,14 +14,15 @@
 # limitations under the License.
 # ==============================================================================
 
-import warnings
 from abc import ABCMeta, abstractmethod
 
 import numpy as np
 
+from .._config import _get_config
 from ..common._base import BaseEstimator
 from ..datatypes import from_table, to_table
 from ..utils import _is_csr
+from ..utils._array_api import _get_sycl_namespace
 from ..utils.validation import _check_array
 
 
@@ -76,10 +77,12 @@ def fit(self, data, sample_weight=None, queue=None):
 
         is_csr = _is_csr(data)
 
-        if data is not None and not is_csr:
-            data = _check_array(data, ensure_2d=False)
-        if sample_weight is not None:
-            sample_weight = _check_array(sample_weight, ensure_2d=False)
+        use_raw_input = _get_config().get("use_raw_input", False) is True
+        if not use_raw_input:
+            if data is not None and not is_csr:
+                data = _check_array(data, ensure_2d=False)
+            if sample_weight is not None:
+                sample_weight = _check_array(sample_weight, ensure_2d=False)
 
         is_single_dim = data.ndim == 1
         data_table, weights_table = to_table(data, sample_weight, queue=queue)

onedal/basic_statistics/incremental_basic_statistics.py

@@ -18,8 +18,10 @@
 
 from daal4py.sklearn._utils import get_dtype
 
+from .._config import _get_config
 from ..datatypes import from_table, to_table
 from ..utils import _check_array
+from ..utils._array_api import _get_sycl_namespace
 from .basic_statistics import BaseBasicStatistics
 
 
@@ -82,6 +84,7 @@ def __getstate__(self):
         self.finalize_fit()
         data = self.__dict__.copy()
         data.pop("_queue", None)
+        data.pop("_input_xp", None)  # module cannot be pickled
 
         return data
 
@@ -104,19 +107,31 @@ def partial_fit(self, X, weights=None, queue=None):
         self : object
             Returns the instance itself.
         """
+        use_raw_input = _get_config().get("use_raw_input", False)
+        sua_iface, xp, _ = _get_sycl_namespace(X)
+        # Saving input array namespace and sua_iface, that will be used in
+        # finalize_fit.
+        self._input_sua_iface = sua_iface
+        self._input_xp = xp
+
+        # All data should use the same sycl queue
+        if use_raw_input and sua_iface:
+            queue = X.sycl_queue
+
         self._queue = queue
         policy = self._get_policy(queue, X)
 
-        X = _check_array(
-            X, dtype=[np.float64, np.float32], ensure_2d=False, force_all_finite=False
-        )
-        if weights is not None:
-            weights = _check_array(
-                weights,
-                dtype=[np.float64, np.float32],
-                ensure_2d=False,
-                force_all_finite=False,
+        if not use_raw_input:
+            X = _check_array(
+                X, dtype=[np.float64, np.float32], ensure_2d=False, force_all_finite=False
             )
+            if weights is not None:
+                weights = _check_array(
+                    weights,
+                    dtype=[np.float64, np.float32],
+                    ensure_2d=False,
+                    force_all_finite=False,
+                )
 
         if not hasattr(self, "_onedal_params"):
             dtype = get_dtype(X)

onedal/cluster/dbscan.py

@@ -18,10 +18,12 @@
 
 from daal4py.sklearn._utils import get_dtype, make2d
 
+from .._config import _get_config
 from ..common._base import BaseEstimator
 from ..common._mixin import ClusterMixin
 from ..datatypes import from_table, to_table
 from ..utils import _check_array
+from ..utils._array_api import _asarray, _get_sycl_namespace
 
 
 class BaseDBSCAN(BaseEstimator, ClusterMixin):
@@ -57,18 +59,26 @@ def _get_onedal_params(self, dtype=np.float32):
         }
 
     def _fit(self, X, y, sample_weight, module, queue):
+        use_raw_input = _get_config().get("use_raw_input", False) is True
         policy = self._get_policy(queue, X)
-        X = _check_array(X, accept_sparse="csr", dtype=[np.float64, np.float32])
-        sample_weight = make2d(sample_weight) if sample_weight is not None else None
-        X_table, sample_weight_table = to_table(X, sample_weight, queue=queue)
 
+        if not use_raw_input:
+            X = _check_array(X, accept_sparse="csr", dtype=[np.float64, np.float32])
+            sample_weight = make2d(sample_weight) if sample_weight is not None else None
+            X = make2d(X)
+        X_table, sample_weight_table = to_table(X, sample_weight, queue=queue)
         params = self._get_onedal_params(X_table.dtype)
         result = module.compute(policy, params, X_table, sample_weight_table)
 
-        self.labels_ = from_table(result.responses).ravel()
-        if result.core_observation_indices is not None:
+        _, xp, _ = _get_sycl_namespace(X)
+        self.labels_ = from_table(result.responses, sycl_queue=queue).ravel()
+        if (
+            result.core_observation_indices is not None
+            and not result.core_observation_indices.kind == "empty"
+        ):
             self.core_sample_indices_ = from_table(
-                result.core_observation_indices
+                result.core_observation_indices,
+                sycl_queue=queue,
             ).ravel()
         else:
             self.core_sample_indices_ = np.array([], dtype=np.intc)

onedal/cluster/kmeans.py

@@ -32,10 +32,12 @@
 from sklearn.metrics.pairwise import euclidean_distances
 from sklearn.utils import check_random_state
 
+from .._config import _get_config
 from ..common._base import BaseEstimator as onedal_BaseEstimator
 from ..common._mixin import ClusterMixin, TransformerMixin
 from ..datatypes import from_table, to_table
 from ..utils import _check_array, _is_arraylike_not_scalar, _is_csr
+from ..utils._array_api import _get_sycl_namespace
 
 
 class _BaseKMeans(onedal_BaseEstimator, TransformerMixin, ClusterMixin, ABC):
@@ -80,7 +82,7 @@ def _get_kmeans_init(self, cluster_count, seed, algorithm):
     def _get_basic_statistics_backend(self, result_options):
         return BasicStatistics(result_options)
 
-    def _tolerance(self, X_table, rtol, is_csr, policy, dtype):
+    def _tolerance(self, X_table, rtol, is_csr, policy, dtype, sua_iface):
         """Compute absolute tolerance from the relative tolerance"""
         if rtol == 0.0:
             return rtol
@@ -94,7 +96,7 @@ def _tolerance(self, X_table, rtol, is_csr, policy, dtype):
         return mean_var * rtol
 
     def _check_params_vs_input(
-        self, X_table, is_csr, policy, default_n_init=10, dtype=np.float32
+        self, X_table, is_csr, policy, default_n_init=10, dtype=np.float32, sua_iface=None
     ):
         # n_clusters
         if X_table.shape[0] < self.n_clusters:
@@ -103,7 +105,7 @@ def _check_params_vs_input(
             )
 
         # tol
-        self._tol = self._tolerance(X_table, self.tol, is_csr, policy, dtype)
+        self._tol = self._tolerance(X_table, self.tol, is_csr, policy, dtype, sua_iface)
 
         # n-init
         # TODO(1.4): Remove
@@ -261,9 +263,14 @@ def _fit_backend(
     def _fit(self, X, module, queue=None):
         policy = self._get_policy(queue, X)
         is_csr = _is_csr(X)
-        X = _check_array(
-            X, dtype=[np.float64, np.float32], accept_sparse="csr", force_all_finite=False
-        )
+
+        if _get_config()["use_raw_input"] is False:
+            X = _check_array(
+                X,
+                dtype=[np.float64, np.float32],
+                accept_sparse="csr",
+                force_all_finite=False,
+            )
         X_table = to_table(X, queue=queue)
         dtype = X_table.dtype
 

onedal/covariance/covariance.py

@@ -17,12 +17,14 @@
 
 import numpy as np
 
-from daal4py.sklearn._utils import daal_check_version, get_dtype
-from onedal.utils import _check_array
+from daal4py.sklearn._utils import daal_check_version
 
+from .._config import _get_config
 from ..common._base import BaseEstimator
 from ..common.hyperparameters import get_hyperparameters
 from ..datatypes import from_table, to_table
+from ..utils import _check_array
+from ..utils._array_api import _get_sycl_namespace
 
 
 class BaseEmpiricalCovariance(BaseEstimator, metaclass=ABCMeta):
@@ -93,8 +95,14 @@ def fit(self, X, y=None, queue=None):
         self : object
             Returns the instance itself.
         """
+        use_raw_input = _get_config()["use_raw_input"]
+        sua_iface, xp, _ = _get_sycl_namespace(X)
+        if use_raw_input and sua_iface:
+            queue = X.sycl_queue
+
         policy = self._get_policy(queue, X)
-        X = _check_array(X, dtype=[np.float64, np.float32])
+        if not use_raw_input:
+            X = _check_array(X, dtype=[np.float64, np.float32])
         X = to_table(X, queue=queue)
         params = self._get_onedal_params(X.dtype)
         hparams = get_hyperparameters("covariance", "compute")
@@ -111,12 +119,14 @@ def fit(self, X, y=None, queue=None):
         else:
             result = self._get_backend("covariance", None, "compute", policy, params, X)
         if daal_check_version((2024, "P", 1)) or (not self.bias):
-            self.covariance_ = from_table(result.cov_matrix)
+            self.covariance_ = from_table(result.cov_matrix, sycl_queue=queue)
         else:
             self.covariance_ = (
-                from_table(result.cov_matrix) * (X.shape[0] - 1) / X.shape[0]
+                from_table(result.cov_matrix, sycl_queue=queue)
+                * (X.shape[0] - 1)
+                / X.shape[0]
             )
 
-        self.location_ = from_table(result.means).ravel()
+        self.location_ = from_table(result.means, sycl_queue=queue).ravel()
 
         return self

onedal/covariance/incremental_covariance.py

@@ -15,10 +15,12 @@
 # ===============================================================================
 import numpy as np
 
-from daal4py.sklearn._utils import daal_check_version, get_dtype
+from daal4py.sklearn._utils import daal_check_version
 
+from .._config import _get_config
 from ..datatypes import from_table, to_table
 from ..utils import _check_array
+from ..utils._array_api import _get_sycl_namespace
 from .covariance import BaseEmpiricalCovariance
 
 
@@ -70,6 +72,7 @@ def __getstate__(self):
         self.finalize_fit()
         data = self.__dict__.copy()
         data.pop("_queue", None)
+        data.pop("_input_xp", None)  # module cannot be pickled
 
         return data
 
@@ -95,10 +98,19 @@ def partial_fit(self, X, y=None, queue=None):
         self : object
             Returns the instance itself.
         """
-        X = _check_array(X, dtype=[np.float64, np.float32], ensure_2d=True)
+        # Saving input array namespace and sua_iface, that will be used in
+        # finalize_fit.
+        sua_iface, xp, _ = _get_sycl_namespace(X)
+        self._input_sua_iface = sua_iface
+        self._input_xp = xp
+
+        use_raw_input = _get_config().get("use_raw_input", False)
+        if use_raw_input and sua_iface:
+            queue = X.sycl_queue
+        if not use_raw_input:
+            X = _check_array(X, dtype=[np.float64, np.float32], ensure_2d=True)
 
         self._queue = queue
-
         policy = self._get_policy(queue, X)
 
         X_table = to_table(X, queue=queue)