diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index e892b5304c..20c3e11ed4 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -356,7 +356,7 @@ cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index); * @param[out] dim return dimension of the index * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim); +cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim); /** * @brief Get size of the CAGRA index @@ -365,7 +365,7 @@ cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim); * @param[out] size return number of vectors in the index * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size); +cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size); /** * @brief Get graph degree of the CAGRA index @@ -374,7 +374,7 @@ cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size); * @param[out] graph_degree return graph degree * @return cuvsError_t */ -cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, uint32_t* graph_degree); +cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree); /** * @brief Returns a view of the CAGRA dataset diff --git a/cpp/include/cuvs/neighbors/ivf_flat.h b/cpp/include/cuvs/neighbors/ivf_flat.h index 9e89ea354d..8211f1ef58 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.h +++ b/cpp/include/cuvs/neighbors/ivf_flat.h @@ -166,10 +166,10 @@ cuvsError_t cuvsIvfFlatIndexCreate(cuvsIvfFlatIndex_t* index); cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index); /** Get the number of clusters/inverted lists */ -uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index); +cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, int64_t* n_lists); /** Get the dimensionality of the data */ -uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index); +cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t* dim); /** * @brief Get the cluster centers corresponding to the lists [n_lists, dim] diff --git a/cpp/include/cuvs/neighbors/ivf_pq.h b/cpp/include/cuvs/neighbors/ivf_pq.h index eef10290c7..3ef7b8b7af 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.h +++ b/cpp/include/cuvs/neighbors/ivf_pq.h @@ -262,10 +262,13 @@ cuvsError_t cuvsIvfPqIndexCreate(cuvsIvfPqIndex_t* index); cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index); /** Get the number of clusters/inverted lists */ -uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index); +cuvsError_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index, int64_t* n_lists); /** Get the dimensionality */ -uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index); +cuvsError_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index, int64_t* dim); + +/** Get the size of the index */ +cuvsError_t cuvsIvfPqIndexGetSize(cuvsIvfPqIndex_t index, int64_t* size); /** * @brief Get the cluster centers corresponding to the lists in the original space diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 7c3061bf93..45a48bbf00 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -426,7 +426,7 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr) }); } -extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim) +extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); @@ -434,7 +434,7 @@ extern "C" cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int* dim) }); } -extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size) +extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); @@ -442,7 +442,7 @@ extern "C" cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* s }); } -extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, uint32_t* graph_degree) +extern "C" cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, int64_t* graph_degree) { return cuvs::core::translate_exceptions([=] { auto index_ptr = reinterpret_cast*>(index->addr); diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp index 91ad9abbed..5d8f7462fa 100644 --- a/cpp/src/neighbors/ivf_flat_c.cpp +++ b/cpp/src/neighbors/ivf_flat_c.cpp @@ -373,50 +373,54 @@ extern "C" cuvsError_t cuvsIvfFlatExtend(cuvsResources_t res, }); } -extern "C" uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index) +extern "C" cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, int64_t* n_lists) { - if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); - } else { - return 0; - } + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); } -extern "C" uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index) +extern "C" cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t* dim) { - if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { - auto index_ptr = - reinterpret_cast*>(index->addr); - return index_ptr->dim(); - } else { - return 0; - } + return cuvs::core::translate_exceptions([=] { + if (index->dtype.code == kDLFloat && index->dtype.bits == 32) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLFloat && index->dtype.bits == 16) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) { + auto index_ptr = + reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + } else { + RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits); + } + }); } extern "C" cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, diff --git a/cpp/src/neighbors/ivf_pq_c.cpp b/cpp/src/neighbors/ivf_pq_c.cpp index 4abd080582..5c10d4111b 100644 --- a/cpp/src/neighbors/ivf_pq_c.cpp +++ b/cpp/src/neighbors/ivf_pq_c.cpp @@ -347,16 +347,28 @@ extern "C" cuvsError_t cuvsIvfPqExtend(cuvsResources_t res, }); } -extern "C" uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index) +extern "C" cuvsError_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index, int64_t* n_lists) { - auto index_ptr = reinterpret_cast*>(index->addr); - return index_ptr->n_lists(); + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *n_lists = index_ptr->n_lists(); + }); } -extern "C" uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index) +extern "C" cuvsError_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index, int64_t* dim) { - auto index_ptr = reinterpret_cast*>(index->addr); - return index_ptr->dim(); + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *dim = index_ptr->dim(); + }); +} + +extern "C" cuvsError_t cuvsIvfPqIndexGetSize(cuvsIvfPqIndex_t index, int64_t* size) +{ + return cuvs::core::translate_exceptions([=] { + auto index_ptr = reinterpret_cast*>(index->addr); + *size = index_ptr->size(); + }); } extern "C" cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor* centers) diff --git a/cpp/tests/neighbors/ann_cagra_c.cu b/cpp/tests/neighbors/ann_cagra_c.cu index ae80cc8986..a11ff01be5 100644 --- a/cpp/tests/neighbors/ann_cagra_c.cu +++ b/cpp/tests/neighbors/ann_cagra_c.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-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. @@ -541,7 +541,7 @@ TEST(CagraC, BuildMergeSearch) cuvsCagraIndex_t index_array[2] = {index_main, index_add}; ASSERT_EQ(cuvsCagraMerge(res, merge_params, index_array, 2, index_merged), CUVS_SUCCESS); - int merged_dim = -1; + int64_t merged_dim = -1; ASSERT_EQ(cuvsCagraIndexGetDims(index_merged, &merged_dim), CUVS_SUCCESS); EXPECT_EQ(merged_dim, 2); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index aa082482dc..99b0ecc71e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -36,6 +36,7 @@ import com.nvidia.cuvs.internal.common.CloseableRMMAllocation; import com.nvidia.cuvs.internal.common.CompositeCloseableHandle; import com.nvidia.cuvs.internal.panama.*; + import java.io.FileInputStream; import java.io.InputStream; import java.io.OutputStream; diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 28dfc5c020..1ffb5ad559 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -106,13 +106,6 @@ public static void checkCudaError(int value, String caller) { } } - private static final long UNSIGNED_INT_MASK = 0xFFFFFFFFL; - - public static long dereferenceUnsignedInt(MemorySegment ptr) { - assert ptr.byteSize() == 4; - return ptr.get(uint32_t, 0) & UNSIGNED_INT_MASK; - } - /** * Java analog to CUDA's cudaMemcpyKind, used for cudaMemcpy() calls. * @see CUDA Runtime API diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 29e47a5014..9b62f71b2c 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -134,10 +134,10 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index) - cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int32_t* dim) - cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, uint32_t* size) + cuvsError_t cuvsCagraIndexGetDims(cuvsCagraIndex_t index, int64_t* dim) + cuvsError_t cuvsCagraIndexGetSize(cuvsCagraIndex_t index, int64_t* size) cuvsError_t cuvsCagraIndexGetGraphDegree(cuvsCagraIndex_t index, - uint32_t* degree) + int64_t* degree) cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor * graph) cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 0a8fec17b4..8d47a1d908 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -275,18 +275,18 @@ cdef class Index: @property def dim(self): - cdef int32_t dim + cdef int64_t dim check_cuvs(cuvsCagraIndexGetDims(self.index, &dim)) return dim @property def graph_degree(self): - cdef uint32_t degree + cdef int64_t degree check_cuvs(cuvsCagraIndexGetGraphDegree(self.index, °ree)) return degree def __len__(self): - cdef uint32_t size + cdef int64_t size check_cuvs(cuvsCagraIndexGetSize(self.index, &size)) return size diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd index 6e30ac0a07..470234d9cf 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pxd @@ -15,7 +15,7 @@ # # cython: language_level=3 -from libc.stdint cimport uint32_t, uintptr_t +from libc.stdint cimport int64_t, uint32_t, uintptr_t from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t @@ -62,9 +62,10 @@ cdef extern from "cuvs/neighbors/ivf_flat.h" nogil: cuvsError_t cuvsIvfFlatIndexDestroy(cuvsIvfFlatIndex_t index) - uint32_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index) + cuvsError_t cuvsIvfFlatIndexGetNLists(cuvsIvfFlatIndex_t index, + int64_t * n_lists) - uint32_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index) + cuvsError_t cuvsIvfFlatIndexGetDim(cuvsIvfFlatIndex_t index, int64_t * dim) cuvsError_t cuvsIvfFlatIndexGetCenters(cuvsIvfFlatIndex_t index, DLManagedTensor * centers) diff --git a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx index 3a803e5599..996a39d4fe 100644 --- a/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_flat/ivf_flat.pyx @@ -188,12 +188,16 @@ cdef class Index: @property def n_lists(self): """ The number of inverted lists (clusters) """ - return cuvsIvfFlatIndexGetNLists(self.index) + cdef int64_t n_lists = 0 + cuvsIvfFlatIndexGetNLists(self.index, &n_lists) + return n_lists @property def dim(self): """ dimensionality of the cluster centers """ - return cuvsIvfFlatIndexGetDim(self.index) + cdef int64_t dim = 0 + cuvsIvfFlatIndexGetDim(self.index, &dim) + return dim @property def centers(self): diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd index 897e4951a9..5ca7b97602 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pxd @@ -15,7 +15,7 @@ # # cython: language_level=3 -from libc.stdint cimport uint32_t, uintptr_t +from libc.stdint cimport int64_t, uint32_t, uintptr_t from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t @@ -84,9 +84,12 @@ cdef extern from "cuvs/neighbors/ivf_pq.h" nogil: cuvsError_t cuvsIvfPqIndexDestroy(cuvsIvfPqIndex_t index) - uint32_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index) + cuvsError_t cuvsIvfPqIndexGetNLists(cuvsIvfPqIndex_t index, + int64_t * n_lists) - uint32_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index) + cuvsError_t cuvsIvfPqIndexGetDim(cuvsIvfPqIndex_t index, int64_t * dim) + + cuvsError_t cuvsIvfPqIndexGetSize(cuvsIvfPqIndex_t index, int64_t * size) cuvsError_t cuvsIvfPqIndexGetCenters(cuvsIvfPqIndex_t index, DLManagedTensor * centers) diff --git a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx index f784327d33..5ffa0b0c88 100644 --- a/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/cuvs/cuvs/neighbors/ivf_pq/ivf_pq.pyx @@ -249,12 +249,21 @@ cdef class Index: @property def n_lists(self): """ The number of inverted lists (clusters) """ - return cuvsIvfPqIndexGetNLists(self.index) + cdef int64_t n_lists + check_cuvs(cuvsIvfPqIndexGetNLists(self.index, &n_lists)) + return n_lists @property def dim(self): """ dimensionality of the cluster centers """ - return cuvsIvfPqIndexGetDim(self.index) + cdef int64_t dim + check_cuvs(cuvsIvfPqIndexGetDim(self.index, &dim)) + return dim + + def __len__(self): + cdef int64_t size + check_cuvs(cuvsIvfPqIndexGetSize(self.index, &size)) + return size @property def centers(self):