File size: 2,232 Bytes
29858c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
#include <THC/THC.h>
#include "knn_cuda_kernel.h"
extern THCState *state;
int knn(THCudaTensor *ref_tensor, THCudaTensor *query_tensor,
THCudaLongTensor *idx_tensor) {
THCAssertSameGPU(THCudaTensor_checkGPU(state, 3, idx_tensor, ref_tensor, query_tensor));
long batch, ref_nb, query_nb, dim, k;
THArgCheck(THCudaTensor_nDimension(state, ref_tensor) == 3 , 0, "ref_tensor: 3D Tensor expected");
THArgCheck(THCudaTensor_nDimension(state, query_tensor) == 3 , 1, "query_tensor: 3D Tensor expected");
THArgCheck(THCudaLongTensor_nDimension(state, idx_tensor) == 3 , 3, "idx_tensor: 3D Tensor expected");
THArgCheck(THCudaTensor_size(state, ref_tensor, 0) == THCudaTensor_size(state, query_tensor,0), 0, "input sizes must match");
THArgCheck(THCudaTensor_size(state, ref_tensor, 1) == THCudaTensor_size(state, query_tensor,1), 0, "input sizes must match");
THArgCheck(THCudaTensor_size(state, idx_tensor, 2) == THCudaTensor_size(state, query_tensor,2), 0, "input sizes must match");
//ref_tensor = THCudaTensor_newContiguous(state, ref_tensor);
//query_tensor = THCudaTensor_newContiguous(state, query_tensor);
batch = THCudaLongTensor_size(state, ref_tensor, 0);
dim = THCudaTensor_size(state, ref_tensor, 1);
k = THCudaLongTensor_size(state, idx_tensor, 1);
ref_nb = THCudaTensor_size(state, ref_tensor, 2);
query_nb = THCudaTensor_size(state, query_tensor, 2);
float *ref_dev = THCudaTensor_data(state, ref_tensor);
float *query_dev = THCudaTensor_data(state, query_tensor);
long *idx_dev = THCudaLongTensor_data(state, idx_tensor);
// scratch buffer for distances
float *dist_dev = (float*)THCudaMalloc(state, ref_nb * query_nb * sizeof(float));
for (int b = 0; b < batch; b++) {
knn_device(ref_dev + b * dim * ref_nb, ref_nb, query_dev + b * dim * query_nb, query_nb, dim, k,
dist_dev, idx_dev + b * k * query_nb, THCState_getCurrentStream(state));
}
// free buffer
THCudaFree(state, dist_dev);
//printf("aaaaa\n");
// check for errors
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
printf("error in knn: %s\n", cudaGetErrorString(err));
THError("aborting");
}
return 1;
}
|