rename

aten/TH/THGraclus.h

-void THTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col);
+void       THTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col);
 
 void   THByteTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col,   THByteTensor *weight);
 void   THCharTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col,   THCharTensor *weight);

aten/THCC/THCCGraclus.c

+#include <THC/THC.h>
+
+#include "THC.h"
+
+#define THCCTensor_(NAME) TH_CONCAT_4(THCC,Real,Tensor_,NAME)
+
+extern THCState *state;
+
+void THCCTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col) {
+  THCTensor_graclus(state, self, row, col);
+}
+
+#include "generic/THCCGraclus.c"
+#include "THCGenerateAllTypes.h"

aten/THCC/THCCGraclus.h

+void       THCCTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col);
+
+void   THCCByteTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaByteTensor *weight);
+void   THCCCharTensor_graclus(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,   THCudaCharTensor *weight);
+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/THCCGreedy.c

-#include <THC/THC.h>
-
-#include "THCGreedy.h"
-
-#define THCCGreedy_ TH_CONCAT_3(THCC,Real,Greedy)
-#define THCGreedy_ TH_CONCAT_3(TH,CReal,Greedy)
-
-extern THCState *state;
-
-void THCCGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col) {
-  THCGreedy(state, cluster, row, col);
-}
-
-#include "generic/THCCGreedy.c"
-#include "THCGenerateAllTypes.h"

aten/THCC/THCCGreedy.h

-void THCCGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col);
-
-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  THCCHalfGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,    THCudaHalfTensor *weight);
-void  THCCFloatGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,       THCudaTensor *weight);
-void THCCDoubleGreedy(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col, THCudaDoubleTensor *weight);

aten/THCC/THCCGrid.c

 #include <THC/THC.h>
 
-#include "THCGrid.h"
+#include "THC.h"
 
-#define THCCGrid_ TH_CONCAT_3(THCC,Real,Grid)
-#define THCGrid_ TH_CONCAT_3(TH,CReal,Grid)
+#define THCCTensor_(NAME) TH_CONCAT_4(THCC,Real,Tensor_,NAME)
 
 extern THCState *state;
 

aten/THCC/THCCGrid.h

-void   THCCByteGrid(THCudaLongTensor *cluster,   THCudaByteTensor *pos,   THCudaByteTensor *size, THCudaLongTensor *count);
-void   THCCCharGrid(THCudaLongTensor *cluster,   THCudaCharTensor *pos,   THCudaCharTensor *size, THCudaLongTensor *count);
-void  THCCShortGrid(THCudaLongTensor *cluster,  THCudaShortTensor *pos,  THCudaShortTensor *size, THCudaLongTensor *count);
-void    THCCIntGrid(THCudaLongTensor *cluster,    THCudaIntTensor *pos,    THCudaIntTensor *size, THCudaLongTensor *count);
-void   THCCLongGrid(THCudaLongTensor *cluster,   THCudaLongTensor *pos,   THCudaLongTensor *size, THCudaLongTensor *count);
-void  THCCHalfGrid(THCudaLongTensor *cluster,    THCudaHalfTensor *pos,   THCudaHalfTensor *size, THCudaLongTensor *count);
-void  THCCFloatGrid(THCudaLongTensor *cluster,       THCudaTensor *pos,       THCudaTensor *size, THCudaLongTensor *count);
-void THCCDoubleGrid(THCudaLongTensor *cluster, THCudaDoubleTensor *pos, THCudaDoubleTensor *size, THCudaLongTensor *count);
+void   THCCByteTensor_grid(THCudaLongTensor *self,   THCudaByteTensor *pos,   THCudaByteTensor *size, THCudaLongTensor *count);
+void   THCCCharTensor_grid(THCudaLongTensor *self,   THCudaCharTensor *pos,   THCudaCharTensor *size, THCudaLongTensor *count);
+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

+#ifndef THC_GENERIC_FILE
+#define THC_GENERIC_FILE "generic/THCCGraclus.c"
+#else
+
+void THCCTensor_(graclus)(THCudaLongTensor *self, THCudaLongTensor *row, THCudaLongTensor *col,
+                         THCTensor *weight) {
+  THCTensor_(graclus)(state, self, row, col, weight);
+}
+
+#endif

aten/THCC/generic/THCCGreedy.c

-#ifndef THC_GENERIC_FILE
-#define THC_GENERIC_FILE "generic/THCCGreedy.c"
-#else
-
-void THCCGreedy_(THCudaLongTensor *cluster, THCudaLongTensor *row, THCudaLongTensor *col,
-                 THCTensor *weight) {
-  THCGreedy_(state, cluster, row, col, weight);
-}
-
-#endif

aten/THCC/generic/THCCGrid.c

@@ -2,9 +2,9 @@
 #define THC_GENERIC_FILE "generic/THCCGrid.c"
 #else
 
-void THCCGrid_(THCudaLongTensor *cluster, THCTensor *pos, THCTensor *size,
-               THCudaLongTensor *count) {
-  THCGrid_(state, cluster, pos, size, count);
+void THCCTensor_(grid)(THCudaLongTensor *self, THCTensor *pos, THCTensor *size,
+                       THCudaLongTensor *count) {
+  THCTensor_(grid)(state, self, pos, size, count);
 }
 
 #endif

benchmark/ffi.py

+import time
 import torch
 from torch_cluster._ext import ffi
 
-cluster = torch.cuda.LongTensor(5)
-pos = torch.cuda.FloatTensor([[1, 1], [3, 3], [1, 1], [5, 5], [3, 3]])
-size = torch.cuda.FloatTensor([2, 2])
-count = torch.cuda.LongTensor([3, 3])
+cluster = torch.cuda.LongTensor(4)
+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])
+# deg = torch.cuda.LongTensor([2, 3, 3, 2])
 
-func = ffi.THCCFloatGrid
+func = ffi.THCCGreedy
 print(func)
 
-func(cluster, pos, size, count)
+a = 0
+torch.cuda.synchronize()
+t = time.perf_counter()
+# for i in range(100):
+func(cluster, row, col)
+# a += cluster.sum() / cluster.size(0)
+torch.cuda.synchronize()
+print(time.perf_counter() - t)
 print(cluster)

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/THC', 'aten/THCC']
+#     define_macros += [('WITH_CUDA', None)]
+#     extra_objects += ['aten/build/THC.so']
+#     with_cuda = True
 
 ffi = create_extension(
     name='torch_cluster._ext.ffi',