Go API - [WIP]

build.sh

@@ -18,14 +18,15 @@ ARGS=$*
 # scripts, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libcuvs python rust docs tests bench-ann examples --uninstall  -v -g -n --compile-static-lib --allgpuarch --no-mg --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
+VALIDARGS="clean libcuvs python rust go docs tests bench-ann examples --uninstall  -v -g -n --compile-static-lib --allgpuarch --no-mg --no-cpu --cpu-only --no-shared-libs --no-nvtx --show_depr_warn --incl-cache-stats --time -h"
 HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>] [--limit-tests=<targets>] [--limit-bench-ann=<targets>] [--build-metrics=<filename>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libcuvs          - build the cuvs C++ code only. Also builds the C-wrapper library
                       around the C++ code.
    python           - build the cuvs Python package
    rust             - build the cuvs Rust bindings
+   go               - build the cuvs Go bindings
    docs             - build the documentation
    tests            - build the tests
    bench-ann        - build end-to-end ann benchmarks
@@ -459,6 +460,13 @@ if (( ${NUMARGS} == 0 )) || hasArg rust; then
     cargo test
 fi
 
+# Build the cuvs Go bindings
+if (( ${NUMARGS} == 0 )) || hasArg go; then
+    cd ${REPODIR}/go
+    go build ./...
+    go test ./...
+fi
+
 export RAPIDS_VERSION="$(sed -E -e 's/^([0-9]{2})\.([0-9]{2})\.([0-9]{2}).*$/\1.\2.\3/' "${REPODIR}/VERSION")"
 export RAPIDS_VERSION_MAJOR_MINOR="$(sed -E -e 's/^([0-9]{2})\.([0-9]{2})\.([0-9]{2}).*$/\1.\2/' "${REPODIR}/VERSION")"
 

ci/build_go.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+rapids-logger "Create test conda environment"
+. /opt/conda/etc/profile.d/conda.sh
+
+RAPIDS_VERSION="$(rapids-version)"
+
+rapids-dependency-file-generator \
+  --output conda \
+  --file-key go \
+  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" | tee env.yaml
+
+rapids-mamba-retry env create --yes -f env.yaml -n go
+
+# seeing failures on activating the environment here on unbound locals
+# apply workaround from https://github.com/conda/conda/issues/8186#issuecomment-532874667
+set +eu
+conda activate go
+set -eu
+
+rapids-print-env
+
+export CGO_CFLAGS="-I/usr/local/cuda/include -I/home/ajit/miniforge3/envs/cuvs/include"
+export CGO_LDFLAGS="-L/usr/local/cuda/lib64 -L/home/ajit/miniforge3/envs/cuvs/lib -lcudart -lcuvs -lcuvs_c"
+
+rapids-logger "Downloading artifacts from previous jobs"
+CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+
+# installing libcuvs/libraft will speed up the rust build substantially
+rapids-mamba-retry install \
+  --channel "${CPP_CHANNEL}" \
+  libcuvs  \
+  libraft  \
+  cuvs
+
+bash ./build.sh go

examples/c/CMakeLists.txt

@@ -13,6 +13,8 @@
 
 cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)
 
+set(CMAKE_CUDA_ARCHITECTURES 89)
+
 # ------------- configure rapids-cmake --------------#
 
 include(../cmake/thirdparty/fetch_rapids.cmake)

go/brute_force/brute_force.go

@@ -0,0 +1,64 @@
+package brute_force
+
+// #include <cuvs/neighbors/brute_force.h>
+import "C"
+
+import (
+	"errors"
+	"unsafe"
+
+	cuvs "github.com/rapidsai/cuvs/go"
+)
+
+type bruteForceIndex struct {
+	index   C.cuvsBruteForceIndex_t
+	trained bool
+}
+
+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
+}
+
+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 {
+	CMetric, exists := cuvs.CDistances[metric]
+
+	if !exists {
+		return errors.New("cuvs: invalid distance metric")
+	}
+
+	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsBruteForceBuild(C.cuvsResources_t(Resources.Resource), (*C.DLManagedTensor)(unsafe.Pointer(Dataset.C_tensor)), C.cuvsDistanceType(CMetric), C.float(metric_arg), index.index)))
+	if err != nil {
+		return err
+	}
+	index.trained = true
+
+	return nil
+}
+
+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")
+	}
+
+	prefilter := C.cuvsFilter{
+		addr:  0,
+		_type: C.NO_FILTER,
+	}
+
+	return cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsBruteForceSearch(C.ulong(resources.Resource), index.index, (*C.DLManagedTensor)(unsafe.Pointer(queries.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(neighbors.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(distances.C_tensor)), prefilter)))
+}

go/brute_force/brute_force_test.go

@@ -0,0 +1,110 @@
+package brute_force
+
+import (
+	"math/rand/v2"
+	"testing"
+
+	cuvs "github.com/rapidsai/cuvs/go"
+)
+
+func TestCagra(t *testing.T) {
+	const (
+		nDataPoints = 1024
+		nFeatures   = 16
+		nQueries    = 4
+		k           = 4
+		epsilon     = 0.001
+	)
+
+	resource, _ := cuvs.NewResource(nil)
+	defer resource.Close()
+
+	testDataset := make([][]float32, nDataPoints)
+	for i := range testDataset {
+		testDataset[i] = make([]float32, nFeatures)
+		for j := range testDataset[i] {
+			testDataset[i][j] = rand.Float32()
+		}
+	}
+
+	dataset, err := cuvs.NewTensor(testDataset)
+	if err != nil {
+		t.Fatalf("error creating dataset tensor: %v", err)
+	}
+	defer dataset.Close()
+
+	index, _ := CreateIndex()
+	defer index.Close()
+
+	// use the first 4 points from the dataset as queries : will test that we get them back
+	// as their own nearest neighbor
+	queries, _ := cuvs.NewTensor(testDataset[:nQueries])
+	defer queries.Close()
+
+	neighbors, err := cuvs.NewTensorOnDevice[int64](&resource, []int64{int64(nQueries), int64(k)})
+	if err != nil {
+		t.Fatalf("error creating neighbors tensor: %v", err)
+	}
+	defer neighbors.Close()
+
+	distances, err := cuvs.NewTensorOnDevice[float32](&resource, []int64{int64(nQueries), int64(k)})
+	if err != nil {
+		t.Fatalf("error creating distances tensor: %v", err)
+	}
+	defer distances.Close()
+
+	if _, err := dataset.ToDevice(&resource); err != nil {
+		t.Fatalf("error moving dataset to device: %v", err)
+	}
+
+	if err := BuildIndex(resource, &dataset, cuvs.DistanceL2, 2.0, index); err != nil {
+		t.Fatalf("error building index: %v", err)
+	}
+
+	if err := resource.Sync(); err != nil {
+		t.Fatalf("error syncing resource: %v", err)
+	}
+
+	if _, err := queries.ToDevice(&resource); err != nil {
+		t.Fatalf("error moving queries to device: %v", err)
+	}
+
+	err = SearchIndex(resource, *index, &queries, &neighbors, &distances)
+	if err != nil {
+		t.Fatalf("error searching index: %v", err)
+	}
+
+	if _, err := neighbors.ToHost(&resource); err != nil {
+		t.Fatalf("error moving neighbors to host: %v", err)
+	}
+
+	if _, err := distances.ToHost(&resource); err != nil {
+		t.Fatalf("error moving distances to host: %v", err)
+	}
+
+	if err := resource.Sync(); err != nil {
+		t.Fatalf("error syncing resource: %v", err)
+	}
+
+	neighborsSlice, err := neighbors.Slice()
+	if err != nil {
+		t.Fatalf("error getting neighbors slice: %v", err)
+	}
+
+	for i := range neighborsSlice {
+		if neighborsSlice[i][0] != int64(i) {
+			t.Error("wrong neighbor, expected", i, "got", neighborsSlice[i][0])
+		}
+	}
+
+	distancesSlice, err := distances.Slice()
+	if err != nil {
+		t.Fatalf("error getting distances slice: %v", err)
+	}
+
+	for i := range distancesSlice {
+		if distancesSlice[i][0] >= epsilon || distancesSlice[i][0] <= -epsilon {
+			t.Error("distance should be close to 0, got", distancesSlice[i][0])
+		}
+	}
+}

go/cagra/cagra.go

@@ -0,0 +1,62 @@
+package cagra
+
+// #include <cuvs/neighbors/cagra.h>
+import "C"
+
+import (
+	"errors"
+	"unsafe"
+
+	cuvs "github.com/rapidsai/cuvs/go"
+)
+
+type CagraIndex struct {
+	index   C.cuvsCagraIndex_t
+	trained bool
+}
+
+func CreateIndex() (*CagraIndex, error) {
+	var index C.cuvsCagraIndex_t
+	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexCreate(&index)))
+	if err != nil {
+		return nil, err
+	}
+
+	return &CagraIndex{index: index}, nil
+}
+
+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 {
+		return err
+	}
+	index.trained = true
+	return nil
+}
+
+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")
+	}
+	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraExtend(C.ulong(Resources.Resource), params.params, (*C.DLManagedTensor)(unsafe.Pointer(additional_dataset.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(return_dataset.C_tensor)), index.index)))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (index *CagraIndex) Close() error {
+	err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraIndexDestroy(index.index)))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+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")
+	}
+
+	return cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraSearch(C.cuvsResources_t(Resources.Resource), params.params, index.index, (*C.DLManagedTensor)(unsafe.Pointer(queries.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(neighbors.C_tensor)), (*C.DLManagedTensor)(unsafe.Pointer(distances.C_tensor)))))
+}