From 8cbe1bd60fe582962110011ffa955f61c66b402b Mon Sep 17 00:00:00 2001 From: NIne-WIngEd Date: Sat, 15 Aug 2026 21:32:20 -0500 Subject: [PATCH] Use RAFT backend for make_blobs --- python/cuml/cuml/datasets/CMakeLists.txt | 3 +- python/cuml/cuml/datasets/_blobs.pyx | 116 ++++++++++++ python/cuml/cuml/datasets/blobs.py | 138 ++++++++++++++ python/cuml/tests/test_make_blobs.py | 228 ++++++++++++++++++++++- 4 files changed, 483 insertions(+), 2 deletions(-) create mode 100644 python/cuml/cuml/datasets/_blobs.pyx diff --git a/python/cuml/cuml/datasets/CMakeLists.txt b/python/cuml/cuml/datasets/CMakeLists.txt index 622b20ce9f..12ac1e68f8 100644 --- a/python/cuml/cuml/datasets/CMakeLists.txt +++ b/python/cuml/cuml/datasets/CMakeLists.txt @@ -1,12 +1,13 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= set(cython_sources "") add_module_gpu_default("arima.pyx" ${arima_algo} ${datasets_algo}) +add_module_gpu_default("_blobs.pyx" ${datasets_algo}) add_module_gpu_default("regression.pyx" ${regression_algo} ${datasets_algo}) rapids_cython_create_modules( diff --git a/python/cuml/cuml/datasets/_blobs.pyx b/python/cuml/cuml/datasets/_blobs.pyx new file mode 100644 index 0000000000..32ae67c08f --- /dev/null +++ b/python/cuml/cuml/datasets/_blobs.pyx @@ -0,0 +1,116 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +import cupy as cp + +from cuml.internals import get_handle + +from libc.stddef cimport size_t +from libc.stdint cimport int64_t, uint64_t, uintptr_t +from libcpp cimport bool +from pylibraft.common.handle cimport handle_t + + +cdef extern from "cuml/datasets/make_blobs.hpp" namespace "ML" nogil: + void cpp_make_blobs "ML::Datasets::make_blobs" ( + const handle_t& handle, + float* out, + int64_t* labels, + int64_t n_rows, + int64_t n_cols, + int64_t n_clusters, + bool row_major, + const float* centers, + const float* cluster_std, + const float cluster_std_scalar, + bool shuffle, + float center_box_min, + float center_box_max, + uint64_t seed) except + + + void cpp_make_blobs "ML::Datasets::make_blobs" ( + const handle_t& handle, + double* out, + int64_t* labels, + int64_t n_rows, + int64_t n_cols, + int64_t n_clusters, + bool row_major, + const double* centers, + const double* cluster_std, + const double cluster_std_scalar, + bool shuffle, + double center_box_min, + double center_box_max, + uint64_t seed) except + + + +def make_blobs( + n_samples, + n_features, + n_centers, + centers, + cluster_std, + center_box_min, + center_box_max, + shuffle, + random_state, + order, + dtype, +): + dtype = cp.dtype(dtype) + + h = get_handle() + cdef handle_t* h_ptr = h.getHandle() + + X = cp.empty((n_samples, n_features), dtype=dtype, order=order) + y = cp.empty(n_samples, dtype=cp.int64) + + cdef uintptr_t x_p = X.data.ptr + cdef uintptr_t y_p = y.data.ptr + cdef uintptr_t ctr_p = 0 + if centers is not None: + ctr_p = centers.data.ptr + + cdef bool row_c = order == "C" + + if dtype == cp.dtype("float32"): + cpp_make_blobs( + h_ptr[0], + x_p, + y_p, + n_samples, + n_features, + n_centers, + row_c, + ctr_p, + 0, + cluster_std, + shuffle, + center_box_min, + center_box_max, + random_state, + ) + elif dtype == cp.dtype("float64"): + cpp_make_blobs( + h_ptr[0], + x_p, + y_p, + n_samples, + n_features, + n_centers, + row_c, + ctr_p, + 0, + cluster_std, + shuffle, + center_box_min, + center_box_max, + random_state, + ) + else: + raise ValueError("RAFT make_blobs only supports float32 and float64.") + + return X, y.astype(dtype, copy=False) diff --git a/python/cuml/cuml/datasets/blobs.py b/python/cuml/cuml/datasets/blobs.py index e87076dcb8..672e86d9a6 100644 --- a/python/cuml/cuml/datasets/blobs.py +++ b/python/cuml/cuml/datasets/blobs.py @@ -5,6 +5,7 @@ import numbers from collections.abc import Iterable +from random import getrandbits import cupy as cp import numpy as np @@ -70,6 +71,115 @@ def _get_centers(rs, centers, center_box, n_samples, n_features, dtype): return centers, n_centers +def _make_blobs_raft( + n_samples, + n_features, + centers, + cluster_std, + center_box, + shuffle, + random_state, + return_centers, + order, + dtype, +): + n_samples = int(n_samples) + n_features = int(n_features) + dt = cp.dtype(dtype) + + if n_samples <= 0 or n_features <= 0: + raise ValueError("`n_samples` and `n_features` must be positive.") + + if cluster_std < 0: + raise ValueError("`cluster_std` must be non-negative.") + + gen_ctr = centers is None or isinstance(centers, numbers.Integral) + + if centers is None: + n_ctr, ctr = 3, None + elif isinstance(centers, numbers.Integral): + n_ctr, ctr = int(centers), None + if n_ctr <= 0: + raise ValueError("`centers` must be greater than 0.") + else: + ctr = cp.asarray(centers, dtype=dt, order=order) + + if ctr.ndim != 2: + raise ValueError("`centers` must be a 2D array.") + if ctr.shape[1] != n_features: + raise ValueError( + "Expected `n_features` to be equal to" + " the length of axis 1 of centers array" + ) + + n_ctr = ctr.shape[0] + if n_ctr == 0: + raise ValueError("`centers` must contain at least one center.") + + if gen_ctr: + try: + lo, hi = center_box + except (TypeError, ValueError): + raise ValueError( + "`center_box` must contain exactly two values." + ) from None + + if lo > hi: + raise ValueError( + "`center_box` minimum must not exceed its maximum." + ) + else: + lo = hi = 0.0 + + if return_centers and gen_ctr: + rs = _create_rs_generator(random_state=random_state) + ctr, n_ctr = _get_centers( + rs, + centers, + center_box, + n_samples, + n_features, + dt, + ) + + if ctr is not None: + ctr = cp.asarray(ctr, dtype=dt, order=order) + + if order == "C" and not ctr.flags["C_CONTIGUOUS"]: + ctr = cp.ascontiguousarray(ctr) + elif order == "F" and not ctr.flags["F_CONTIGUOUS"]: + ctr = cp.asfortranarray(ctr) + + if random_state is None: + seed = getrandbits(64) + else: + seed = int(random_state) + + if not 0 <= seed <= (1 << 64) - 1: + raise ValueError("`random_state` must be between 0 and 2**64 - 1.") + + from cuml.datasets._blobs import make_blobs as cpp_blobs + + X, y = cpp_blobs( + n_samples=n_samples, + n_features=n_features, + n_centers=n_ctr, + centers=ctr, + cluster_std=float(cluster_std), + center_box_min=float(lo), + center_box_max=float(hi), + shuffle=bool(shuffle), + random_state=seed, + order=order, + dtype=dt, + ) + + if return_centers: + return X, y, ctr if gen_ctr else centers + + return X, y + + @nvtx.annotate(message="datasets.make_blobs", domain="cuml_python") @cuml.internals.mlfunc(array_arg=None) def make_blobs( @@ -151,6 +261,34 @@ def make_blobs( -------- make_classification: a more intricate variant """ + dt = cp.dtype(dtype) + + use_cpp = ( + isinstance(n_samples, numbers.Integral) + and isinstance(n_features, numbers.Integral) + and n_samples > 0 + and n_features > 0 + and isinstance(cluster_std, numbers.Real) + and isinstance(random_state, (type(None), int)) + and shuffle is True + and order in ("C", "F") + and dt in (cp.dtype("float32"), cp.dtype("float64")) + ) + + if use_cpp: + return _make_blobs_raft( + n_samples=n_samples, + n_features=n_features, + centers=centers, + cluster_std=cluster_std, + center_box=center_box, + shuffle=shuffle, + random_state=random_state, + return_centers=return_centers, + order=order, + dtype=dt, + ) + generator = _create_rs_generator(random_state=random_state) centers, n_centers = _get_centers( diff --git a/python/cuml/tests/test_make_blobs.py b/python/cuml/tests/test_make_blobs.py index eafbd3e71d..5a7bc74954 100644 --- a/python/cuml/tests/test_make_blobs.py +++ b/python/cuml/tests/test_make_blobs.py @@ -1,8 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. +# +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # import cupy as cp +import numpy as np import pytest import cuml @@ -80,3 +82,227 @@ def test_make_blobs_scalar_parameters( assert cp.unique(labels).shape == (centers,), ( "unexpected number of clusters" ) + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize("order", ["F", "C"]) +def test_make_blobs_native_reproducible(dtype, order): + kw = { + "n_samples": 128, + "n_features": 4, + "centers": 4, + "cluster_std": 0.3, + "random_state": 1234, + "order": order, + "dtype": dtype, + } + + x1, y1 = cuml.make_blobs(**kw) + x2, y2 = cuml.make_blobs(**kw) + + cp.testing.assert_array_equal(x1, x2) + cp.testing.assert_array_equal(y1, y2) + + assert x1.dtype == cp.dtype(dtype) + assert y1.dtype == cp.dtype(dtype) + assert x1.flags[f"{order}_CONTIGUOUS"] + assert cp.unique(y1).shape == (4,) + + +def test_make_blobs_matches_native_bridge(): + from cuml.datasets._blobs import make_blobs as raft_blobs + + x1, y1 = cuml.make_blobs( + n_samples=96, + n_features=3, + centers=3, + cluster_std=0.25, + center_box=(-2.0, 5.0), + random_state=2026, + order="C", + dtype="float32", + ) + + x2, y2 = raft_blobs( + n_samples=96, + n_features=3, + n_centers=3, + centers=None, + cluster_std=0.25, + center_box_min=-2.0, + center_box_max=5.0, + shuffle=True, + random_state=2026, + order="C", + dtype=cp.dtype("float32"), + ) + + cp.testing.assert_array_equal(x1, x2) + cp.testing.assert_array_equal(y1, y2) + + +@pytest.mark.parametrize("order", ["F", "C"]) +@pytest.mark.parametrize("xp", [cp, np]) +def test_make_blobs_explicit_centers(order, xp): + ctr = xp.asarray( + [[-2.0, -2.0], [2.0, 2.0]], + dtype=xp.float64, + order="C", + ) + + out, labels, got = cuml.make_blobs( + n_samples=96, + n_features=2, + centers=ctr, + cluster_std=0.1, + random_state=7, + return_centers=True, + order=order, + dtype="float32", + ) + + cp.testing.assert_array_equal(got, cp.asarray(ctr)) + assert got.dtype == ctr.dtype + assert got.flags["C_CONTIGUOUS"] + assert out.flags[f"{order}_CONTIGUOUS"] + assert cp.unique(labels).shape == (2,) + + +def test_make_blobs_empty_centers(): + ctr = cp.empty((0, 2), dtype=cp.float32) + + with pytest.raises(ValueError, match="at least one center"): + cuml.make_blobs( + n_samples=8, + n_features=2, + centers=ctr, + ) + + +@pytest.mark.parametrize("order", ["F", "C"]) +def test_make_blobs_generated_return_centers(order): + from cuml.datasets._blobs import make_blobs as raft_blobs + + x1, y1, ctr = cuml.make_blobs( + n_samples=96, + n_features=2, + centers=2, + cluster_std=0.1, + center_box=(-4.0, 4.0), + random_state=7, + return_centers=True, + order=order, + dtype="float32", + ) + + x2, y2 = raft_blobs( + n_samples=96, + n_features=2, + n_centers=2, + centers=ctr, + cluster_std=0.1, + center_box_min=0.0, + center_box_max=0.0, + shuffle=True, + random_state=7, + order=order, + dtype=cp.dtype("float32"), + ) + + assert ctr.shape == (2, 2) + assert ctr.flags[f"{order}_CONTIGUOUS"] + + cp.testing.assert_array_equal(x1, x2) + cp.testing.assert_array_equal(y1, y2) + + +def test_make_blobs_shuffle_false_keeps_block_labels(): + ctr = cp.asarray( + [[-2.0], [2.0]], + dtype=cp.float32, + ) + + out, labels = cuml.make_blobs( + n_samples=6, + n_features=1, + centers=ctr, + cluster_std=0.0, + shuffle=False, + random_state=7, + order="C", + dtype="float32", + ) + + exp_x = cp.asarray( + [[-2.0], [-2.0], [-2.0], [2.0], [2.0], [2.0]], + dtype=cp.float32, + ) + exp_y = cp.asarray( + [0.0, 0.0, 0.0, 1.0, 1.0, 1.0], + dtype=cp.float32, + ) + + cp.testing.assert_array_equal(out, exp_x) + cp.testing.assert_array_equal(labels, exp_y) + + +def test_make_blobs_cluster_std_sequence_compatibility(): + ctr = cp.asarray( + [[-2.0, -2.0], [2.0, 2.0]], + dtype=cp.float32, + ) + + out, labels = cuml.make_blobs( + n_samples=64, + n_features=2, + centers=ctr, + cluster_std=[0.0, 0.0], + shuffle=True, + random_state=9, + dtype="float32", + ) + + exp = ctr[labels.astype(cp.int64)] + cp.testing.assert_array_equal(out, exp) + + +def test_make_blobs_zero_samples_compatibility(): + with pytest.raises(ValueError): + cuml.make_blobs( + n_samples=0, + n_features=2, + random_state=7, + ) + + +def test_make_blobs_zero_features_compatibility(): + x, y = cuml.make_blobs( + n_samples=8, + n_features=0, + random_state=7, + ) + + assert x.shape == (8, 0) + assert y.shape == (8,) + + +def test_make_blobs_random_state_compatibility(): + kw = { + "n_samples": 64, + "n_features": 3, + "centers": 3, + "cluster_std": 0.2, + "dtype": "float32", + } + + x1, y1 = cuml.make_blobs( + **kw, + random_state=cp.random.RandomState(17), + ) + x2, y2 = cuml.make_blobs( + **kw, + random_state=cp.random.RandomState(17), + ) + + cp.testing.assert_array_equal(x1, x2) + cp.testing.assert_array_equal(y1, y2)