minor cleanup, ready for release
[RBC.git] / rbc.cu
diff --git a/rbc.cu b/rbc.cu
index 51607f9..c0494e6 100644 (file)
--- a/rbc.cu
+++ b/rbc.cu
@@ -16,7 +16,7 @@
 #include "kernelWrap.h"
 #include "sKernelWrap.h"
 
-void queryRBC(const matrix q, const rbcStruct rbcS, unint *NNs){
+void queryRBC(const matrix q, const rbcStruct rbcS, unint *NNs, real* NNdists){
   unint m = q.r;
   unint numReps = rbcS.dr.r;
   unint compLength;
@@ -55,8 +55,64 @@ void queryRBC(const matrix q, const rbcStruct rbcS, unint *NNs){
   checkErr( cudaMalloc( (void**)&dqMap, compLength*sizeof(*dqMap) ) );
   cudaMemcpy( dqMap, qMap, compLength*sizeof(*dqMap), cudaMemcpyHostToDevice );
   
-  computeNNs(rbcS.dx, rbcS.dxMap, dq, dqMap, dcP, NNs, compLength);
+  computeNNs(rbcS.dx, rbcS.dxMap, dq, dqMap, dcP, NNs, NNdists, compLength);
   
+  free(qMap);
+  cudaFree(dqMap);
+  freeCompPlan(&dcP);
+  cudaFree(dq.mat);
+  free(cM.mat);
+  free(repIDsQ);
+  free(distToRepsQ);
+  free(groupCountQ);
+}
+
+//This function is very similar to queryRBC, with a couple of basic changes to handle
+//k-nn.  
+void kqueryRBC(const matrix q, const rbcStruct rbcS, intMatrix NNs, matrix NNdists){
+  unint m = q.r;
+  unint numReps = rbcS.dr.r;
+  unint compLength;
+  compPlan dcP;
+  unint *qMap, *dqMap;
+  qMap = (unint*)calloc(PAD(m+(BLOCK_SIZE-1)*PAD(numReps)),sizeof(*qMap));
+  matrix dq;
+  copyAndMove(&dq, &q);
+  
+  charMatrix cM;
+  cM.r=cM.c=numReps; cM.pr=cM.pc=cM.ld=PAD(numReps);
+  cM.mat = (char*)calloc( cM.pr*cM.pc, sizeof(*cM.mat) );
+  
+  unint *repIDsQ;
+  repIDsQ = (unint*)calloc( m, sizeof(*repIDsQ) );
+  real *distToRepsQ;
+  distToRepsQ = (real*)calloc( m, sizeof(*distToRepsQ) );
+  unint *groupCountQ;
+  groupCountQ = (unint*)calloc( PAD(numReps), sizeof(*groupCountQ) );
+  
+  computeReps(dq, rbcS.dr, repIDsQ, distToRepsQ);
+
+  //How many points are assigned to each group?
+  computeCounts(repIDsQ, m, groupCountQ);
+  
+  //Set up the mapping from groups to queries (qMap).
+  buildQMap(q, qMap, repIDsQ, numReps, &compLength);
+  
+  // Setup the computation matrix.  Currently, the computation matrix is 
+  // just the identity matrix: each query assigned to a particular 
+  // representative is compared only to that representative's points.  
+
+  // NOTE: currently, idIntersection is the *only* computation matrix 
+  // that will work properly with k-nn search (this is not true for 1-nn above).
+  idIntersection(cM);
+
+  initCompPlan(&dcP, cM, groupCountQ, rbcS.groupCount, numReps);
+
+  checkErr( cudaMalloc( (void**)&dqMap, compLength*sizeof(*dqMap) ) );
+  cudaMemcpy( dqMap, qMap, compLength*sizeof(*dqMap), cudaMemcpyHostToDevice );
+  
+  computeKNNs(rbcS.dx, rbcS.dxMap, dq, dqMap, dcP, NNs, NNdists, compLength);
+
   free(qMap);
   freeCompPlan(&dcP);
   cudaFree(dq.mat);
@@ -68,21 +124,20 @@ void queryRBC(const matrix q, const rbcStruct rbcS, unint *NNs){
 
 
 void buildRBC(const matrix x, rbcStruct *rbcS, unint numReps, unint s){
-  //const matrix dr, intMatrix xmap, unint *counts, unint s){
   unint n = x.pr;
   intMatrix xmap;
 
   setupReps(x, rbcS, numReps);
   copyAndMove(&rbcS->dx, &x);
   
-  xmap.r=numReps; xmap.pr=PAD(numReps); xmap.c=s; xmap.pr=xmap.ld=PAD(s);
+  xmap.r=numReps; xmap.pr=PAD(numReps); xmap.c=s; xmap.pc=xmap.ld=PAD(s);
   xmap.mat = (unint*)calloc( xmap.pr*xmap.pc, sizeof(*xmap.mat) );
   copyAndMoveI(&rbcS->dxMap, &xmap);
-  rbcS->groupCount = (uint*)calloc( PAD(numReps), sizeof(*rbcS->groupCount) );
+  rbcS->groupCount = (unint*)calloc( PAD(numReps), sizeof(*rbcS->groupCount) );
   
   //Figure out how much fits into memory
-  unint memFree, memTot;
-  cuMemGetInfo(&memFree, &memTot);
+  size_t memFree, memTot;
+  cudaMemGetInfo(&memFree, &memTot);
   memFree = (unint)(((float)memFree)*MEM_USABLE);
   /* mem needed per rep:
    *  n*sizeof(real) - dist mat
@@ -94,7 +149,7 @@ void buildRBC(const matrix x, rbcStruct *rbcS, unint numReps, unint s){
    */
   unint ptsAtOnce = DPAD(memFree/((n+1)*sizeof(real) + n*sizeof(char) + (n+1)*sizeof(unint) + 2*MEM_USED_IN_SCAN(n)));
   if(!ptsAtOnce){
-    fprintf(stderr,"error: %d is not enough memory to build the RBC.. exiting\n", memFree);
+    fprintf(stderr,"error: %lu is not enough memory to build the RBC.. exiting\n", (unsigned long)memFree);
     exit(1);
   }
 
@@ -247,7 +302,6 @@ void idIntersection(charMatrix cM){
 }
 
 
-
 void fullIntersection(charMatrix cM){
   unint i,j;
   for(i=0;i<cM.r;i++){
@@ -258,22 +312,40 @@ void fullIntersection(charMatrix cM){
 }
 
 
-void computeNNs(matrix dx, intMatrix dxMap, matrix dq, unint *dqMap, compPlan dcP, unint *NNs, unint compLength){
-  real *dMins;
+void computeNNs(matrix dx, intMatrix dxMap, matrix dq, unint *dqMap, compPlan dcP, unint *NNs, real *NNdists, unint compLength){
+  real *dNNdists;
   unint *dMinIDs;
   
-  checkErr( cudaMalloc((void**)&dMins,compLength*sizeof(*dMins)) );
+  checkErr( cudaMalloc((void**)&dNNdists,compLength*sizeof(*dNNdists)) );
   checkErr( cudaMalloc((void**)&dMinIDs,compLength*sizeof(*dMinIDs)) );
 
-  planNNWrap(dq, dqMap, dx, dxMap, dMins, dMinIDs, dcP, compLength);
-  cudaMemcpy( NNs, dMinIDs, dq.r*sizeof(*NNs), cudaMemcpyDeviceToHost);
-  
-      
-  cudaFree(dMins);
+  planNNWrap(dq, dqMap, dx, dxMap, dNNdists, dMinIDs, dcP, compLength );
+  cudaMemcpy( NNs, dMinIDs, dq.r*sizeof(*NNs), cudaMemcpyDeviceToHost );
+  cudaMemcpy( NNdists, dNNdists, dq.r*sizeof(*dNNdists), cudaMemcpyDeviceToHost );
+
+  cudaFree(dNNdists);
   cudaFree(dMinIDs);
 }
 
 
+void computeKNNs(matrix dx, intMatrix dxMap, matrix dq, unint *dqMap, compPlan dcP, intMatrix NNs, matrix NNdists, unint compLength){
+  matrix dNNdists;
+  intMatrix dMinIDs;
+  dNNdists.r=compLength; dNNdists.pr=compLength; dNNdists.c=K; dNNdists.pc=K; dNNdists.ld=dNNdists.pc;
+  dMinIDs.r=compLength; dMinIDs.pr=compLength; dMinIDs.c=K; dMinIDs.pc=K; dMinIDs.ld=dMinIDs.pc;
+
+  checkErr( cudaMalloc((void**)&dNNdists.mat,dNNdists.pr*dNNdists.pc*sizeof(*dNNdists.mat)) );
+  checkErr( cudaMalloc((void**)&dMinIDs.mat,dMinIDs.pr*dMinIDs.pc*sizeof(*dMinIDs.mat)) );
+
+  planKNNWrap(dq, dqMap, dx, dxMap, dNNdists, dMinIDs, dcP, compLength);
+  cudaMemcpy( NNs.mat, dMinIDs.mat, dq.r*K*sizeof(*NNs.mat), cudaMemcpyDeviceToHost );
+  cudaMemcpy( NNdists.mat, dNNdists.mat, dq.r*K*sizeof(*NNdists.mat), cudaMemcpyDeviceToHost );
+
+  cudaFree(dNNdists.mat);
+  cudaFree(dMinIDs.mat);
+}
+
+
 //This calls the dist1Kernel wrapper, but has it compute only 
 //a submatrix of the all-pairs distance matrix.  In particular,
 //only distances from dr[start,:].. dr[start+length-1] to all of x
@@ -349,6 +421,11 @@ void initCompPlan(compPlan *dcP, charMatrix cM, unint *groupCountQ, unint *group
   checkErr( cudaMalloc( (void**)&dcP->qGroupToXGroup, sQGroupToXGroup*sizeof(*dcP->qGroupToXGroup) ) );
   cudaMemcpy( dcP->qGroupToXGroup, cP.qGroupToXGroup, sQGroupToXGroup*sizeof(*dcP->qGroupToXGroup), cudaMemcpyHostToDevice );
   dcP->ld = cP.ld;
+
+  free(cP.numGroups);
+  free(cP.groupCountX);
+  free(cP.qToQGroup);
+  free(cP.qGroupToXGroup);
 }