add nn_descent #325

Intron7 · 2025-02-07T12:11:43Z

This PR will add cuvs.neighbors.nn_descent to rsc. However nn_descent doesn't add self. nn_descent will be supported in cuvs 25.02

Intron7 · 2025-02-11T10:07:28Z

After discussing with @flying-sheep we decided to add self to nn_descent.

ilan-gold

Also, why not use string templating for kernels?

Something like

dist_kernel = kernel_code_inner = r"""
extern "C" __global__
void computeDistances_{NAME}(const float* pca,
                    float* out,
                    const unsigned int* pairs,
                    const long long int n_samples,
                    const long long int n_pcs,
                    const long long int n_neighbors)
{
    long long int i1 = blockDim.x * blockIdx.x + threadIdx.x;
    if(i1 >= n_samples){
        return;
    }
    {PREPARE}
    for (long long int j = 0; j < n_neighbors; j++){
        {OUTER_LOOP}
        for (long long int d = 0; d < n_pcs; d++) {
            {INNER_LOOP}
        }
        out[i1 * n_neighbors + j] = {DIST_OP}
    }
}
"""

and then replacing the variables as needed?

dist_kernel.format(name="cosine", ...)

This'd make adding new kernels hopefully less work going forward and also make things a bit easier to read, I think.

ilan-gold · 2025-02-11T13:18:41Z

src/rapids_singlecell/preprocessing/_kernels/_nn_descent.py

+
+kernel_code = r"""
+extern "C" __global__
+void computeDistances(const float* pca,


does this only work with pca? should it be called something else? it looks like just normal euclidean distance

(this applies everywhere about argument names being pca etc.)

ilan-gold · 2025-02-11T13:21:50Z

src/rapids_singlecell/preprocessing/_neighbors.py

+    neighbors = cp.asarray(neighbors)
+    distances = cp.asarray(distances)


just out of curiosity, what are these if not cupy arrays?

raft device arrays. You cant slice them or anything they are very bare bones

src/rapids_singlecell/preprocessing/_neighbors.py

ilan-gold · 2025-02-11T13:26:01Z

src/rapids_singlecell/preprocessing/_neighbors.py

+    if metric == "euclidean":
+        metric_to_use = "sqeuclidean"
+    else:
+        metric_to_use = metric
+    idxparams = nn_descent.IndexParams(graph_degree=k, metric=metric_to_use)


Suggested change

if metric == "euclidean":

metric_to_use = "sqeuclidean"

else:

metric_to_use = metric

idxparams = nn_descent.IndexParams(graph_degree=k, metric=metric_to_use)

idxparams = nn_descent.IndexParams(graph_degree=k, metric="sqeuclidean" if metric == "euclidean" else metric)

Or why not just

Suggested change

if metric == "euclidean":

metric_to_use = "sqeuclidean"

else:

metric_to_use = metric

idxparams = nn_descent.IndexParams(graph_degree=k, metric=metric_to_use)

if metric == "euclidean":

metric = "sqeuclidean"

idxparams = nn_descent.IndexParams(graph_degree=k, metric=metric)

and then don't change the others?

I later need to know if the metric is sqeuclidean or euclidean to do the sqrt

Ilan’s first suggestion still applies

src/rapids_singlecell/preprocessing/_neighbors.py

ilan-gold · 2025-02-11T13:26:54Z

src/rapids_singlecell/preprocessing/_neighbors.py

+    )
+    if metric == "euclidean":
+        distances = cp.sqrt(distances)
+    if metric in ("cosine", "euclidean", "sqeuclidean"):


can't be sqeuclidean?

Suggested change

if metric in ("cosine", "euclidean", "sqeuclidean"):

if metric in ("cosine", "euclidean"):

Confused logic? How is the possible? see above?

src/rapids_singlecell/preprocessing/_neighbors.py

ilan-gold · 2025-02-11T13:46:26Z

src/rapids_singlecell/preprocessing/_kernels/_nn_descent.py

+
+    float sum_i1 = 0.0f;
+    for (long long int d = 0; d < n_pcs; d++) {
+        sum_i1 += pca[i1 * n_pcs + d] * pca[i1 * n_pcs + d];


Is https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__SINGLE.html?highlight=pow#_CPPv44powfff not feasible? Seems like two accesses to the same thing is not great, but would be very curious to find out if it actually is fine (i.e., how cheap this is).

Same thing below re: data access.

Co-authored-by: Ilan Gold <[email protected]>

flying-sheep

Looks pretty good, but you should document that inner_product isn’t an actual distance metric.

Talking about distances and metrics could lead users into thinking that it is one unless you add a warning about it.

flying-sheep · 2025-02-11T16:10:53Z

src/rapids_singlecell/preprocessing/_kernels/_nn_descent.py

+                    const unsigned int* pairs,
+                    const long long int n_samples,


Suggested change

const unsigned int* pairs,

const long long int n_samples,

const std::uint32_t* pairs,

const std::uint64_t n_samples,

and so on

flying-sheep · 2025-02-11T16:15:38Z

src/rapids_singlecell/preprocessing/_kernels/_nn_descent.py

+}
+"""
+
+# Create the RawKernel object


remove these comments, they are obvious. As said: Comments should contain the “why”, the code should contain the “what”. ... = cp.RawKernel(...) already says “create the RawKernel object” loudly.

flying-sheep · 2025-02-11T16:17:51Z

src/rapids_singlecell/preprocessing/_neighbors.py

@@ -21,7 +21,7 @@
    from anndata import AnnData

 AnyRandom = None | int | np.random.RandomState
-_Alogithms = Literal["brute", "ivfflat", "ivfpq", "cagra"]
+_Alogithms = Literal[_Alogithms_bbknn, Literal["nn_descent"]]


Suggested change

_Alogithms = Literal[_Alogithms_bbknn, Literal["nn_descent"]]

_Algorithms = _Algorithms_bbknn | Literal["nn_descent"]

or Union if you have to

flying-sheep · 2025-02-11T16:19:55Z

src/rapids_singlecell/preprocessing/_neighbors.py

+    if metric == "euclidean":
+        metric_to_use = "sqeuclidean"
+    else:
+        metric_to_use = metric
+    idxparams = nn_descent.IndexParams(graph_degree=k, metric=metric_to_use)


Ilan’s first suggestion still applies

flying-sheep · 2025-02-11T16:21:14Z

src/rapids_singlecell/preprocessing/_neighbors.py

+        add_self_neighbors = cp.arange(X.shape[0], dtype=cp.uint32)
+        neighbors = cp.concatenate(
+            (add_self_neighbors[:, None], neighbors[:, :-1]), axis=1
+        )
+        add_self_distances = cp.zeros((X.shape[0], 1), dtype=cp.float32)
+        distances = cp.concatenate((add_self_distances, distances[:, :-1]), axis=1)


Add a comment here explaining why you add self.

A valid comment would e.g. be something like

API xyz expects the kNN indices to contain a self column. The real distance metrics don‘t add this column since all self-distances are 0 in distance metrics, but inner_product, which isn’t an actual distance metric as it has non-0 values in the self column, does add it.

Co-authored-by: Ilan Gold <[email protected]>

add nn_descent

518764e

Intron7 added the run-gpu-ci runs GPU CI label Feb 7, 2025

Intron7 requested a review from ilan-gold February 7, 2025 12:12

github-actions bot removed the run-gpu-ci runs GPU CI label Feb 7, 2025

fix Literal

caeb583

Intron7 added the run-gpu-ci runs GPU CI label Feb 7, 2025

github-actions bot removed the run-gpu-ci runs GPU CI label Feb 7, 2025

Intron7 and others added 3 commits February 11, 2025 11:05

add self

d304459

add test

fdb5932

Merge branch 'main' into add-nn_descent

034d58b

ilan-gold requested changes Feb 11, 2025

View reviewed changes

Intron7 and others added 2 commits February 11, 2025 17:03

Update src/rapids_singlecell/preprocessing/_neighbors.py

d619faa

Co-authored-by: Ilan Gold <[email protected]>

Update src/rapids_singlecell/preprocessing/_neighbors.py

061d333

Co-authored-by: Ilan Gold <[email protected]>

flying-sheep reviewed Feb 11, 2025

View reviewed changes

Update src/rapids_singlecell/preprocessing/_neighbors.py

d1a8ab0

Co-authored-by: Ilan Gold <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add nn_descent #325

add nn_descent #325

Intron7 commented Feb 7, 2025 •

edited

Loading

Intron7 commented Feb 11, 2025

ilan-gold left a comment

ilan-gold Feb 11, 2025

ilan-gold Feb 11, 2025

ilan-gold Feb 11, 2025

Intron7 Feb 11, 2025

ilan-gold Feb 11, 2025

ilan-gold Feb 11, 2025

Intron7 Feb 11, 2025

flying-sheep Feb 11, 2025 •

edited

Loading

ilan-gold Feb 11, 2025

ilan-gold Feb 11, 2025

ilan-gold Feb 11, 2025

ilan-gold Feb 11, 2025

flying-sheep left a comment

flying-sheep Feb 11, 2025

flying-sheep Feb 11, 2025 •

edited

Loading

flying-sheep Feb 11, 2025

flying-sheep Feb 11, 2025 •

edited

Loading

flying-sheep Feb 11, 2025

		neighbors = cp.asarray(neighbors)
		distances = cp.asarray(distances)

	if metric in ("cosine", "euclidean", "sqeuclidean"):
	if metric in ("cosine", "euclidean"):

	_Alogithms = Literal[_Alogithms_bbknn, Literal["nn_descent"]]
	_Algorithms = _Algorithms_bbknn \| Literal["nn_descent"]

add nn_descent #325

Are you sure you want to change the base?

add nn_descent #325

Conversation

Intron7 commented Feb 7, 2025 • edited Loading

Intron7 commented Feb 11, 2025

ilan-gold left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flying-sheep Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flying-sheep left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flying-sheep Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

flying-sheep Feb 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Intron7 commented Feb 7, 2025 •

edited

Loading

flying-sheep Feb 11, 2025 •

edited

Loading

flying-sheep Feb 11, 2025 •

edited

Loading

flying-sheep Feb 11, 2025 •

edited

Loading