setup for cuda and cpu

Author: rusty1s

aten/cluster.py

@@ -0,0 +1,32 @@
+import torch
+
+import cluster_cpu
+import cluster_cuda
+
+
+def grid(pos, size, start=None, end=None):
+    lib = cluster_cuda if pos.is_cuda else cluster_cpu
+    start = pos.t().min(dim=1)[0] if start is None else start
+    end = pos.t().max(dim=1)[0] if end is None else end
+    return lib.grid(pos, size, start, end)
+
+
+def graclus(row, col, num_nodes):
+    return cluster_cpu.graclus(row, col, num_nodes)
+
+
+device = torch.device('cuda')
+pos = torch.tensor([[1, 1], [3, 3], [5, 5], [7, 7]], device=device)
+size = torch.tensor([2, 2], device=device)
+print('pos', pos.tolist())
+print('size', size.tolist())
+cluster = grid(pos, size)
+print('result', cluster.tolist(), cluster.dtype, cluster.device)
+
+row = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 3, 3])
+col = torch.tensor([1, 2, 0, 2, 3, 0, 1, 3, 1, 2])
+print(row)
+print(col)
+print('-----------------')
+cluster = graclus(row, col, 4)
+print(cluster)

aten/cpu/cluster.py

@@ -1,8 +1,8 @@
 #include <torch/torch.h>
 
-#include "degree.cpp"
-#include "loop.cpp"
-#include "perm.cpp"
+#include "../include/degree.cpp"
+#include "../include/loop.cpp"
+#include "../include/perm.cpp"
 
 at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes) {
   std::tie(row, col) = remove_self_loops(row, col);

aten/cpu/graclus.cpp

@@ -1,5 +1,4 @@
-#ifndef DEGREE_CPP
-#define DEGREE_CPP
+#include "degree.h"
 
 #include <torch/torch.h>
 
@@ -9,5 +8,3 @@ inline at::Tensor degree(at::Tensor index, int num_nodes,
   auto one = at::full(zero.type(), {index.size(0)}, 1);
   return zero.scatter_add_(0, index, one);
 }
-
-#endif // DEGREE_CPP

aten/cpu/setup.py

@@ -0,0 +1,9 @@
+#ifndef DEGREE_INC
+#define DEGREE_INC
+
+#include <torch/torch.h>
+
+inline at::Tensor degree(at::Tensor index, int num_nodes,
+                         at::ScalarType scalar_type);
+
+#endif // DEGREE_INC

aten/cuda/cluster.py

@@ -1,5 +1,4 @@
-#ifndef LOOP_CPP
-#define LOOP_CPP
+#include "loop.h"
 
 #include <torch/torch.h>
 
@@ -8,5 +7,3 @@ inline std::tuple<at::Tensor, at::Tensor> remove_self_loops(at::Tensor row,
   auto mask = row != col;
   return {row.masked_select(mask), col.masked_select(mask)};
 }
-
-#endif // LOOP_CPP

aten/cuda/setup.py

@@ -0,0 +1,9 @@
+#ifndef LOOP_INC
+#define LOOP_INC
+
+#include <torch/torch.h>
+
+inline std::tuple<at::Tensor, at::Tensor> remove_self_loops(at::Tensor row,
+                                                            at::Tensor col);
+
+#endif // LOOP_INC