propose kernel, deg and cumDeg computation

aten/THC/THCColor.cu

@@ -14,8 +14,6 @@ __global__ void assignColorKernel(int64_t *color, curandStateMtgp32 *state, uint
 }
 
 int THCGreedy_assignColor(THCState *state, THCudaLongTensor *color) {
-  THCAssertSameGPU(THCudaLongTensor_checkGPU(state, 1, color));
-
   int64_t *colorData = THCudaLongTensor_data(state, color);
   ptrdiff_t nNodes = THCudaLongTensor_nElement(state, color);
   uint8_t* d_done; cudaMalloc(&d_done, sizeof(uint8_t)); cudaMemset(d_done, 1, sizeof(uint8_t));

aten/THC/THCDegree.cu

+void THCDegree(THCState *state, THCudaLongTensor *self, THCudaLongTensor *index) {
+  int nEdges = THCudaLongTensor_nElement(state, index);
+  THCudaLongTensor *one = THCudaLongTensor_newWithSize1d(state, nEdges);
+  THCudaLongTensor_fill(state, one, 1);
+
+  THCudaLongTensor_fill(state, self, 0);
+  THCudaLongTensor_scatterAdd(state, self, 0, index, one);
+
+  THCudaLongTensor_free(state, one);
+}

aten/THC/THCGreedy.cu

 #include "THCGreedy.h"
 
+#include "common.cuh"
+#include "THCDegree.cu"
 #include "THCColor.cu"
 
+__global__ void proposeKernel(int64_t *tensor, int64_t *color, int64_t *row, int64_t *col,
+                               int64_t *deg, int64_t *cumDeg, ptrdiff_t nNodes) {
+  KERNEL_LOOP(i, nNodes) {
+    if (color[i] != -1) continue; // Only visit blue nodes.
+    ptrdiff_t c;
+    for (ptrdiff_t e = cumDeg[i]; e < cumDeg[i] + deg[i]; e++) {
+      c = col[e];
+      if (color[c] == -2) { tensor[i] = c; break; }  // Propose to first red node.
+    }
+    if (tensor[i] < 0) color[i] = i;  // Mark node as dead.
+  }
+}
+
+void THCGreedy_propose(THCState *state, THCudaLongTensor *tensor, THCudaLongTensor *color,
+                        THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,
+                        THCudaLongTensor *cumDeg) {
+  ptrdiff_t nNodes = THCudaLongTensor_nElement(state, color);
+  int64_t *tensorData = THCudaLongTensor_data(state, tensor);
+  int64_t *colorData = THCudaLongTensor_data(state, color);
+  int64_t *rowData = THCudaLongTensor_data(state, row);
+  int64_t *colData = THCudaLongTensor_data(state, col);
+  int64_t *degData = THCudaLongTensor_data(state, deg);
+  int64_t *cumDegData = THCudaLongTensor_data(state, cumDeg);
+  KERNEL_RUN(proposeKernel, nNodes, tensorData, colorData, rowData, colData, degData, cumDegData);
+}
+
 void THCGreedy(THCState *state, THCudaLongTensor *cluster, THCudaLongTensor *row,
-               THCudaLongTensor *col, THCudaLongTensor *deg) {
-  THCAssertSameGPU(THCudaLongTensor_checkGPU(state, 4, cluster, row, col, deg));
+               THCudaLongTensor *col) {
+  THCAssertSameGPU(THCudaLongTensor_checkGPU(state, 4, cluster, row, col));
+
+  int nNodes = THCudaLongTensor_nElement(state, cluster);
 
   THCudaLongTensor_fill(state, cluster, -1);
+  THCudaLongTensor *prop = THCudaLongTensor_newClone(state, cluster);
+
+  THCudaLongTensor *deg = THCudaLongTensor_newWithSize1d(state, nNodes);
+  THCDegree(state, deg, row);
+
+  THCudaLongTensor *cumDeg = THCudaLongTensor_newWithSize1d(state, nNodes);
+  THCudaLongTensor_cumsum(state, cumDeg, deg, 0);
+
   THCGreedy_assignColor(state, cluster);
+  THCGreedy_propose(state, prop, cluster, row, col, deg, cumDeg);
+
   /* while(!THCGreedy_assignColor(state, cluster)) { */
   /*   printf("DRIN"); */
-  /* }; */
-
-  // Fill cluster with -1
-  // assign color to clusters < 0 (return done)
-  // Generate proposal vector with length of nodes (init?)
   // call propose step
   // call response step
+  /* }; */
+  THCudaLongTensor_free(state, prop);
+  THCudaLongTensor_free(state, deg);
+  THCudaLongTensor_free(state, cumDeg);
 }
 
 #include "generic/THCGreedy.cu"

aten/THC/THCGreedy.h

@@ -10,7 +10,7 @@ extern "C" {
 #endif  // __cplusplus
 
 void THCGreedy(THCState *state, THCudaLongTensor *cluster, THCudaLongTensor *row,
-               THCudaLongTensor *col, THCudaLongTensor *deg);
+               THCudaLongTensor *col);
 
 #include "generic/THCGreedy.h"
 #include "THC/THCGenerateAllTypes.h"

aten/THC/generic/THCGreedy.cu

@@ -3,7 +3,7 @@
 #else
 
 void THCGreedy_(THCState *state, THCudaLongTensor *cluster, THCudaLongTensor *row,
-                THCudaLongTensor *col, THCudaLongTensor *deg, THCTensor *weight) {
+                THCudaLongTensor *col, THCTensor *weight) {
   printf("THCGreedy dynamic drin");
 }
 

aten/THC/generic/THCGreedy.h

@@ -3,6 +3,6 @@
 #else
 
 void THCGreedy_(THCState *state, THCudaLongTensor *cluster, THCudaLongTensor *row,
-                THCudaLongTensor *col, THCudaLongTensor *deg, THCTensor *weight);
+                THCudaLongTensor *col, THCTensor *weight);
 
 #endif  // THC_GENERIC_FILE

aten/THCC/THCCGreedy.c

@@ -7,9 +7,8 @@
 
 extern THCState *state;
 
-void THCCGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,
-                THCudaLongTensor *deg) {
-  THCGreedy(state, cluster, row, col, deg);
+void THCCGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col) {
+  THCGreedy(state, cluster, row, col);
 }
 
 #include "generic/THCCGreedy.c"

aten/THCC/THCCGreedy.h

-void THCCGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg);
+void THCCGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col);
 
-void   THCCByteGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,   THCudaByteTensor *weight);
-void   THCCCharGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,   THCudaCharTensor *weight);
-void  THCCShortGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,  THCudaShortTensor *weight);
-void    THCCIntGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,    THCudaIntTensor *weight);
-void   THCCLongGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,   THCudaLongTensor *weight);
-void  THCCFloatGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg,       THCudaTensor *weight);
-void THCCDoubleGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *deg, THCudaDoubleTensor *weight);
+void   THCCByteGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaByteTensor *weight);
+void   THCCCharGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaCharTensor *weight);
+void  THCCShortGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,  THCudaShortTensor *weight);
+void    THCCIntGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,    THCudaIntTensor *weight);
+void   THCCLongGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaLongTensor *weight);
+void  THCCFloatGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,       THCudaTensor *weight);
+void THCCDoubleGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaDoubleTensor *weight);

aten/THCC/generic/THCCGreedy.c

@@ -3,8 +3,8 @@
 #else
 
 void THCCGreedy_(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,
-                 THCudaLongTensor *deg, THCTensor *weight) {
-  THCGreedy_(state, cluster, row, col, deg, weight);
+                 THCTensor *weight) {
+  THCGreedy_(state, cluster, row, col, weight);
 }
 
 #endif