bruteforce works

rapidsai · ajit283 · Jun 14, 2024 · Jun 14, 2024 · Jun 19, 2024 · Jun 19, 2024
commit b47cf5bc98eb4b03955f1d357ef3a93431fc9ad4
diff --git a/go/cuvs/bruteforce.go b/go/cuvs/bruteforce.go
@@ -43,6 +43,10 @@ func DestroyIndex(index Index) {
 
 func BuildIndex(Resources C.cuvsResources_t, Dataset *C.DLManagedTensor, metric string, metric_arg float32, index *Index) {
 
+	if Dataset.dl_tensor.device.device_type != C.kDLCUDA {
+		panic("Dataset must be on GPU")
+	}
+
 	// Data := unsafe.Pointer(Dataset)
 
 	// CheckCuvs(C.cuvsRMMAlloc(Resources, &Data, 2400))

diff --git a/go/cuvs/bruteforce_test.go b/go/cuvs/bruteforce_test.go
@@ -1,56 +1,63 @@
 package common
 
 import (
+	"math/rand"
 	"testing"
+	"time"
+	"unsafe"
 )
 
 func TestBruteForce(t *testing.T) {
 
 	resource := NewResource(nil)
 
-	dataset := NewManagedTensor(true, []int{1, 2}, []float32{1, 2}, false)
-	bytes := GetBytes(dataset)
-	println("size:")
-	println(bytes)
-	println(dataset.dl_tensor.dtype.code)
+	rand.Seed(time.Now().UnixNano())
 
-	index := CreateIndex()
-
-	BuildIndex(resource.resource, dataset, "L2Expanded", 2.0, index)
-	println("got here")
-
-	println("index created")
-
-	index.trained = true
+	NDataPoints := 16
+	NFeatures := 8
 
-	queries := NewManagedTensor(true, []int{1, 2}, []float32{1, 2}, false)
+	TestDataset := make([]float32, NDataPoints*NFeatures)
+	for index := range TestDataset {
+		TestDataset[index] = rand.Float32()
+		// TestDataset[index] = float32(index)
+	}
 
-	neighbors := NewManagedTensor(true, []int{1, 1}, []int64{3, 3}, true)
+	dataset := NewManagedTensor(true, []int{NDataPoints, NFeatures}, TestDataset, false)
 
-	// distances := NewManagedTensor(true, []int{1, 1}, []float32{1, 4}, false)
-
-	println("tensors created")
+	index := CreateIndex()
+	// use the first 4 points from the dataset as queries : will test that we get them back
+	// as their own nearest neighbor
 
-	println(queries.dl_tensor.data)
+	NQueries := 4
+	K := 4
+	queries := NewManagedTensor(true, []int{NQueries, NFeatures}, TestDataset[:(NQueries*NFeatures)], false)
+	neighbors := NewManagedTensor(true, []int{NQueries, K}, make([]int64, NQueries*K), true)
+	distances := NewManagedTensor(true, []int{NQueries, K}, make([]float32, NQueries*K), false)
 
-	queries_device := ToDevice(queries, &resource)
+	ToDevice(neighbors, &resource)
+	ToDevice(distances, &resource)
+	ToDevice(dataset, &resource)
 
-	println("queries tensor transferred")
+	BuildIndex(resource.resource, dataset, "L2Expanded", 2.0, index)
+	Sync(resource.resource)
 
-	// neighbors_device := ToDevice(neighbors, &resource)
+	ToDevice(queries, &resource)
 
-	// distances_device := ToDevice(distances, &resource)
+	index.trained = true
 
-	println(queries_device.dl_tensor.data)
+	SearchIndex(resource.resource, *index, queries, neighbors, distances)
 
-	println("tensors on device")
+	ToHost(neighbors, &resource)
+	ToHost(distances, &resource)
 
-	// SearchIndex(resource.resource, *index, queries_device, neighbors_device, distances_device)
+	Sync(resource.resource)
 
-	println("search done")
+	p := (*int64)(unsafe.Pointer(uintptr(neighbors.dl_tensor.data) + uintptr(K*8*3)))
 
-	p := (*float32)(neighbors.dl_tensor.data)
+	d := (*float32)(distances.dl_tensor.data)
 
 	println(*p)
 
+	println(*d)
+
 }
diff --git a/go/cuvs/dlpack.go b/go/cuvs/dlpack.go
@@ -35,11 +35,11 @@ func NewManagedTensor[T any](from_cai bool, shape []int, data []T, use_int64 boo
 
 	var devicetype C.DLDeviceType
 
-	if from_cai {
-		devicetype = C.kDLCUDA
-	} else {
-		devicetype = C.kDLCPU
-	}
+	// if from_cai {
+	devicetype = C.kDLCUDA
+	// } else {
+	// devicetype = C.kDLCPU
+	// }
 
 	device := C.DLDevice{
 		device_type: devicetype,
@@ -94,7 +94,7 @@ func GetBytes(t *C.DLManagedTensor) int {
 	return bytes
 }
 
-func ToDevice(t *C.DLManagedTensor, res *Resource) *C.DLManagedTensor {
+func ToDevice(t *C.DLManagedTensor, res *Resource) {
 	bytes := GetBytes(t)
 
 	// device_data := &C.void{}
@@ -112,8 +112,8 @@ func ToDevice(t *C.DLManagedTensor, res *Resource) *C.DLManagedTensor {
 	// println("new host data location:")
 	// println(t.dl_tensor.data)
 
-	// CheckCuvs(C.cuvsRMMAlloc(res.resource, &DeviceDataPointer, C.size_t(bytes)))
-	CheckCuda(C.cudaMalloc(&DeviceDataPointer, C.size_t(bytes)))
+	CheckCuvs(C.cuvsRMMAlloc(res.resource, &DeviceDataPointer, C.size_t(bytes)))
+	// CheckCuda(C.cudaMalloc(&DeviceDataPointer, C.size_t(bytes)))
 
 	println("device data pointer:")
 	println(DeviceDataPointer)
@@ -134,9 +134,11 @@ func ToDevice(t *C.DLManagedTensor, res *Resource) *C.DLManagedTensor {
 
 	CheckCuda(
 		C.cudaMemcpy(
+			// DeviceDataPointer,
 			DeviceDataPointer,
-			// t.dl_tensor.data,
-			unsafe.Pointer(&hostData[0]),
+			// unsafe.Pointer(&hostData[0]),
+			t.dl_tensor.data,
+			// unsafe.Pointer(&hostData[0]),
 			// DeviceDataPointer,
 			// C.size_t(bytes),
 			C.size_t(bytes),
@@ -149,7 +151,65 @@ func ToDevice(t *C.DLManagedTensor, res *Resource) *C.DLManagedTensor {
 	println("done")
 
 	t.dl_tensor.data = DeviceDataPointer
+	println("normal transfer done")
+	// return t
+
+}
+
+func ToHost(t *C.DLManagedTensor, res *Resource) {
+	bytes := GetBytes(t)
+
+	// var DeviceDataPointer unsafe.Pointer
+
+	addr := (C.malloc(C.size_t(bytes)))
+
+	// hostData := make([]float32, 2)
+
+	CheckCuda(
+		C.cudaMemcpy(
+			// DeviceDataPointer,
+			addr,
+			// unsafe.Pointer(&hostData[0]),
+			t.dl_tensor.data,
+			// unsafe.Pointer(&hostData[0]),
+			// DeviceDataPointer,
+			// C.size_t(bytes),
+			C.size_t(bytes),
+			C.cudaMemcpyDeviceToHost,
+
+			// res.get_cuda_stream(),
+			// GetCudaStream(res.resource),
+		))
+
+	t.dl_tensor.data = addr
+}
+
+func TestCuda() {
+
+	var DeviceDataPointer unsafe.Pointer
+	CheckCuda(C.cudaMalloc(&DeviceDataPointer, C.size_t(8)))
+
+	array := make([]float32, 2)
+
+	for i := range array {
+		array[i] = float32(i)
+	}
+
+	CheckCuda(
+		C.cudaMemcpy(
+			// DeviceDataPointer,
+			DeviceDataPointer,
+			unsafe.Pointer(&array[0]),
+			// unsafe.Pointer(&hostData[0]),
+			// DeviceDataPointer,
+			// C.size_t(bytes),
+			C.size_t(8),
+			C.cudaMemcpyHostToDevice,
+
+			// res.get_cuda_stream(),
+			// GetCudaStream(res.resource),
+		))
 
-	return t
+	println("cuda test done")
 
 }