add docstrings

rapidsai · ajit283 · Jun 14, 2024 · Jun 14, 2024 · Jun 19, 2024 · Jun 19, 2024
commit f7fac3586ab9d37564141061236a27fa7ffc2f51
diff --git a/examples/c/CMakeLists.txt b/examples/c/CMakeLists.txt
@@ -13,7 +13,6 @@
 
 cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)
 
-set(CMAKE_CUDA_ARCHITECTURES 89)
 
 # ------------- configure rapids-cmake --------------#
 

diff --git a/go/brute_force/brute_force.go b/go/brute_force/brute_force.go
@@ -10,31 +10,42 @@ import (
 	cuvs "github.com/rapidsai/cuvs/go"
 )
 
-type bruteForceIndex struct {
+// Brute Force KNN Index
+type BruteForceIndex struct {
 	index   C.cuvsBruteForceIndex_t
 	trained bool
 }
 
-func CreateIndex() (*bruteForceIndex, error) {
+// Creates a new empty Brute Force KNN Index
+func CreateIndex() (*BruteForceIndex, error) {
 	var index C.cuvsBruteForceIndex_t
 
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsBruteForceIndexCreate(&index)))
 	if err != nil {
 		return nil, err
 	}
 
-	return &bruteForceIndex{index: index, trained: false}, nil
+	return &BruteForceIndex{index: index, trained: false}, nil
 }
 
-func (index *bruteForceIndex) Close() error {
+// Destroys the Brute Force KNN Index
+func (index *BruteForceIndex) Close() error {
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsBruteForceIndexDestroy(index.index)))
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
-func BuildIndex[T any](Resources cuvs.Resource, Dataset *cuvs.Tensor[T], metric cuvs.Distance, metric_arg float32, index *bruteForceIndex) error {
+// Builds a new Brute Force KNN Index from the dataset for efficient search.
+//
+// # Arguments
+//
+// * `Resources` - Resources to use
+// * `Dataset` - A row-major matrix on either the host or device to index
+// * `metric` - Distance type to use for building the index
+// * `metric_arg` - Value of `p` for Minkowski distances - set to 2.0 if not applicable
+func BuildIndex[T any](Resources cuvs.Resource, Dataset *cuvs.Tensor[T], metric cuvs.Distance, metric_arg float32, index *BruteForceIndex) error {
 	CMetric, exists := cuvs.CDistances[metric]
 
 	if !exists {
@@ -50,7 +61,15 @@ func BuildIndex[T any](Resources cuvs.Resource, Dataset *cuvs.Tensor[T], metric
 	return nil
 }
 
-func SearchIndex[T any](resources cuvs.Resource, index bruteForceIndex, queries *cuvs.Tensor[T], neighbors *cuvs.Tensor[int64], distances *cuvs.Tensor[T]) error {
+// Perform a Nearest Neighbors search on the Index
+//
+// # Arguments
+//
+// * `Resources` - Resources to use
+// * `queries` - Tensor in device memory to query for
+// * `neighbors` - Tensor in device memory that receives the indices of the nearest neighbors
+// * `distances` - Tensor in device memory that receives the distances of the nearest neighbors
+func SearchIndex[T any](resources cuvs.Resource, index BruteForceIndex, queries *cuvs.Tensor[T], neighbors *cuvs.Tensor[int64], distances *cuvs.Tensor[T]) error {
 	if !index.trained {
 		return errors.New("index needs to be built before calling search")
 	}

diff --git a/go/cagra/cagra.go b/go/cagra/cagra.go
@@ -10,11 +10,13 @@ import (
 	cuvs "github.com/rapidsai/cuvs/go"
 )
 
+// Cagra ANN Index
 type CagraIndex struct {
 	index   C.cuvsCagraIndex_t
 	trained bool
 }
 
+// Creates a new empty Cagra Index
 func CreateIndex() (*CagraIndex, error) {
 	var index C.cuvsCagraIndex_t
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexCreate(&index)))
@@ -25,6 +27,14 @@ func CreateIndex() (*CagraIndex, error) {
 	return &CagraIndex{index: index}, nil
 }
 
+// Builds a new Index from the dataset for efficient search.
+//
+// # Arguments
+//
+// * `Resources` - Resources to use
+// * `params` - Parameters for building the index
+// * `dataset` - A row-major Tensor on either the host or device to index
+// * `index` - CagraIndex to build
 func BuildIndex[T any](Resources cuvs.Resource, params *IndexParams, dataset *cuvs.Tensor[T], index *CagraIndex) error {
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraBuild(C.ulong(Resources.Resource), params.params, (*C.DLManagedTensor)(unsafe.Pointer(dataset.C_tensor)), index.index)))
 	if err != nil {
@@ -34,6 +44,15 @@ func BuildIndex[T any](Resources cuvs.Resource, params *IndexParams, dataset *cu
 	return nil
 }
 
+// Extends the index with additional data
+//
+// # Arguments
+//
+// * `Resources` - Resources to use
+// * `params` - Parameters for extending the index
+// * `additional_dataset` - A row-major Tensor on the device to extend the index with
+// * `return_dataset` - A row-major Tensor on the device that will receive the extended dataset
+// * `index` - CagraIndex to extend
 func ExtendIndex[T any](Resources cuvs.Resource, params *ExtendParams, additional_dataset *cuvs.Tensor[T], return_dataset *cuvs.Tensor[T], index *CagraIndex) error {
 	if !index.trained {
 		return errors.New("index needs to be built before calling extend")
@@ -45,6 +64,7 @@ func ExtendIndex[T any](Resources cuvs.Resource, params *ExtendParams, additiona
 	return nil
 }
 
+// Destroys the Cagra Index
 func (index *CagraIndex) Close() error {
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexDestroy(index.index)))
 	if err != nil {
@@ -53,6 +73,15 @@ func (index *CagraIndex) Close() error {
 	return nil
 }
 
+// Perform a Approximate Nearest Neighbors search on the Index
+//
+// # Arguments
+//
+// * `Resources` - Resources to use
+// * `params` - Parameters to use in searching the index
+// * `queries` - A tensor in device memory to query for
+// * `neighbors` - Tensor in device memory that receives the indices of the nearest neighbors
+// * `distances` - Tensor in device memory that receives the distances of the nearest neighbors
 func SearchIndex[T any](Resources cuvs.Resource, params *SearchParams, index *CagraIndex, queries *cuvs.Tensor[T], neighbors *cuvs.Tensor[uint32], distances *cuvs.Tensor[T]) error {
 	if !index.trained {
 		return errors.New("index needs to be built before calling search")

diff --git a/go/cagra/extend_params.go b/go/cagra/extend_params.go
@@ -7,10 +7,12 @@ import (
 	cuvs "github.com/rapidsai/cuvs/go"
 )
 
+// Parameters to extend CAGRA Index
 type ExtendParams struct {
 	params C.cuvsCagraExtendParams_t
 }
 
+// Creates a new ExtendParams
 func CreateExtendParams() (*ExtendParams, error) {
 	var params C.cuvsCagraExtendParams_t
 
@@ -24,6 +26,12 @@ func CreateExtendParams() (*ExtendParams, error) {
 	return ExtendParams, nil
 }
 
+// The additional dataset is divided into chunks and added to the graph.
+// This is the knob to adjust the tradeoff between the recall and operation throughput.
+// Large chunk sizes can result in high throughput, but use more
+// working memory (O(max_chunk_size*degree^2)).
+// This can also degrade recall because no edges are added between the nodes in the same chunk.
+// Auto select when 0.
 func (p *ExtendParams) SetMaxChunkSize(max_chunk_size uint32) (*ExtendParams, error) {
 	p.params.max_chunk_size = C.uint32_t(max_chunk_size)
 	return p, nil

diff --git a/go/cagra/index_params.go b/go/cagra/index_params.go
@@ -13,6 +13,7 @@ type IndexParams struct {
 	params C.cuvsCagraIndexParams_t
 }
 
+// Supplemental parameters to build CAGRA Index
 type CompressionParams struct {
 	params C.cuvsCagraCompressionParams_t
 }
@@ -31,6 +32,7 @@ var cBuildAlgos = map[BuildAlgo]int{
 	AutoSelect: C.AUTO_SELECT,
 }
 
+// Creates a new CompressionParams
 func CreateCompressionParams() (*CompressionParams, error) {
 	var params C.cuvsCagraCompressionParams_t
 
@@ -46,42 +48,54 @@ func CreateCompressionParams() (*CompressionParams, error) {
 	return &CompressionParams{params: params}, nil
 }
 
+// The bit length of the vector element after compression by PQ.
 func (p *CompressionParams) SetPQBits(pq_bits uint32) (*CompressionParams, error) {
 	p.params.pq_bits = C.uint32_t(pq_bits)
 
 	return p, nil
 }
 
+// The dimensionality of the vector after compression by PQ. When zero,
+// an optimal value is selected using a heuristic.
 func (p *CompressionParams) SetPQDim(pq_dim uint32) (*CompressionParams, error) {
 	p.params.pq_dim = C.uint32_t(pq_dim)
 
 	return p, nil
 }
 
+// Vector Quantization (VQ) codebook size - number of "coarse cluster
+// centers". When zero, an optimal value is selected using a heuristic.
 func (p *CompressionParams) SetVQNCenters(vq_n_centers uint32) (*CompressionParams, error) {
 	p.params.vq_n_centers = C.uint32_t(vq_n_centers)
 
 	return p, nil
 }
 
+// The number of iterations searching for kmeans centers (both VQ & PQ
+// phases).
 func (p *CompressionParams) SetKMeansNIters(kmeans_n_iters uint32) (*CompressionParams, error) {
 	p.params.kmeans_n_iters = C.uint32_t(kmeans_n_iters)
 
 	return p, nil
 }
 
+// The fraction of data to use during iterative kmeans building (VQ
+// phase). When zero, an optimal value is selected using a heuristic.
 func (p *CompressionParams) SetVQKMeansTrainsetFraction(vq_kmeans_trainset_fraction float64) (*CompressionParams, error) {
 	p.params.vq_kmeans_trainset_fraction = C.double(vq_kmeans_trainset_fraction)
 
 	return p, nil
 }
 
+// The fraction of data to use during iterative kmeans building (PQ
+// phase). When zero, an optimal value is selected using a heuristic.
 func (p *CompressionParams) SetPQKMeansTrainsetFraction(pq_kmeans_trainset_fraction float64) (*CompressionParams, error) {
 	p.params.pq_kmeans_trainset_fraction = C.double(pq_kmeans_trainset_fraction)
 
 	return p, nil
 }
 
+// Creates a new IndexParams
 func CreateIndexParams() (*IndexParams, error) {
 	var params C.cuvsCagraIndexParams_t
 
@@ -95,17 +109,20 @@ func CreateIndexParams() (*IndexParams, error) {
 	return IndexParams, nil
 }
 
+// Degree of input graph for pruning
 func (p *IndexParams) SetIntermediateGraphDegree(intermediate_graph_degree uintptr) (*IndexParams, error) {
 	p.params.intermediate_graph_degree = C.size_t(intermediate_graph_degree)
 	return p, nil
 }
 
+// Degree of output graph
 func (p *IndexParams) SetGraphDegree(intermediate_graph_degree uintptr) (*IndexParams, error) {
 	p.params.graph_degree = C.size_t(intermediate_graph_degree)
 
 	return p, nil
 }
 
+// ANN algorithm to build knn graph
 func (p *IndexParams) SetBuildAlgo(build_algo BuildAlgo) (*IndexParams, error) {
 	CBuildAlgo, exists := cBuildAlgos[build_algo]
 
@@ -117,18 +134,21 @@ func (p *IndexParams) SetBuildAlgo(build_algo BuildAlgo) (*IndexParams, error) {
 	return p, nil
 }
 
+// Number of iterations to run if building with NN_DESCENT
 func (p *IndexParams) SetNNDescentNiter(nn_descent_niter uint32) (*IndexParams, error) {
 	p.params.nn_descent_niter = C.ulong(nn_descent_niter)
 
 	return p, nil
 }
 
+// Compression parameters
 func (p *IndexParams) SetCompression(compression *CompressionParams) (*IndexParams, error) {
 	p.params.compression = C.cuvsCagraCompressionParams_t(compression.params)
 
 	return p, nil
 }
 
+// Destroys IndexParams
 func (p *IndexParams) Close() error {
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexParamsDestroy(p.params)))
 	if err != nil {

diff --git a/go/cagra/search_params.go b/go/cagra/search_params.go
@@ -9,6 +9,7 @@ import (
 	cuvs "github.com/rapidsai/cuvs/go"
 )
 
+// Supplemental parameters to search CAGRA Index
 type SearchParams struct {
 	params C.cuvsCagraSearchParams_t
 }
@@ -30,6 +31,7 @@ const (
 	HashmapModeAuto
 )
 
+// Creates a new SearchParams
 func CreateSearchParams() (*SearchParams, error) {
 	var params C.cuvsCagraSearchParams_t
 
@@ -41,21 +43,27 @@ func CreateSearchParams() (*SearchParams, error) {
 	return &SearchParams{params: params}, nil
 }
 
+// Maximum number of queries to search at the same time (batch size). Auto select when 0
 func (p *SearchParams) SetMaxQueries(max_queries uintptr) (*SearchParams, error) {
 	p.params.max_queries = C.size_t(max_queries)
 	return p, nil
 }
 
+// Number of intermediate search results retained during the search.
+// This is the main knob to adjust trade off between accuracy and search speed.
+// Higher values improve the search accuracy
 func (p *SearchParams) SetItopkSize(itopk_size uintptr) (*SearchParams, error) {
 	p.params.itopk_size = C.size_t(itopk_size)
 	return p, nil
 }
 
+// Upper limit of search iterations. Auto select when 0.
 func (p *SearchParams) SetMaxIterations(max_iterations uintptr) (*SearchParams, error) {
 	p.params.max_iterations = C.size_t(max_iterations)
 	return p, nil
 }
 
+// Which search implementation to use.
 func (p *SearchParams) SetAlgo(algo SearchAlgo) (*SearchParams, error) {
 	CAlgo := C.SINGLE_CTA
 
@@ -77,26 +85,31 @@ func (p *SearchParams) SetAlgo(algo SearchAlgo) (*SearchParams, error) {
 	return p, nil
 }
 
+// Number of threads used to calculate a single distance. 4, 8, 16, or 32.
 func (p *SearchParams) SetTeamSize(team_size uintptr) (*SearchParams, error) {
 	p.params.team_size = C.size_t(team_size)
 	return p, nil
 }
 
+// Lower limit of search iterations.
 func (p *SearchParams) SetMinIterations(min_iterations uintptr) (*SearchParams, error) {
 	p.params.min_iterations = C.size_t(min_iterations)
 	return p, nil
 }
 
+// How many nodes to search at once. Auto select when 0.
 func (p *SearchParams) SetSearchWidth(search_width uintptr) (*SearchParams, error) {
 	p.params.search_width = C.size_t(search_width)
 	return p, nil
 }
 
+// Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0.
 func (p *SearchParams) SetThreadBlockSize(thread_block_size uintptr) (*SearchParams, error) {
 	p.params.thread_block_size = C.size_t(thread_block_size)
 	return p, nil
 }
 
+// Hashmap type. Auto selection when AUTO.
 func (p *SearchParams) SetHashmapMode(hashmap_mode HashmapMode) (*SearchParams, error) {
 	CHashMode := C.AUTO_HASH
 
@@ -116,26 +129,31 @@ func (p *SearchParams) SetHashmapMode(hashmap_mode HashmapMode) (*SearchParams,
 	return p, nil
 }
 
+// Lower limit of hashmap bit length. More than 8.
 func (p *SearchParams) SetHashmapMinBitlen(hashmap_min_bitlen uintptr) (*SearchParams, error) {
 	p.params.hashmap_min_bitlen = C.size_t(hashmap_min_bitlen)
 	return p, nil
 }
 
+// Upper limit of hashmap fill rate. More than 0.1, less than 0.9.
 func (p *SearchParams) SetHashmapMaxFillRate(hashmap_max_fill_rate float32) (*SearchParams, error) {
 	p.params.hashmap_max_fill_rate = C.float(hashmap_max_fill_rate)
 	return p, nil
 }
 
+// Number of iterations of initial random seed node selection. 1 or more.
 func (p *SearchParams) SetNumRandomSamplings(num_random_samplings uint32) (*SearchParams, error) {
 	p.params.num_random_samplings = C.uint32_t(num_random_samplings)
 	return p, nil
 }
 
+// Bit mask used for initial random seed node selection.
 func (p *SearchParams) SetRandXorMask(rand_xor_mask uint64) (*SearchParams, error) {
 	p.params.rand_xor_mask = C.uint64_t(rand_xor_mask)
 	return p, nil
 }
 
+// Destroys SearchParams
 func (p *SearchParams) Close() error {
 	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraSearchParamsDestroy(p.params)))
 	if err != nil {
Original file line number	Diff line number	Diff line change
Expand Up		@@ -13,7 +13,6 @@

		cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)

		set(CMAKE_CUDA_ARCHITECTURES 89)

		# ------------- configure rapids-cmake --------------#

Expand Down