rename and bugfixes

aten/THC/generic/THCGrid.h

@@ -2,7 +2,7 @@
 #define THC_GENERIC_FILE "generic/THCGrid.h"
 #else
 
-void THCGrid_(THCState *state, THCudaLongTensor *cluster, THCTensor *pos, THCTensor *size,
-              THCudaLongTensor *count);
+void THCTensor_(grid)(THCState *state, THCudaLongTensor *self, THCTensor *pos, THCTensor *size,
+                      THCudaLongTensor *count);
 
 #endif  // THC_GENERIC_FILE

aten/THCC/THCCGraclus.h

@@ -5,6 +5,5 @@ void   THCCCharTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THC
 void  THCCShortTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,  THCudaShortTensor *weight);
 void    THCCIntTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,    THCudaIntTensor *weight);
 void   THCCLongTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaLongTensor *weight);
-void   THCCHalfTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaHalfTensor *weight);
 void  THCCFloatTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,       THCudaTensor *weight);
 void THCCDoubleTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col, THCudaDoubleTensor *weight);

aten/THCC/THCCGrid.h

@@ -3,6 +3,5 @@ void   THCCCharTensor_grid(THCudaLongTensor *self,   THCudaCharTensor *pos,   TH
 void  THCCShortTensor_grid(THCudaLongTensor *self,  THCudaShortTensor *pos,  THCudaShortTensor *size, THCudaLongTensor *count);
 void    THCCIntTensor_grid(THCudaLongTensor *self,    THCudaIntTensor *pos,    THCudaIntTensor *size, THCudaLongTensor *count);
 void   THCCLongTensor_grid(THCudaLongTensor *self,   THCudaLongTensor *pos,   THCudaLongTensor *size, THCudaLongTensor *count);
-void   THCCHalfTensor_grid(THCudaLongTensor *self,   THCudaHalfTensor *pos,   THCudaHalfTensor *size, THCudaLongTensor *count);
 void  THCCFloatTensor_grid(THCudaLongTensor *self,       THCudaTensor *pos,       THCudaTensor *size, THCudaLongTensor *count);
 void THCCDoubleTensor_grid(THCudaLongTensor *self, THCudaDoubleTensor *pos, THCudaDoubleTensor *size, THCudaLongTensor *count);

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

aten/THCC/generic/THCCGrid.c

@@ -4,7 +4,7 @@
 
 void THCCTensor_(grid)(THCudaLongTensor *self, THCTensor *pos, THCTensor *size,
                        THCudaLongTensor *count) {
-  THCTensor_(grid)(state, self, pos, size, count);
+  /* THCTensor_(grid)(state, self, pos, size, count); */
 }
 
 #endif

build.py

@@ -17,15 +17,15 @@ define_macros = []
 extra_objects = []
 with_cuda = False
 
-# if torch.cuda.is_available():
-#     subprocess.call(['./build_new.sh', osp.dirname(torch.__file__)])
-
-#     headers += ['aten/THCC/THCC{}.h'.format(f) for f in files]
-#     sources += ['aten/THCC/THCC{}.c'.format(f) for f in files]
-#     include_dirs += ['aten/THC', 'aten/THCC']
-#     define_macros += [('WITH_CUDA', None)]
-#     extra_objects += ['aten/build/THC.so']
-#     with_cuda = True
+if torch.cuda.is_available():
+    subprocess.call(['./build_new.sh', osp.dirname(torch.__file__)])
+
+    headers += ['aten/THCC/THCC{}.h'.format(f) for f in files]
+    sources += ['aten/THCC/THCC{}.c'.format(f) for f in files]
+    include_dirs += ['aten/THCC']
+    define_macros += [('WITH_CUDA', None)]
+    extra_objects += ['aten/build/THC.so']
+    with_cuda = True
 
 ffi = create_extension(
     name='torch_cluster._ext.ffi',

test/test_graclus.py

@@ -48,3 +48,11 @@ def test_graclus_cluster_cpu(tensor, i):
 
     cluster = graclus_cluster(row, col, weight)
     assert_correct_graclus(row, col, cluster)
+
+
+def test_graclus_cluster_gpu():
+    row = torch.cuda.LongTensor([0, 0, 1, 1, 1, 2, 2, 2, 3, 3])
+    col = torch.cuda.LongTensor([1, 2, 0, 2, 3, 0, 1, 3, 1, 2])
+
+    cluster = graclus_cluster(row, col)
+    print(cluster.cpu().tolist())