weighted proposal boilerplate

aten/THC/THCGraclus.cu

@@ -6,30 +6,36 @@
 #include "THCPropose.cuh"
 #include "THCResponse.cuh"
 
+#define THC_TENSOR_GRACLUS(state, self, row, CODE) { \
+  int nNodes = THCudaLongTensor_nElement(state, self); \
+  THCudaLongTensor_fill(state, self, -1); \
+\
+  THCudaLongTensor *prop = THCudaLongTensor_newWithSize1d(state, nNodes); \
+  THCudaLongTensor_fill(state, prop, -1); \
+\
+  THCudaLongTensor *degree = THCudaLongTensor_newWithSize1d(state, nNodes); \
+  THCudaLongTensor_degree(state, degree, row); \
+\
+  THCudaLongTensor *cumDegree = THCudaLongTensor_newWithSize1d(state, nNodes); \
+  THCudaLongTensor_cumDegree(state, cumDegree, row); \
+\
+  CODE \
+\
+  THCudaLongTensor_free(state, prop); \
+  THCudaLongTensor_free(state, degree); \
+  THCudaLongTensor_free(state, cumDegree); \
+}
+
 void THCTensor_graclus(THCState *state, THCudaLongTensor *self, THCudaLongTensor *row,
                        THCudaLongTensor *col) {
   THCAssertSameGPU(THCudaLongTensor_checkGPU(state, 3, self, row, col));
 
-  int nNodes = THCudaLongTensor_nElement(state, self);
-  THCudaLongTensor_fill(state, self, -1);
-
-  THCudaLongTensor *prop = THCudaLongTensor_newWithSize1d(state, nNodes);
-  THCudaLongTensor_fill(state, prop, -1);
-
-  THCudaLongTensor *degree = THCudaLongTensor_newWithSize1d(state, nNodes);
-  THCudaLongTensor_degree(state, degree, row);
-
-  THCudaLongTensor *cumDegree = THCudaLongTensor_newWithSize1d(state, nNodes);
-  THCudaLongTensor_cumDegree(state, cumDegree, row);
-
-  while(!THCudaLongTensor_color(state, self)) {
-    THCTensor_propose(state, self, prop, row, col, degree, cumDegree);
-    THCTensor_response(state, self, prop, row, col, degree, cumDegree);
-  }
-
-  THCudaLongTensor_free(state, prop);
-  THCudaLongTensor_free(state, degree);
-  THCudaLongTensor_free(state, cumDegree);
+  THC_TENSOR_GRACLUS(state, self, row,
+    while(!THCudaLongTensor_color(state, self)) {
+      THCTensor_propose(state, self, prop, row, col, degree, cumDegree);
+      THCTensor_response(state, self, prop, row, col, degree, cumDegree);
+    }
+  )
 }
 
 #include "generic/THCGraclus.cu"

aten/THC/THCPropose.cuh

@@ -17,6 +17,15 @@ __global__ void proposeKernel(int64_t *color, int64_t *prop, int64_t *row, int64
   }
 }
 
+template<typename T>
+__global__ void weightedProposeKernel(int64_t *color, int64_t *prop, int64_t *row, int64_t *col,
+                                      T *weight, int64_t *degree, int64_t *cumDegree,
+                                      ptrdiff_t nNodes) {
+  KERNEL_LOOP(i, nNodes) {
+  }
+}
+
+
 void THCTensor_propose(THCState *state, THCudaLongTensor *color, THCudaLongTensor *prop,
                        THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *degree,
                        THCudaLongTensor *cumDegree) {
@@ -26,4 +35,7 @@ void THCTensor_propose(THCState *state, THCudaLongTensor *color, THCudaLongTenso
              THCudaLongTensor_data(state, degree), THCudaLongTensor_data(state, cumDegree));
 }
 
+#include "generic/THCPropose.cuh"
+#include "THC/THCGenerateAllTypes.h"
+
 #endif  // THC_PROPOSE_INC

aten/THC/THCResponse.cuh

@@ -3,15 +3,6 @@
 
 #include "common.cuh"
 
-    /* if (color[i] != -1) { continue; }  // Only visit blue nodes. */
-    /* 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. */
-    /*   if (color[c] == -2) { prop[i] = c; break; }  // Propose to first red neighbor. */
-    /* } */
-    /* if (isDead) { color[i] = i; }  // Mark node as dead. */
-
 __global__ void responseKernel(int64_t *color, int64_t *prop, int64_t *row, int64_t *col,
                                int64_t *degree, int64_t *cumDegree, ptrdiff_t nNodes) {
   KERNEL_LOOP(i, nNodes) {
@@ -31,6 +22,14 @@ __global__ void responseKernel(int64_t *color, int64_t *prop, int64_t *row, int6
   }
 }
 
+template<typename T>
+__global__ void weightedResponseKernel(int64_t *color, int64_t *prop, int64_t *row, int64_t *col,
+                                       T *weight, int64_t *degree, int64_t *cumDegree,
+                                       ptrdiff_t nNodes) {
+  KERNEL_LOOP(i, nNodes) {
+  }
+}
+
 void THCTensor_response(THCState *state, THCudaLongTensor *color, THCudaLongTensor *prop,
                         THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *degree,
                         THCudaLongTensor *cumDegree) {
@@ -40,4 +39,7 @@ void THCTensor_response(THCState *state, THCudaLongTensor *color, THCudaLongTens
              THCudaLongTensor_data(state, degree), THCudaLongTensor_data(state, cumDegree));
 }
 
+#include "generic/THCResponse.cuh"
+#include "THC/THCGenerateAllTypes.h"
+
 #endif  // THC_RESPONSE_INC

aten/THC/generic/THCGraclus.cu

@@ -4,6 +4,15 @@
 
 void THCTensor_(graclus)(THCState *state, THCudaLongTensor *self, THCudaLongTensor *row,
                          THCudaLongTensor *col, THCTensor *weight) {
+  THCAssertSameGPU(THCTensor_(checkGPU)(state, 4, self, row, col, weight));
+
+  THC_TENSOR_GRACLUS(state, self, row,
+    /* while(!THCudaLongTensor_color(state, self)) { */
+      THCudaLongTensor_color(state, self);
+      THCTensor_(propose)(state, self, prop, row, col, weight, degree, cumDegree);
+      THCTensor_(response)(state, self, prop, row, col, weight, degree, cumDegree);
+    /* } */
+  )
 }
 
 #endif  // THC_GENERIC_FILE

aten/THC/generic/THCPropose.cuh

+#ifndef THC_GENERIC_FILE
+#define THC_GENERIC_FILE "generic/THCPropose.cuh"
+#else
+
+void THCTensor_(propose)(THCState *state, THCudaLongTensor *color, THCudaLongTensor *prop,
+                         THCudaLongTensor *row, THCudaLongTensor *col, THCTensor *weight,
+                         THCudaLongTensor *degree, THCudaLongTensor *cumDegree) {
+  KERNEL_RUN(weightedProposeKernel, THCudaLongTensor_nElement(state, color),
+             THCudaLongTensor_data(state, color), THCudaLongTensor_data(state, prop),
+             THCudaLongTensor_data(state, row), THCudaLongTensor_data(state, col),
+             THCTensor_(data)(state, weight), THCudaLongTensor_data(state, degree),
+             THCudaLongTensor_data(state, cumDegree));
+}
+
+#endif  // THC_GENERIC_FILE
+

aten/THC/generic/THCResponse.cuh

+#ifndef THC_GENERIC_FILE
+#define THC_GENERIC_FILE "generic/THCResponse.cuh"
+#else
+
+void THCTensor_(response)(THCState *state, THCudaLongTensor *color, THCudaLongTensor *prop,
+                          THCudaLongTensor *row, THCudaLongTensor *col, THCTensor *weight,
+                          THCudaLongTensor *degree, THCudaLongTensor *cumDegree) {
+  KERNEL_RUN(weightedResponseKernel, THCudaLongTensor_nElement(state, color),
+             THCudaLongTensor_data(state, color), THCudaLongTensor_data(state, prop),
+             THCudaLongTensor_data(state, row), THCudaLongTensor_data(state, col),
+             THCTensor_(data)(state, weight), THCudaLongTensor_data(state, degree),
+             THCudaLongTensor_data(state, cumDegree));
+}
+
+#endif  // THC_GENERIC_FILE