Export symbols (#132)

* Export symbols in the DLL (similar to rusty1s/pytorch_sparse#198)

* Export symbols in the DLL (similar to rusty1s/pytorch_sparse#198)

* Export symbols in the DLL (similar to rusty1s/pytorch_sparse#198)

* Update setup.py

Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>

Author: YoannPitarch

CMakeLists.txt

@@ -15,8 +15,8 @@ endif()
 find_package(Python3 COMPONENTS Development)
 find_package(Torch REQUIRED)
 
-file(GLOB HEADERS csrc/cluster.h)
-file(GLOB OPERATOR_SOURCES csrc/cpu/*.h csrc/cpu/*.cpp csrc/*.cpp)
+file(GLOB HEADERS csrc/*.h)
+file(GLOB OPERATOR_SOURCES csrc/*.* csrc/cpu/*.*)
 if(WITH_CUDA)
   file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} csrc/cuda/*.h csrc/cuda/*.cu)
 endif()

csrc/cluster.h

@@ -2,30 +2,38 @@
 
 #include <torch/extension.h>
 
-int64_t cuda_version();
+#include "macros.h"
 
-torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, double ratio,
+namespace cluster {
+CLUSTER_API int64_t cuda_version() noexcept;
+
+namespace detail {
+CLUSTER_INLINE_VARIABLE int64_t _cuda_version = cuda_version();
+} // namespace detail
+} // namespace cluster
+
+CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, double ratio,
                   bool random_start);
 
-torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
+CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
                       torch::optional<torch::Tensor> optional_weight);
 
-torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
+CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
                    torch::optional<torch::Tensor> optional_start,
                    torch::optional<torch::Tensor> optional_end);
 
-torch::Tensor knn(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
+CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
                   torch::Tensor ptr_y, int64_t k, bool cosine);
 
-torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
+CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
                       torch::Tensor ptr_y);
 
-torch::Tensor radius(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
+CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
                      torch::Tensor ptr_y, double r, int64_t max_num_neighbors);
 
-std::tuple<torch::Tensor, torch::Tensor>
+CLUSTER_API std::tuple<torch::Tensor, torch::Tensor>
 random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
             int64_t walk_length, double p, double q);
 
-torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
+CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
                                int64_t count, double factor);

csrc/cpu/fps_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 torch::Tensor fps_cpu(torch::Tensor src, torch::Tensor ptr, torch::Tensor ratio,
                       bool random_start);

csrc/cpu/graclus_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 torch::Tensor graclus_cpu(torch::Tensor rowptr, torch::Tensor col,
                           torch::optional<torch::Tensor> optional_weight);

csrc/cpu/grid_cpu.h

@@ -1,7 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
-
+#include "../extensions.h"
 torch::Tensor grid_cpu(torch::Tensor pos, torch::Tensor size,
                        torch::optional<torch::Tensor> optional_start,
                        torch::optional<torch::Tensor> optional_end);

csrc/cpu/knn_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
                       torch::optional<torch::Tensor> ptr_x,

csrc/cpu/radius_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
                          torch::optional<torch::Tensor> ptr_x,

csrc/cpu/rw_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::Tensor>
 random_walk_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,

csrc/cpu/sampler_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 torch::Tensor neighbor_sampler_cpu(torch::Tensor start, torch::Tensor rowptr,
                                    int64_t count, double factor);

csrc/cpu/utils.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 #define CHECK_CPU(x) AT_ASSERTM(x.device().is_cpu(), #x " must be CPU tensor")
 #define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")