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

@@ -22,6 +22,7 @@ jobs:
       - conda-python-tests
       - docs-build
       - rust-build
+      - go-build
       - wheel-build-libcuvs
       - wheel-build-cuvs
       - wheel-tests-cuvs
@@ -58,13 +59,15 @@ jobs:
           - '!notebooks/**'
           - '!python/**'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
         test_notebooks:
           - '**'
           - '!.devcontainer/**'
           - '!.pre-commit-config.yaml'
           - '!README.md'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
         test_python:
           - '**'
@@ -75,6 +78,7 @@ jobs:
           - '!img/**'
           - '!notebooks/**'
           - '!rust/**'
+          - '!go/**'
           - '!thirdparty/LICENSES/**'
   checks:
     secrets: inherit
@@ -136,6 +140,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-libcuvs:
     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 java 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 java 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
    java             - build the cuvs Java bindings
    docs             - build the documentation
    tests            - build the tests
@@ -447,6 +448,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
+
 # Build the cuvs Java bindings
 if (( ${NUMARGS} == 0 )) || hasArg java; then
     if ! hasArg libcuvs; then

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${CONDA_PREFIX}/include"
+export CGO_LDFLAGS="-L${CONDA_PREFIX}/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

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -23,6 +23,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -23,6 +23,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -24,6 +24,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-aarch64=13.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -24,6 +24,7 @@ dependencies:
 - dlpack>=0.8,<1.0
 - doxygen>=1.8.20
 - gcc_linux-64=13.*
+- go
 - graphviz
 - ipython
 - libclang==16.0.6

dependencies.yaml

@@ -21,6 +21,7 @@ files:
       - rapids_build
       - run_py_cuvs
       - rust
+      - go
       - test_libcuvs
       - test_python_common
       - test_py_cuvs
@@ -81,6 +82,14 @@ files:
       - cuda_version
       - rapids_build
       - rust
+  go:
+    output: none
+    includes:
+      - clang
+      - cuda
+      - cuda_version
+      - rapids_build
+      - go
   py_build_libcuvs:
     output: pyproject
     pyproject_dir: python/libcuvs
@@ -470,6 +479,11 @@ dependencies:
         packages:
           - make
           - rust
+  go:
+    common:
+      - output_types: [conda]
+        packages:
+          - go
   build_wheels:
     common:
       - output_types: [requirements, pyproject]

go/brute_force/brute_force.go

@@ -0,0 +1,83 @@
+package brute_force
+
+// #include <cuvs/neighbors/brute_force.h>
+import "C"
+
+import (
+	"errors"
+	"unsafe"
+
+	cuvs "github.com/rapidsai/cuvs/go"
+)
+
+// Brute Force KNN Index
+type BruteForceIndex struct {
+	index   C.cuvsBruteForceIndex_t
+	trained bool
+}
+
+// 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
+}
+
+// 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
+}
+
+// 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 {
+		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
+}
+
+// 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")
+	}
+
+	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)))
+}