first weighted graclus gpu try

aten/THC/THCNumerics.cuh

@@ -16,6 +16,7 @@
 template<typename T>
 struct THCNumerics {
   static inline __host__ __device__ T div(T a, T b) { return a / b; }
+  static inline __host__ __device__ bool gt(T a, T b) { return a > b; }
   static inline __host__ __device__ int floor(T a) { return a; }
 };
 
@@ -23,6 +24,7 @@ struct THCNumerics {
 template<>
 struct THCNumerics<half> {
   static inline __host__ __device__ half div(half a, half b) { return f2h(h2f(a) / h2f(b)); }
+  static inline __host__ __device__ bool gt(half a, half b) { return h2f(a) > h2f(b); }
   static inline __host__ __device__ int floor(half a) { return (int) h2f(a); }
 };
 #endif  // CUDA_HALF_TENSOR

aten/THC/THCPropose.cuh

@@ -2,6 +2,7 @@
 #define THC_PROPOSE_INC
 
 #include "common.cuh"
+#include "THCNumerics.cuh"
 
 __global__ void proposeKernel(int64_t *color, int64_t *prop, int64_t *row, int64_t *col,
                               int64_t *degree, int64_t *cumDegree, ptrdiff_t nNodes) {
@@ -22,6 +23,24 @@ __global__ void weightedProposeKernel(int64_t *color, int64_t *prop, int64_t *ro
                                       T *weight, int64_t *degree, int64_t *cumDegree,
                                       ptrdiff_t nNodes) {
   KERNEL_LOOP(i, nNodes) {
+    if (color[i] != -1) { continue; }  // Only visit blue nodes.
+    ptrdiff_t c; bool isDead = true;
+    T maxWeight, tmp;
+    int64_t matchedValue;
+    for (ptrdiff_t e = cumDegree[i] - degree[i]; e < cumDegree[i]; e++) {
+      maxWeight = ScalarConvert<int, T>::to(0);
+      matchedValue = -1;
+      c = col[e];
+      tmp = weight[e];
+      if (isDead && color[c] < 0) { isDead = false; }  // Unmatched neighbor found.
+      // Propose to current maximum weighted red neighbor.
+      if (color[c] == -2 && THCNumerics<T>::gt(tmp, maxWeight)) {
+        matchedValue = c;
+        maxWeight = tmp;
+      }
+    }
+    if (matchedValue >= 0) { prop[i] = matchedValue; }
+    if (isDead) { color[i] = i; }  // Mark node as dead.
   }
 }
 

aten/THC/THCResponse.cuh

@@ -7,8 +7,7 @@ __global__ void responseKernel(int64_t *color, int64_t *prop, int64_t *row, int6
                                int64_t *degree, int64_t *cumDegree, ptrdiff_t nNodes) {
   KERNEL_LOOP(i, nNodes) {
     if (color[i] != -2) { continue; }  // Only visit red nodes.
-    ptrdiff_t c;  // int64_t neighborColor, minValue;
-    bool isDead = true;
+    ptrdiff_t c; bool isDead = true;
     for (ptrdiff_t e = cumDegree[i] - degree[i]; e < cumDegree[i]; e++) {
       c = col[e];
       if (isDead && color[c] < 0) { isDead = false; }  // Unmatched neighbor found.
@@ -27,6 +26,27 @@ __global__ void weightedResponseKernel(int64_t *color, int64_t *prop, int64_t *r
                                        T *weight, int64_t *degree, int64_t *cumDegree,
                                        ptrdiff_t nNodes) {
   KERNEL_LOOP(i, nNodes) {
+    if (color[i] != -2) { continue; }  // Only visit red nodes.
+    ptrdiff_t c; bool isDead = true;
+    T maxWeight, tmp;
+    ptrdiff_t matchedValue;
+    for (ptrdiff_t e = cumDegree[i] - degree[i]; e < cumDegree[i]; e++) {
+      maxWeight = ScalarConvert<int, T>::to(0);
+      matchedValue = -1;
+      c = col[e];
+      tmp = weight[e];
+      if (isDead && color[c] < 0) { isDead = false; }  // Unmatched neighbor found.
+      // Match maximum weighted blue neighbor, who proposed to i.
+      if (color[c] == -1 && prop[c] == i && THCNumerics<T>::gt(tmp, maxWeight)) {
+        matchedValue = c;
+        maxWeight = tmp;
+      }
+    }
+    if (matchedValue >= 0) {
+      color[i] = min(i, matchedValue);
+      color[c] = min(i, matchedValue);
+    }
+    if (isDead) { color[i] = i; }  // Mark node as dead.
   }
 }
 

aten/THC/generic/THCGraclus.cu

@@ -7,11 +7,10 @@ void THCTensor_(graclus)(THCState *state, THCudaLongTensor *self, THCudaLongTens
   THCAssertSameGPU(THCTensor_(checkGPU)(state, 4, self, row, col, weight));
 
   THC_TENSOR_GRACLUS(state, self, row,
-    /* while(!THCudaLongTensor_color(state, self)) { */
-      THCudaLongTensor_color(state, self);
+    while(!THCudaLongTensor_color(state, self)) {
       THCTensor_(propose)(state, self, prop, row, col, weight, degree, cumDegree);
       THCTensor_(response)(state, self, prop, row, col, weight, degree, cumDegree);
-    /* } */
+    }
   )
 }
 

aten/THCC/generic/THCCGraclus.c

@@ -4,7 +4,7 @@
 
 void THCCTensor_(graclus)(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,
                          THCTensor *weight) {
-  /* THCTensor_(graclus)(state, self, row, col, weight); */
+  THCTensor_(graclus)(state, self, row, col, weight);
 }
 
 #endif