[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

@@ -180,30 +180,35 @@ 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:
-                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)
+            if not get_config()["use_raw_input"] == True:
+                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)
+                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
+                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
-            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
+            else:
+                return _run_on_device(func, obj, *args, **kwargs)
 
         if freefunc:
 

onedal/basic_statistics/basic_statistics.py

@@ -19,9 +19,11 @@
 
 import numpy as np
 
+from .._config import _get_config
 from ..common._base import BaseEstimator
 from ..datatypes import _convert_to_supported, from_table, to_table
 from ..utils import _is_csr
+from ..utils._array_api import _get_sycl_namespace
 from ..utils.validation import _check_array
 
 
@@ -72,23 +74,37 @@ def __init__(self, result_options="all", algorithm="by_default"):
         super().__init__(result_options, algorithm)
 
     def fit(self, data, sample_weight=None, queue=None):
+        use_raw_input = _get_config()["use_raw_input"]
+        # All data should use the same sycl queue.
+        sua_iface, xp, _ = _get_sycl_namespace(data)
+        # TODO:
+        # update support_input_format.
+        if use_raw_input and sua_iface:
+            queue = data.sycl_queue
         policy = self._get_policy(queue, data, sample_weight)
-
-        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)
-
+        if not use_raw_input:
+            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)
+        # TODO
+        # use xp for dtype.
         data, sample_weight = _convert_to_supported(policy, data, sample_weight)
         is_single_dim = data.ndim == 1
-        data_table, weights_table = to_table(data, sample_weight)
+        data_table = to_table(data, sua_iface=sua_iface)
+        weights_table = to_table(sample_weight, sua_iface=sua_iface)
 
         dtype = data.dtype
         raw_result = self._compute_raw(data_table, weights_table, policy, dtype, is_csr)
         for opt, raw_value in raw_result.items():
-            value = from_table(raw_value).ravel()
+            # value = from_table(raw_value.responses, sua_iface=sua_iface, sycl_queue=queue, xp=xp).reshape(-1)
+            value = xp.ravel(
+                from_table(
+                    raw_value.responses, sua_iface=sua_iface, sycl_queue=queue, xp=xp
+                )
+            )
             if is_single_dim:
                 setattr(self, opt, value[0])
             else:

sklearnex/_config.py

@@ -42,6 +42,7 @@ def set_config(
     target_offload=None,
     allow_fallback_to_host=None,
     allow_sklearn_after_onedal=None,
+    use_raw_input=None,
     **sklearn_configs,
 ):
     """Set global configuration
@@ -75,6 +76,8 @@ def set_config(
         local_config["allow_fallback_to_host"] = allow_fallback_to_host
     if allow_sklearn_after_onedal is not None:
         local_config["allow_sklearn_after_onedal"] = allow_sklearn_after_onedal
+    if use_raw_input is not None:
+        local_config["use_raw_input"] = use_raw_input
 
 
 @contextmanager

sklearnex/_device_offload.py

@@ -58,44 +58,49 @@ def _get_backend(obj, queue, method_name, *data):
 
 
 def dispatch(obj, method_name, branches, *args, **kwargs):
-    q = _get_global_queue()
-    has_usm_data_for_args, q, hostargs = _transfer_to_host(q, *args)
-    has_usm_data_for_kwargs, q, hostvalues = _transfer_to_host(q, *kwargs.values())
-    hostkwargs = dict(zip(kwargs.keys(), hostvalues))
-
-    backend, q, patching_status = _get_backend(obj, q, method_name, *hostargs)
-    has_usm_data = has_usm_data_for_args or has_usm_data_for_kwargs
-    if backend == "onedal":
-        # Host args only used before onedal backend call.
-        # Device will be offloaded when onedal backend will be called.
-        patching_status.write_log(queue=q, transferred_to_host=False)
-        return branches[backend](obj, *hostargs, **hostkwargs, queue=q)
-    if backend == "sklearn":
-        if (
-            "array_api_dispatch" in get_config()
-            and get_config()["array_api_dispatch"]
-            and "array_api_support" in obj._get_tags()
-            and obj._get_tags()["array_api_support"]
-            and not has_usm_data
-        ):
-            # USM ndarrays are also excluded for the fallback Array API. Currently, DPNP.ndarray is
-            # not compliant with the Array API standard, and DPCTL usm_ndarray Array API is compliant,
-            # except for the linalg module. There is no guarantee that stock scikit-learn will
-            # work with such input data. The condition will be updated after DPNP.ndarray and
-            # DPCTL usm_ndarray enabling for conformance testing and these arrays supportance
-            # of the fallback cases.
-            # If `array_api_dispatch` enabled and array api is supported for the stock scikit-learn,
-            # then raw inputs are used for the fallback.
-            patching_status.write_log(transferred_to_host=False)
-            return branches[backend](obj, *args, **kwargs)
-        else:
-            patching_status.write_log()
-            return branches[backend](obj, *hostargs, **hostkwargs)
-    raise RuntimeError(
-        f"Undefined backend {backend} in " f"{obj.__class__.__name__}.{method_name}"
-    )
+    if not get_config()["use_raw_input"] == True:
+        q = _get_global_queue()
+        has_usm_data_for_args, q, hostargs = _transfer_to_host(q, *args)
+        has_usm_data_for_kwargs, q, hostvalues = _transfer_to_host(q, *kwargs.values())
+        hostkwargs = dict(zip(kwargs.keys(), hostvalues))
+
+        backend, q, patching_status = _get_backend(obj, q, method_name, *hostargs)
+        has_usm_data = has_usm_data_for_args or has_usm_data_for_kwargs
+        if backend == "onedal":
+            # Host args only used before onedal backend call.
+            # Device will be offloaded when onedal backend will be called.
+            patching_status.write_log(queue=q, transferred_to_host=False)
+            return branches[backend](obj, *hostargs, **hostkwargs, queue=q)
+        if backend == "sklearn":
+            if (
+                "array_api_dispatch" in get_config()
+                and get_config()["array_api_dispatch"]
+                and "array_api_support" in obj._get_tags()
+                and obj._get_tags()["array_api_support"]
+                and not has_usm_data
+            ):
+                # USM ndarrays are also excluded for the fallback Array API. Currently, DPNP.ndarray is
+                # not compliant with the Array API standard, and DPCTL usm_ndarray Array API is compliant,
+                # except for the linalg module. There is no guarantee that stock scikit-learn will
+                # work with such input data. The condition will be updated after DPNP.ndarray and
+                # DPCTL usm_ndarray enabling for conformance testing and these arrays supportance
+                # of the fallback cases.
+                # If `array_api_dispatch` enabled and array api is supported for the stock scikit-learn,
+                # then raw inputs are used for the fallback.
+                patching_status.write_log(transferred_to_host=False)
+                return branches[backend](obj, *args, **kwargs)
+            else:
+                patching_status.write_log()
+                return branches[backend](obj, *hostargs, **hostkwargs)
+        raise RuntimeError(
+            f"Undefined backend {backend} in " f"{obj.__class__.__name__}.{method_name}"
+        )
+    else:
+        return branches["onedal"](obj, *args, **kwargs)
 
 
+# TODO:
+# wrap output.
 def wrap_output_data(func):
     """
     Converts and moves the output arrays of the decorated function