Go API - [WIP]

.github/workflows/build.yaml

@@ -47,6 +47,19 @@ jobs:
       node_type: "gpu-v100-latest-1"
       run_script: "ci/build_rust.sh"
       sha: ${{ inputs.sha }}
+  go-build:
+    needs: cpp-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: ${{ inputs.build_type || 'branch' }}
+      branch: ${{ inputs.branch }}
+      arch: "amd64"
+      date: ${{ inputs.date }}
+      container_image: "rapidsai/ci-conda:latest"
+      node_type: "gpu-v100-latest-1"
+      run_script: "ci/build_go.sh"
+      sha: ${{ inputs.sha }}
   python-build:
     needs: [cpp-build]
     secrets: inherit

.github/workflows/pr.yaml

@@ -21,6 +21,7 @@ jobs:
       - conda-python-tests
       - docs-build
       - rust-build
+      - go-build
       - wheel-build-cuvs
       - wheel-tests-cuvs
       - devcontainer
@@ -44,13 +45,15 @@ jobs:
           - '!notebooks/**'
           - '!python/**'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
         test_notebooks:
           - '**'
           - '!.devcontainer/**'
           - '!.pre-commit-config.yaml'
           - '!README.md'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
         test_python:
           - '**'
@@ -61,6 +64,7 @@ jobs:
           - '!img/**'
           - '!notebooks/**'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
   checks:
     secrets: inherit
@@ -122,6 +126,16 @@ jobs:
       arch: "amd64"
       container_image: "rapidsai/ci-conda:latest"
       run_script: "ci/build_rust.sh"
+  go-build:
+    needs: conda-cpp-build
+    secrets: inherit
+    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+    with:
+      build_type: pull-request
+      node_type: "gpu-v100-latest-1"
+      arch: "amd64"
+      container_image: "rapidsai/ci-conda:latest"
+      run_script: "ci/build_go.sh"
   wheel-build-cuvs:
     needs: checks
     secrets: inherit

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,40 @@
+#!/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"
+export CC=clang
+
+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

dependencies.yaml

@@ -23,6 +23,7 @@ files:
       - test_py_cuvs
       - cupy
       - rust
+      - go
   bench_ann:
     output: conda
     matrix:
@@ -74,6 +75,13 @@ files:
       - rapids_build
       - cuda
       - rust
+  go:
+    output: none
+    includes:
+      - cuda_version
+      - rapids_build
+      - cuda
+      - rust
   py_build_cuvs:
     output: pyproject
     pyproject_dir: python/cuvs
@@ -411,6 +419,15 @@ dependencies:
           # clang/liblclang only needed for bindgen support
           - clang
           - libclang
+  go:
+    common:
+      - output_types: [conda]
+        packages:
+          - make
+          - go
+          # clang/liblclang only needed for bindgen support
+          - clang
+          - libclang
   build_wheels:
     common:
       - output_types: [requirements, pyproject]

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])
+		}
+	}
+}