Fix UMAP transform illegal memory access error when data_on_host=True (…

…#6259)

Fixes #6216 by identifying whether the original input data is on host or device and conditionally builds the brute force index (required for a separate `transform()` call) for the correct matrix view.

- [x] Identify and fix root cause
- [x] Clean up implementation
- [x] Implement unit test
- [x] Document fix

Closes #6216

Authors:
  - Simon Adorf (https://github.com/csadorf)
  - Dante Gama Dessavre (https://github.com/dantegd)

Approvers:
  - William Hicks (https://github.com/wphicks)
  - Victor Lafargue (https://github.com/viclafargue)

URL: #6259

cpp/src/umap/knn_graph/algo.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -92,12 +92,22 @@ inline void launcher(const raft::handle_t& handle,
                      cudaStream_t stream)
 {
   if (params->build_algo == ML::UMAPParams::graph_build_algo::BRUTE_FORCE_KNN) {
-    auto idx = cuvs::neighbors::brute_force::build(
-      handle,
-      raft::make_device_matrix_view<const float, int64_t>(inputsA.X, inputsA.n, inputsA.d),
-      params->metric,
-      params->p);
-
+    cudaPointerAttributes attr;
+    RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, inputsA.X));
+    float* ptr = reinterpret_cast<float*>(attr.devicePointer);
+    auto idx   = [&]() {
+      if (ptr != nullptr) {  // inputsA on device
+        return cuvs::neighbors::brute_force::build(
+          handle,
+          {params->metric, params->p},
+          raft::make_device_matrix_view<const float, int64_t>(inputsA.X, inputsA.n, inputsA.d));
+      } else {  // inputsA on host
+        return cuvs::neighbors::brute_force::build(
+          handle,
+          {params->metric, params->p},
+          raft::make_host_matrix_view<const float, int64_t>(inputsA.X, inputsA.n, inputsA.d));
+      }
+    }();
     cuvs::neighbors::brute_force::search(
       handle,
       idx,

python/cuml/cuml/tests/test_umap.py

@@ -842,7 +842,10 @@ def test_umap_distance_metrics_fit_transform_trust_on_sparse_input(
 
 @pytest.mark.parametrize("data_on_host", [True, False])
 @pytest.mark.parametrize("num_clusters", [0, 3, 5])
-def test_umap_trustworthiness_on_batch_nnd(data_on_host, num_clusters):
+@pytest.mark.parametrize("fit_then_transform", [False, True])
+def test_umap_trustworthiness_on_batch_nnd(
+    data_on_host, num_clusters, fit_then_transform
+):
 
     digits = datasets.load_digits()
 
@@ -853,9 +856,15 @@ def test_umap_trustworthiness_on_batch_nnd(data_on_host, num_clusters):
         build_kwds={"nnd_n_clusters": num_clusters},
     )
 
-    cuml_embedding = cuml_model.fit_transform(
-        digits.data, convert_dtype=True, data_on_host=data_on_host
-    )
+    if fit_then_transform:
+        cuml_model.fit(
+            digits.data, convert_dtype=True, data_on_host=data_on_host
+        )
+        cuml_embedding = cuml_model.transform(digits.data)
+    else:
+        cuml_embedding = cuml_model.fit_transform(
+            digits.data, convert_dtype=True, data_on_host=data_on_host
+        )
 
     cuml_trust = trustworthiness(digits.data, cuml_embedding, n_neighbors=10)