wip

rapidsai · ajit283 · Jun 14, 2024 · Jun 14, 2024 · Jun 19, 2024 · Jun 19, 2024
commit 7961798a868b1775e7763a3dbe4e640322e4596d
diff --git a/go/cuvs/bruteforce.go b/go/cuvs/bruteforce.go
@@ -0,0 +1,67 @@
+package common
+
+// #cgo CFLAGS: -I/usr/local/cuda/include -I/home/ajit/miniforge3/include
+// #cgo LDFLAGS: -L/usr/local/cuda/lib64 -L/home/ajit/miniforge3/lib -lcudart  -lcuvs -lcuvs_c
+// #include <cuda_runtime_api.h>
+// #include <cuvs/core/c_api.h>
+// #include <cuvs/distance/pairwise_distance.h>
+// #include <cuvs/neighbors/brute_force.h>
+// #include <cuvs/neighbors/ivf_flat.h>
+// #include <cuvs/neighbors/cagra.h>
+// #include <cuvs/neighbors/ivf_pq.h>
+import "C"
+import (
+	"unsafe"
+)
+
+type Index struct {
+	index   C.cuvsBruteForceIndex_t
+	trained bool
+}
+
+func CreateIndex() *Index {
+
+	index := (C.cuvsBruteForceIndex_t)(C.malloc(C.size_t(unsafe.Sizeof(C.cuvsBruteForceIndex{}))))
+
+	err := C.cuvsBruteForceIndexCreate(&index)
+
+	CheckCuvs(err)
+
+	return &Index{index: index, trained: false}
+
+}
+
+func DestroyIndex(index Index) {
+	err := C.cuvsBruteForceIndexDestroy(index.index)
+	CheckCuvs(err)
+
+}
+
+func BuildIndex(Resources C.cuvsResources_t, Dataset *C.DLManagedTensor, metric string, metric_arg float32, index C.cuvsBruteForceIndex_t) {
+
+	// Data := unsafe.Pointer(Dataset)
+
+	// C.cuvsRMMAlloc(Resources, &Data, 24)
+
+	CMetric := C.cuvsDistanceType(0)
+
+	switch metric {
+	case "L2Expanded":
+		CMetric = C.L2Expanded
+	default:
+		panic("Unsupported metric")
+	}
+
+	CheckCuvs(C.cuvsBruteForceBuild(Resources, Dataset, CMetric, C.float(metric_arg), index))
+
+}
+
+func SearchIndex(resources C.cuvsResources_t, index Index, queries *ManagedTensor, neighbors *ManagedTensor, distances *ManagedTensor) {
+
+	if !index.trained {
+		panic("Index needs to be built before calling search.")
+	}
+
+	CheckCuvs(C.cuvsBruteForceSearch(resources, index.index, queries, neighbors, distances))
+
+}
diff --git a/go/cuvs/bruteforce_test.go b/go/cuvs/bruteforce_test.go
@@ -0,0 +1,33 @@
+package common
+
+import (
+	"testing"
+)
+
+func TestBruteForce(t *testing.T) {
+
+	resource := NewResource(nil)
+
+	dataset := NewManagedTensor(true, []int{1, 1}, []float32{1, 2})
+
+	println(dataset.dl_tensor.dtype.code)
+
+	index := CreateIndex()
+
+	BuildIndex(resource.resource, dataset, "L2Expanded", 2.0, index.index)
+
+	index.trained = true
+
+	queries := NewManagedTensor(true, []int{1, 1}, []float32{1, 2})
+
+	neighbors := NewManagedTensor(true, []int{1, 1}, []float32{3, 3})
+
+	distances := NewManagedTensor(true, []int{1, 1}, []float32{1, 4})
+
+	SearchIndex(resource.resource, *index, queries, neighbors, distances)
+
+	p := (*float32)(neighbors.dl_tensor.data)
+
+	println(*p)
+
+}
diff --git a/go/cuvs/dlpack.go b/go/cuvs/dlpack.go
@@ -0,0 +1,64 @@
+package common
+
+// #include <stdio.h>
+// #include <stdlib.h>
+// #include <dlpack/dlpack.h>  // Replace with the actual header file containing DLManagedTensor
+// void delete_tensor(DLManagedTensor *tensor){
+//     free(tensor->dl_tensor.shape);
+//     tensor->manager_ctx = NULL;
+//     free(tensor);
+// }
+import "C"
+import (
+	"unsafe"
+)
+
+type ManagedTensor = C.DLManagedTensor
+
+func NewManagedTensor(from_cai bool, shape []int, data []float32) *C.DLManagedTensor {
+
+	if len(shape) < 2 {
+		panic("shape must be atleast 2")
+	}
+
+	dlm := (*C.DLManagedTensor)(C.malloc(C.size_t(unsafe.Sizeof(C.DLManagedTensor{}))))
+
+	if dlm == nil {
+		panic("memory allocation failed")
+	}
+
+	var devicetype C.DLDeviceType
+
+	if from_cai {
+		devicetype = C.kDLCUDA
+	} else {
+		devicetype = C.kDLCPU
+	}
+
+	device := C.DLDevice{
+		device_type: devicetype,
+		device_id:   0,
+	}
+
+	dtype := C.DLDataType{
+		bits:  C.uchar(32),
+		lanes: C.ushort(1),
+		code:  C.kDLFloat,
+	}
+
+	dlm.dl_tensor.data = unsafe.Pointer(&data[0])
+
+	dlm.dl_tensor.device = device
+
+	dlm.dl_tensor.dtype = dtype
+	dlm.dl_tensor.ndim = C.int(len(shape))
+	dlm.dl_tensor.shape = (*C.int64_t)(unsafe.Pointer(&shape[0]))
+	dlm.dl_tensor.strides = nil
+	dlm.dl_tensor.byte_offset = 0
+
+	dlm.manager_ctx = nil
+	dlm.deleter = (*[0]byte)(C.delete_tensor)
+
+	return dlm
+
+}
diff --git a/go/cuvs/dlpack_test.go b/go/cuvs/dlpack_test.go
@@ -0,0 +1,11 @@
+package common
+
+import (
+	"testing"
+)
+
+func TestCreateNewTensor(t *testing.T) {
+
+	// tensor := NewManagedTensor(false, []int{1}, []int32{1})
+
+}
diff --git a/go/cuvs/exceptions.go b/go/cuvs/exceptions.go
@@ -0,0 +1,17 @@
+package common
+
+// #cgo CFLAGS: -I/usr/local/cuda/include -I/home/ajit/miniforge3/include
+// #cgo LDFLAGS: -L/usr/local/cuda/lib64 -L/home/ajit/miniforge3/lib -lcudart -lcudart -lcuvs
+// #include <cuvs/core/c_api.h>
+// #include <cuvs/distance/pairwise_distance.h>
+// #include <cuvs/neighbors/brute_force.h>
+// #include <cuvs/neighbors/ivf_flat.h>
+// #include <cuvs/neighbors/cagra.h>
+// #include <cuvs/neighbors/ivf_pq.h>
+import "C"
+
+func CheckCuvs(error C.cuvsError_t) {
+	if error == C.CUVS_ERROR {
+		panic(C.GoString(C.cuvsGetLastErrorText()))
+	}
+}
diff --git a/go/cuvs/resources.go b/go/cuvs/resources.go
@@ -0,0 +1,34 @@
+package common
+
+// #cgo CFLAGS: -I/usr/local/cuda/include -I/home/ajit/miniforge3/include
+// #cgo LDFLAGS: -L/usr/local/cuda/lib64 -L/home/ajit/miniforge3/lib -lcudart -lcudart -lcuvs
+// #include <cuda_runtime_api.h>
+// #include <cuvs/core/c_api.h>
+// #include <cuvs/distance/pairwise_distance.h>
+// #include <cuvs/neighbors/brute_force.h>
+// #include <cuvs/neighbors/ivf_flat.h>
+// #include <cuvs/neighbors/cagra.h>
+// #include <cuvs/neighbors/ivf_pq.h>
+import "C"
+
+type Resource struct {
+	resource C.cuvsResources_t
+}
+
+// func NewResource() *Resource {
+func NewResource(stream C.cudaStream_t) Resource {
+
+	res := C.cuvsResources_t(0)
+
+	CheckCuvs(C.cuvsResourcesCreate(&res))
+
+	if stream != nil {
+		CheckCuvs(C.cuvsStreamSet(res, stream))
+	}
+
+	return Resource{resource: res}
+}
+
+func Sync(r C.cuvsResources_t) {
+	CheckCuvs(C.cuvsStreamSync(r))
+}
diff --git a/go/go.mod b/go/go.mod
@@ -0,0 +1,3 @@
+module rapidsai/cuvs
+
+go 1.22.4