outsourced cuda code

build.py

@@ -20,8 +20,8 @@ with_cuda = False
 if torch.cuda.is_available():
     subprocess.call(['./build.sh', osp.dirname(torch.__file__)])
 
-    headers += ['torch_cluster/src/cuda.h']
-    sources += ['torch_cluster/src/cuda.c']
+    headers += ['torch_cluster/src/{}_cuda.h'.format(f) for f in files]
+    sources += ['torch_cluster/src/{}_cuda.c'.format(f) for f in files]
     include_dirs += ['torch_cluster/kernel']
     define_macros += [('WITH_CUDA', None)]
     extra_objects += ['torch_cluster/build/kernel.so']

torch_cluster/src/generic/cuda.c → torch_cluster/src/generic/grid_cuda.c

 #ifndef THC_GENERIC_FILE
-#define THC_GENERIC_FILE "generic/cuda.c"
+#define THC_GENERIC_FILE "generic/grid_cuda.c"
 #else
 
 void cluster_(grid)(int C, THCudaLongTensor *output, THCTensor *position, THCTensor *size, THCudaLongTensor *count) {

torch_cluster/src/generic/serial_cuda.c

+#ifndef THC_GENERIC_FILE
+#define THC_GENERIC_FILE "generic/serial_cuda.c"
+#else
+
+void cluster_(serial)(THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col, THCTensor *weight, THCudaLongTensor *degree) {
+}
+
+#endif
+

torch_cluster/src/cuda.c → torch_cluster/src/grid_cuda.c

@@ -7,17 +7,17 @@
 
 extern THCState *state;
 
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateFloatType.h"
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateDoubleType.h"
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateByteType.h"
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateCharType.h"
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateShortType.h"
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateIntType.h"
-#include "generic/cuda.c"
+#include "generic/grid_cuda.c"
 #include "THCGenerateLongType.h"

torch_cluster/src/serial_cuda.c

+#include <THC/THC.h>
+
+#include "kernel.h"
+
+#define cluster_(NAME) TH_CONCAT_4(cluster_, NAME, _cuda_, Real)
+#define cluster_kernel_(NAME) TH_CONCAT_4(cluster_, NAME, _kernel_, Real)
+
+extern THCState *state;
+
+void cluster_serial_cuda(THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *degree) {
+}
+
+#include "generic/serial_cuda.c"
+#include "THCGenerateFloatType.h"
+#include "generic/serial_cuda.c"
+#include "THCGenerateDoubleType.h"
+#include "generic/serial_cuda.c"
+#include "THCGenerateByteType.h"
+#include "generic/serial_cuda.c"
+#include "THCGenerateCharType.h"
+#include "generic/serial_cuda.c"
+#include "THCGenerateShortType.h"
+#include "generic/serial_cuda.c"
+#include "THCGenerateIntType.h"
+#include "generic/serial_cuda.c"
+#include "THCGenerateLongType.h"

torch_cluster/src/serial_cuda.h

+void cluster_serial_cuda(THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col, THCudaLongTensor *degree);
+
+void cluster_serial_cuda_Float (THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col,       THCudaTensor *weight, THCudaLongTensor *degree);
+void cluster_serial_cuda_Double(THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col, THCudaDoubleTensor *weight, THCudaLongTensor *degree);
+void cluster_serial_cuda_Byte  (THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaByteTensor *weight, THCudaLongTensor *degree);
+void cluster_serial_cuda_Char  (THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaCharTensor *weight, THCudaLongTensor *degree);
+void cluster_serial_cuda_Short (THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col,  THCudaShortTensor *weight, THCudaLongTensor *degree);
+void cluster_serial_cuda_Int   (THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col,    THCudaIntTensor *weight, THCudaLongTensor *degree);
+void cluster_serial_cuda_Long  (THCudaLongTensor *output, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaLongTensor *weight, THCudaLongTensor *degree);