rebuild to new extension api

Author: rusty1s

.travis.yml

@@ -6,14 +6,23 @@ matrix:
     - python: 2.7
     - python: 3.5
     - python: 3.6
+addons:
+  apt:
+    sources:
+      - ubuntu-toolchain-r-test
+    packages:
+      - gcc-4.9
+      - g++-4.9
+before_install:
+  - export CC="gcc-4.9"
+  - export CXX="g++-4.9"
 install:
   - if [[ $TRAVIS_PYTHON_VERSION == 2.7 ]]; then pip install http://download.pytorch.org/whl/cpu/torch-0.4.1-cp27-cp27mu-linux_x86_64.whl; fi
   - if [[ $TRAVIS_PYTHON_VERSION == 3.5 ]]; then pip install http://download.pytorch.org/whl/cpu/torch-0.4.1-cp35-cp35m-linux_x86_64.whl; fi
   - if [[ $TRAVIS_PYTHON_VERSION == 3.6 ]]; then pip install http://download.pytorch.org/whl/cpu/torch-0.4.1-cp36-cp36m-linux_x86_64.whl; fi
   - pip install pycodestyle
   - pip install flake8
   - pip install codecov
-  - pip install cffi
 script:
   - pycodestyle .
   - flake8 .

MANIFEST.in

@@ -1,7 +1,3 @@
 include LICENSE
-include build.py
-include build.sh
-
-recursive-include aten *
-
-recursive-exclude torch_cluster/_ext *
+recursive-include cpu *
+recursive-include cuda *

README.md

@@ -39,10 +39,11 @@ $ echo $CPATH
 Then run:
 
 ```
-pip install cffi torch-cluster
+pip install torch-cluster
 ```
 
 If you are running into any installation problems, please create an [issue](https://github.com/rusty1s/pytorch_cluster/issues).
+Be sure to import `torch` first before using this package to resolve symbols the dynamic linker must see.
 
 ## Graclus
 
@@ -62,7 +63,7 @@ cluster = graclus_cluster(row, col, weight)
 
 ```
 print(cluster)
-tensor([ 0,  0,  1])
+tensor([0, 0, 1])
 ```
 
 ## VoxelGrid
@@ -81,7 +82,7 @@ cluster = grid_cluster(pos, size)
 
 ```
 print(cluster)
-tensor([ 0,  5,  3,  0,  1])
+tensor([0, 5, 3, 0, 1])
 ```
 
 ## Running tests

build.py

@@ -0,0 +1,48 @@
+#include <torch/torch.h>
+
+// #include "../include/degree.cpp"
+// #include "../include/loop.cpp"
+// #include "../include/perm.cpp"
+
+at::Tensor graclus(at::Tensor row, at::Tensor col, int64_t num_nodes) {
+  // std::tie(row, col) = remove_self_loops(row, col);
+  // std::tie(row, col) = randperm(row, col, num_nodes);
+  // auto deg = degree(row, num_nodes, row.type().scalarType());
+
+  auto cluster = at::full(num_nodes, -1, row.options());
+
+  // auto *row_data = row.data<int64_t>();
+  // auto *col_data = col.data<int64_t>();
+  // auto *deg_data = deg.data<int64_t>();
+  // auto *cluster_data = cluster.data<int64_t>();
+
+  // int64_t e_idx = 0, d_idx, r, c;
+  // while (e_idx < row.size(0)) {
+  //   r = row_data[e_idx];
+  //   if (cluster_data[r] < 0) {
+  //     cluster_data[r] = r;
+  //     for (d_idx = 0; d_idx < deg_data[r]; d_idx++) {
+  //       c = col_data[e_idx + d_idx];
+  //       if (cluster_data[c] < 0) {
+  //         cluster_data[r] = std::min(r, c);
+  //         cluster_data[c] = std::min(r, c);
+  //         break;
+  //       }
+  //     }
+  //   }
+  //   e_idx += deg_data[r];
+  // }
+
+  return cluster;
+}
+
+at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor weight,
+                            int64_t num_nodes) {
+  auto cluster = at::full(num_nodes, -1, row.options());
+  return cluster;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("graclus", &graclus, "Graclus (CPU)");
+  m.def("weighted_graclus", &weighted_graclus, "Weighted Graclus (CPU)");
+}

build.sh

@@ -0,0 +1,23 @@
+#include <torch/torch.h>
+
+at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start,
+                at::Tensor end) {
+  size = size.toType(pos.type());
+  start = start.toType(pos.type());
+  end = end.toType(pos.type());
+
+  pos = pos - start.view({1, -1});
+  auto num_voxels = ((end - start) / size).toType(at::kLong);
+  num_voxels = (num_voxels + 1).cumsum(0);
+  num_voxels -= num_voxels.data<int64_t>()[0];
+  num_voxels.data<int64_t>()[0] = 1;
+
+  auto cluster = pos / size.view({1, -1});
+  cluster = cluster.toType(at::kLong);
+  cluster *= num_voxels.view({1, -1});
+  cluster = cluster.sum(1);
+
+  return cluster;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("grid", &grid, "Grid (CPU)"); }

cpu/graclus.cpp

@@ -0,0 +1,15 @@
+#include <torch/torch.h>
+
+at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start,
+                at::Tensor end);
+
+at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes);
+
+at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor weight,
+                            int num_nodes);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("grid", &grid, "Grid (CUDA)");
+  m.def("graclus", &graclus, "Graclus (CUDA)");
+  m.def("weighted_graclus", &weighted_graclus, "Weightes Graclus (CUDA)");
+}

cpu/grid.cpp

@@ -0,0 +1,27 @@
+#include <torch/torch.h>
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
+
+at::Tensor graclus_cuda(at::Tensor row, at::Tensor col, int64_t num_nodes);
+
+at::Tensor weighted_graclus_cuda(at::Tensor row, at::Tensor col,
+                                 at::Tensor weight, int64_t num_nodes);
+
+at::Tensor graclus(at::Tensor row, at::Tensor col, int64_t num_nodes) {
+  CHECK_CUDA(row);
+  CHECK_CUDA(col);
+  return graclus_cuda(row, col, num_nodes);
+}
+
+at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor weight,
+                            int64_t num_nodes) {
+  CHECK_CUDA(row);
+  CHECK_CUDA(col);
+  CHECK_CUDA(weight);
+  return graclus_cuda(row, col, num_nodes);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("graclus", &graclus, "Graclus (CUDA)");
+  m.def("weighted_graclus", &weighted_graclus, "Weighted Graclus (CUDA)");
+}

cuda/cluster.cpp

@@ -0,0 +1,73 @@
+#include <ATen/ATen.h>
+
+at::Tensor graclus_cuda(at::Tensor row, at::Tensor col, int64_t num_nodes) {
+  auto cluster = at::full(num_nodes, -1, row.options());
+  return cluster;
+}
+
+at::Tensor weighted_graclus_cuda(at::Tensor row, at::Tensor col,
+                                 at::Tensor weight, int64_t num_nodes) {
+  auto cluster = at::full(num_nodes, -1, row.options());
+  return cluster;
+}
+
+// #include "color.cuh"
+// #include "common.cuh"
+
+// at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes) {
+//   // Remove self-loops.
+//   auto mask = row != col;
+//   row = row.masked_select(mask);
+//   col.masked_select(mask);
+
+//   // Sort by row index.
+//   at::Tensor perm;
+//   std::tie(row, perm) = row.sort();
+//   col = col.index_select(0, perm);
+
+//   // Generate helper vectors.
+//   auto cluster = at::full(row.type(), {num_nodes}, -1);
+//   auto prop = at::full(row.type(), {num_nodes}, -1);
+//   auto deg = degree(row, num_nodes);
+//   auto cum_deg = deg.cumsum(0);
+
+//   color(cluster);
+
+//   /* while (!color(cluster)) { */
+//   /*   propose(cluster, prop, row, col, weight, deg, cum_deg); */
+//   /*   response(cluster, prop, row, col, weight, deg, cum_deg); */
+//   /* } */
+
+//   return cluster;
+// }
+
+// at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor
+// weight,
+//                             int num_nodes) {
+//   // Remove self-loops.
+//   auto mask = row != col;
+//   row = row.masked_select(mask);
+//   col = col.masked_select(mask);
+//   weight = weight.masked_select(mask);
+
+//   // Sort by row index.
+//   at::Tensor perm;
+//   std::tie(row, perm) = row.sort();
+//   col = col.index_select(0, perm);
+//   weight = weight.index_select(0, perm);
+
+//   // Generate helper vectors.
+//   auto cluster = at::full(row.type(), {num_nodes}, -1);
+//   auto prop = at::full(row.type(), {num_nodes}, -1);
+//   auto deg = degree(row, num_nodes);
+//   auto cum_deg = deg.cumsum(0);
+
+//   color(cluster);
+
+//   /* while (!color(cluster)) { */
+//   /*   weighted_propose(cluster, prop, row, col, weight, deg, cum_deg); */
+//   /*   weighted_response(cluster, prop, row, col, weight, deg, cum_deg); */
+//   /* } */
+
+//   return cluster;
+// }