Merge pull request #68 from liaopeiyuan/cpu_radius

C++ CPU for radius and knn
rusty1s · Jun 22, 2020 · b75a19e · b75a19e
2 parents 32fa325 + 29f9716
commit b75a19e
Show file tree

Hide file tree

Showing 22 changed files with 2,863 additions and 201 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -59,7 +59,8 @@ install:
   - conda install pytorch=${TORCH_VERSION} ${TOOLKIT} -c pytorch --yes
   - source script/torch.sh
   - pip install flake8 codecov
-  - python setup.py install
+  - pip install scipy==1.4.1
+  - pip install .[test]
 script:
   - flake8 .
   - python setup.py test

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.0)
 project(torchcluster)
 set(CMAKE_CXX_STANDARD 14)
-set(TORCHCLUSTER_VERSION 1.5.4)
+set(TORCHCLUSTER_VERSION 1.5.5)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
 

diff --git a/csrc/cpu/knn_cpu.cpp b/csrc/cpu/knn_cpu.cpp
@@ -0,0 +1,54 @@
+#include "knn_cpu.h"
+
+#include "utils.h"
+#include "utils/neighbors.cpp"
+
+torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
+                      torch::optional<torch::Tensor> ptr_x,
+                      torch::optional<torch::Tensor> ptr_y, int64_t k,
+                      int64_t num_workers) {
+
+  CHECK_CPU(x);
+  CHECK_INPUT(x.dim() == 2);
+  CHECK_CPU(y);
+  CHECK_INPUT(y.dim() == 2);
+
+  if (ptr_x.has_value()) {
+    CHECK_CPU(ptr_x.value());
+    CHECK_INPUT(ptr_x.value().dim() == 1);
+  }
+  if (ptr_y.has_value()) {
+    CHECK_CPU(ptr_y.value());
+    CHECK_INPUT(ptr_y.value().dim() == 1);
+  }
+
+  std::vector<size_t> *out_vec = new std::vector<size_t>();
+
+  AT_DISPATCH_ALL_TYPES(x.scalar_type(), "radius_cpu", [&] {
+    auto x_data = x.data_ptr<scalar_t>();
+    auto y_data = y.data_ptr<scalar_t>();
+    auto x_vec = std::vector<scalar_t>(x_data, x_data + x.numel());
+    auto y_vec = std::vector<scalar_t>(y_data, y_data + y.numel());
+
+    if (!ptr_x.has_value()) {
+      nanoflann_neighbors<scalar_t>(y_vec, x_vec, out_vec, 0, x.size(-1), 0,
+                                    num_workers, k, 0);
+    } else {
+      auto sx = (ptr_x.value().narrow(0, 1, ptr_x.value().numel() - 1) -
+                 ptr_x.value().narrow(0, 0, ptr_x.value().numel() - 1));
+      auto sy = (ptr_y.value().narrow(0, 1, ptr_y.value().numel() - 1) -
+                 ptr_y.value().narrow(0, 0, ptr_y.value().numel() - 1));
+      auto sx_data = sx.data_ptr<int64_t>();
+      auto sy_data = sy.data_ptr<int64_t>();
+      auto sx_vec = std::vector<long>(sx_data, sx_data + sx.numel());
+      auto sy_vec = std::vector<long>(sy_data, sy_data + sy.numel());
+      batch_nanoflann_neighbors<scalar_t>(y_vec, x_vec, sy_vec, sx_vec, out_vec,
+                                          k, x.size(-1), 0, k, 0);
+    }
+  });
+
+  const int64_t size = out_vec->size() / 2;
+  auto out = torch::from_blob(out_vec->data(), {size, 2},
+                              x.options().dtype(torch::kLong));
+  return out.t().index_select(0, torch::tensor({1, 0}));
+}
diff --git a/csrc/cpu/knn_cpu.h b/csrc/cpu/knn_cpu.h
@@ -0,0 +1,8 @@
+#pragma once
+
+#include <torch/extension.h>
+
+torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
+                      torch::optional<torch::Tensor> ptr_x,
+                      torch::optional<torch::Tensor> ptr_y, int64_t k,
+                      int64_t num_workers);
diff --git a/csrc/cpu/radius_cpu.cpp b/csrc/cpu/radius_cpu.cpp
@@ -0,0 +1,55 @@
+#include "radius_cpu.h"
+
+#include "utils.h"
+#include "utils/neighbors.cpp"
+
+torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
+                         torch::optional<torch::Tensor> ptr_x,
+                         torch::optional<torch::Tensor> ptr_y, double r,
+                         int64_t max_num_neighbors, int64_t num_workers) {
+
+  CHECK_CPU(x);
+  CHECK_INPUT(x.dim() == 2);
+  CHECK_CPU(y);
+  CHECK_INPUT(y.dim() == 2);
+
+  if (ptr_x.has_value()) {
+    CHECK_CPU(ptr_x.value());
+    CHECK_INPUT(ptr_x.value().dim() == 1);
+  }
+  if (ptr_y.has_value()) {
+    CHECK_CPU(ptr_y.value());
+    CHECK_INPUT(ptr_y.value().dim() == 1);
+  }
+
+  std::vector<size_t> *out_vec = new std::vector<size_t>();
+
+  AT_DISPATCH_ALL_TYPES(x.scalar_type(), "radius_cpu", [&] {
+    auto x_data = x.data_ptr<scalar_t>();
+    auto y_data = y.data_ptr<scalar_t>();
+    auto x_vec = std::vector<scalar_t>(x_data, x_data + x.numel());
+    auto y_vec = std::vector<scalar_t>(y_data, y_data + y.numel());
+
+    if (!ptr_x.has_value()) {
+      nanoflann_neighbors<scalar_t>(y_vec, x_vec, out_vec, r, x.size(-1),
+                                    max_num_neighbors, num_workers, 0, 1);
+    } else {
+      auto sx = (ptr_x.value().narrow(0, 1, ptr_x.value().numel() - 1) -
+                 ptr_x.value().narrow(0, 0, ptr_x.value().numel() - 1));
+      auto sy = (ptr_y.value().narrow(0, 1, ptr_y.value().numel() - 1) -
+                 ptr_y.value().narrow(0, 0, ptr_y.value().numel() - 1));
+      auto sx_data = sx.data_ptr<int64_t>();
+      auto sy_data = sy.data_ptr<int64_t>();
+      auto sx_vec = std::vector<long>(sx_data, sx_data + sx.numel());
+      auto sy_vec = std::vector<long>(sy_data, sy_data + sy.numel());
+      batch_nanoflann_neighbors<scalar_t>(y_vec, x_vec, sy_vec, sx_vec, out_vec,
+                                          r, x.size(-1), max_num_neighbors, 0,
+                                          1);
+    }
+  });
+
+  const int64_t size = out_vec->size() / 2;
+  auto out = torch::from_blob(out_vec->data(), {size, 2},
+                              x.options().dtype(torch::kLong));
+  return out.t().index_select(0, torch::tensor({1, 0}));
+}
diff --git a/csrc/cpu/radius_cpu.h b/csrc/cpu/radius_cpu.h
@@ -0,0 +1,8 @@
+#pragma once
+
+#include <torch/extension.h>
+
+torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
+                         torch::optional<torch::Tensor> ptr_x,
+                         torch::optional<torch::Tensor> ptr_y, double r,
+                         int64_t max_num_neighbors, int64_t num_workers);
diff --git a/csrc/cpu/utils/cloud.h b/csrc/cpu/utils/cloud.h
@@ -0,0 +1,64 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <algorithm>
+#include <cmath>
+#include <iomanip>
+#include <iostream>
+#include <map>
+#include <numeric>
+#include <unordered_map>
+#include <vector>
+
+#include <time.h>
+
+template <typename scalar_t> struct PointCloud {
+  std::vector<std::vector<scalar_t> *> pts;
+
+  void set(std::vector<scalar_t> new_pts, int dim) {
+
+    std::vector<std::vector<scalar_t> *> temp(new_pts.size() / dim);
+    for (size_t i = 0; i < new_pts.size(); i++) {
+      if (i % dim == 0) {
+        std::vector<scalar_t> *point = new std::vector<scalar_t>(dim);
+
+        for (size_t j = 0; j < (size_t)dim; j++) {
+          (*point)[j] = new_pts[i + j];
+        }
+        temp[i / dim] = point;
+      }
+    }
+
+    pts = temp;
+  }
+  void set_batch(std::vector<scalar_t> new_pts, size_t begin, long size,
+                 int dim) {
+    std::vector<std::vector<scalar_t> *> temp(size);
+    for (size_t i = 0; i < (size_t)size; i++) {
+      std::vector<scalar_t> *point = new std::vector<scalar_t>(dim);
+      for (size_t j = 0; j < (size_t)dim; j++) {
+        (*point)[j] = new_pts[dim * (begin + i) + j];
+      }
+
+      temp[i] = point;
+    }
+    pts = temp;
+  }
+
+  // Must return the number of data points.
+  inline size_t kdtree_get_point_count() const { return pts.size(); }
+
+  // Returns the dim'th component of the idx'th point in the class:
+  inline scalar_t kdtree_get_pt(const size_t idx, const size_t dim) const {
+    return (*pts[idx])[dim];
+  }
+
+  // Optional bounding-box computation: return false to default to a standard
+  // bbox computation loop.
+  // Return true if the BBOX was already computed by the class and returned in
+  // "bb" so it can be avoided to redo it again. Look at bb.size() to find out
+  // the expected dimensionality (e.g. 2 or 3 for point clouds)
+  template <class BBOX> bool kdtree_get_bbox(BBOX & /* bb */) const {
+    return false;
+  }
+};