enh: Add SPMD interfaces for Covariance (#1697)

* enh: Add SPMD interfaces for Covariance

* linted

Author: ethanglaser

examples/sklearnex/covariance_spmd.py

@@ -0,0 +1,42 @@
+# ==============================================================================
+# Copyright 2024 Intel Corporation
+#
+# 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.
+# ==============================================================================
+
+import dpctl
+import dpctl.tensor as dpt
+import numpy as np
+from mpi4py import MPI
+
+from sklearnex.spmd.covariance import EmpiricalCovariance
+
+
+def get_data(data_seed):
+    ns, nf = 3000, 3
+    drng = np.random.default_rng(data_seed)
+    X = drng.random(size=(ns, nf))
+    return X
+
+
+q = dpctl.SyclQueue("gpu")
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+size = comm.Get_size()
+
+X = get_data(rank)
+dpt_X = dpt.asarray(X, usm_type="device", sycl_queue=q)
+
+cov = EmpiricalCovariance().fit(dpt_X)
+
+print(f"Computed covariance values on rank {rank}:\n", cov.covariance_)

onedal/spmd/__init__.py

@@ -17,6 +17,7 @@
 __all__ = [
     "basic_statistics",
     "cluster",
+    "covariance",
     "decomposition",
     "ensemble",
     "linear_model",

onedal/spmd/covariance/__init__.py

@@ -0,0 +1,19 @@
+# ==============================================================================
+# Copyright 2024 Intel Corporation
+#
+# 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 .covariance import EmpiricalCovariance
+
+__all__ = ["EmpiricalCovariance"]

onedal/spmd/covariance/covariance.py

@@ -0,0 +1,26 @@
+# ==============================================================================
+# Copyright 2024 Intel Corporation
+#
+# 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 onedal.covariance import EmpiricalCovariance as EmpiricalCovariance_Batch
+
+from ..._device_offload import support_usm_ndarray
+from .._common import BaseEstimatorSPMD
+
+
+class EmpiricalCovariance(BaseEstimatorSPMD, EmpiricalCovariance_Batch):
+    @support_usm_ndarray()
+    def fit(self, X, queue=None):
+        return super().fit(X, queue)

setup.py

@@ -565,6 +565,7 @@ def run(self):
 if build_distribute:
     packages_with_tests += [
         "onedal.spmd",
+        "onedal.spmd.covariance",
         "onedal.spmd.decomposition",
         "onedal.spmd.ensemble",
     ]

setup_sklearnex.py

@@ -105,6 +105,7 @@
 if build_distribute:
     packages_with_tests += [
         "sklearnex.spmd",
+        "sklearnex.spmd.covariance",
         "sklearnex.spmd.decomposition",
         "sklearnex.spmd.ensemble",
     ]

sklearnex/spmd/__init__.py

@@ -17,6 +17,7 @@
 __all__ = [
     "basic_statistics",
     "cluster",
+    "covariance",
     "decomposition",
     "ensemble",
     "linear_model",

sklearnex/spmd/covariance/__init__.py

@@ -0,0 +1,19 @@
+# ==============================================================================
+# Copyright 2024 Intel Corporation
+#
+# 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 .covariance import EmpiricalCovariance
+
+__all__ = ["EmpiricalCovariance"]

sklearnex/spmd/covariance/covariance.py

@@ -0,0 +1,21 @@
+# ==============================================================================
+# Copyright 2024 Intel Corporation
+#
+# 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 onedal.spmd.covariance import EmpiricalCovariance
+
+# TODO:
+# Currently it uses `onedal` module interface.
+# Add sklearnex dispatching.

tests/run_examples.py

@@ -147,6 +147,7 @@ def check_library(rule):
 
 req_device = defaultdict(lambda: [])
 req_device["basic_statistics_spmd.py"] = ["gpu"]
+req_device["covariance_spmd.py"] = ["gpu"]
 req_device["dbscan_spmd.py"] = ["gpu"]
 req_device["kmeans_spmd.py"] = ["gpu"]
 req_device["knn_bf_classification_dpnp.py"] = ["gpu"]
@@ -163,6 +164,7 @@ def check_library(rule):
 
 req_library = defaultdict(lambda: [])
 req_library["basic_statistics_spmd.py"] = ["dpctl", "mpi4py"]
+req_library["covariance_spmd.py"] = ["dpctl", "mpi4py"]
 req_library["dbscan_spmd.py"] = ["dpctl", "mpi4py"]
 req_library["basic_statistics_spmd.py"] = ["dpctl", "mpi4py"]
 req_library["kmeans_spmd.py"] = ["dpctl", "mpi4py"]