Update sklearnex gpu to use array_api_dispatch and dpnp (#218)

* Update sklearnex gpu to use array_api_dispatch and dpnp

* ensure GPU ops on same queue

* address comments

* black

Author: ethanglaser

configs/BENCH-CONFIG-SPEC.md

@@ -98,7 +98,7 @@ Configs have the three highest parameter keys:
 | `data`:`distributed_split` | None | None, `rank_based` | Split type used to distribute data between machines in distributed algorithm. `None` type means usage of all data without split on all machines. `rank_based` type splits the data equally between machines with split sequence based on rank id from MPI. |
 |<h3>Algorithm parameters</h3>||||
 | `algorithm`:`library` | None |  | Python module containing measured entity (class or function). |
-| `algorithm`:`device` | `default` | `default`, `cpu`, `gpu` | Device selected for computation. |
+| `algorithm`:`device` | `default` | `default`, `cpu`, `gpu` | Device selected for computation. `sklearnex`+`gpu` cases enable sklearn's `array_api_dispatch` and use `dpnp` data by default (see `sklearn_context` below). |
 
 ## Benchmark-Specific Parameters
 
@@ -109,7 +109,7 @@ Configs have the three highest parameter keys:
 | `algorithm`:`estimator` | None |  | Name of measured estimator. |
 | `algorithm`:`estimator_params` | Empty `dict` |  | Parameters for estimator constructor. |
 | `algorithm`:`batch_size`:`{stage}` | None | Any positive integer | Enables online mode for `{stage}` methods of estimator (sequential calls for each batch). |
-| `algorithm`:`sklearn_context` | None |  | Parameters for sklearn `config_context` used over estimator. |
+| `algorithm`:`sklearn_context` | None |  | Parameters for sklearn `config_context` used over estimator. `array_api_dispatch` requires `SCIPY_ARRAY_API=1`, which scikit-learn_bench sets by default if it is unset in the environment. |
 | `algorithm`:`sklearnex_context` | None |  | Parameters for sklearnex `config_context` used over estimator. Updated by `sklearn_context` if set. |
 | `bench`:`ensure_sklearnex_patching` | True |  | If True, warns about sklearnex patching failures. |
 

configs/common/sklearn.json

@@ -6,12 +6,18 @@
                 { "library": "sklearnex", "device": "cpu" }
             ]
         },
-        "sklearn-ex[cpu,gpu] implementations": {
-            "algorithm": [
-                { "library": "sklearn", "device": "cpu" },
-                { "library": "sklearnex", "device": ["cpu", "gpu"] }
-            ]
-        },
+        "sklearn-ex[cpu,gpu] implementations": [
+            { "algorithm": { "library": "sklearn", "device": "cpu" } },
+            { "algorithm": { "library": "sklearnex", "device": "cpu" } },
+            {
+                "algorithm": {
+                    "library": "sklearnex",
+                    "device": "gpu",
+                    "sklearn_context": { "array_api_dispatch": true }
+                },
+                "data": { "format": "dpnp", "order": "C" }
+            }
+        ],
         "sklearnex spmd implementation": {
             "algorithm": {
                 "library": "sklearnex.spmd",

sklbench/benchmarks/sklearn_estimator.py

@@ -22,6 +22,9 @@
 from importlib.metadata import PackageNotFoundError, version
 from typing import Dict, List, Union
 
+# sklbench uses array API by default; must precede scipy import to take effect
+os.environ.setdefault("SCIPY_ARRAY_API", "1")
+
 import numpy as np
 import pandas as pd
 from sklearn.base import BaseEstimator

sklbench/datasets/transformer.py

@@ -26,7 +26,9 @@
 from ..utils.logger import logger
 
 
-def convert_data(data, dformat: str, order: str, dtype: str, device: str = None):
+def convert_data(
+    data, dformat: str, order: str, dtype: str, device: str = None, sycl_queue=None
+):
     if isinstance(data, csr_matrix) and dformat != "csr_matrix":
         data = data.toarray()
     if dtype == "preserve":
@@ -46,6 +48,11 @@ def convert_data(data, dformat: str, order: str, dtype: str, device: str = None)
     elif dformat == "dpnp":
         import dpnp
 
+        # Pin every subset to one shared queue: sklearnex builds its internal
+        # arrays (e.g. the take() indices in KNN predict) on the device's default
+        # queue, and array_api_dispatch requires all arrays share one queue object.
+        if sycl_queue is not None:
+            return dpnp.asarray(data, dtype=dtype, order=order, sycl_queue=sycl_queue)
         return dpnp.array(data, dtype=dtype, order=order, device=device)
     elif dformat == "dpctl":
         warnings.warn(
@@ -143,6 +150,14 @@ def split_and_transform_data(bench_case, data, data_description):
 
     device = get_bench_case_value(bench_case, "algorithm:device", None)
     common_data_format = get_bench_case_value(bench_case, "data:format", "pandas")
+
+    # Resolve one queue for the device up front so all dpnp subsets share it;
+    # dpnp.array(device=...) per subset can otherwise land on distinct queues.
+    sycl_queue = None
+    if common_data_format == "dpnp" and device is not None:
+        import dpnp
+
+        sycl_queue = dpnp.array([], device=device).sycl_queue
     common_data_order = get_bench_case_value(bench_case, "data:order", "F")
     common_data_dtype = get_bench_case_value(bench_case, "data:dtype", "float32")
 
@@ -177,7 +192,7 @@ def split_and_transform_data(bench_case, data, data_description):
             data_dtype = required_label_dtype
 
         converted_data = convert_data(
-            subset_content, data_format, data_order, data_dtype, device
+            subset_content, data_format, data_order, data_dtype, device, sycl_queue
         )
         data_dict[subset_name] = converted_data
         if not is_label: