Export symbols (#198)

dfalbel · web-flow · commit 3bf43eb09a68 · 2022-02-01T08:53:33.000+01:00
* Mark exported symbols with `SPARSE_API` so they are available in the DLL on WIndows.

* Export symbols by default in the Python library.

* Sync headers with implementation.

* Include the `SPARSE_API` macro.

* Fix linting issue.
diff --git a/csrc/convert.cpp b/csrc/convert.cpp
@@ -19,7 +19,7 @@ PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; }
 #endif
 #endif
 
-torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
+SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
   if (ind.device().is_cuda()) {
 #ifdef WITH_CUDA
     return ind2ptr_cuda(ind, M);
@@ -31,7 +31,7 @@ torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
   }
 }
 
-torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) {
+SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) {
   if (ptr.device().is_cuda()) {
 #ifdef WITH_CUDA
     return ptr2ind_cuda(ptr, E);
diff --git a/csrc/diag.cpp b/csrc/diag.cpp
@@ -19,7 +19,7 @@ PyMODINIT_FUNC PyInit__diag_cpu(void) { return NULL; }
 #endif
 #endif
 
-torch::Tensor non_diag_mask(torch::Tensor row, torch::Tensor col, int64_t M,
+SPARSE_API torch::Tensor non_diag_mask(torch::Tensor row, torch::Tensor col, int64_t M,
                             int64_t N, int64_t k) {
   if (row.device().is_cuda()) {
 #ifdef WITH_CUDA
diff --git a/csrc/ego_sample.cpp b/csrc/ego_sample.cpp
@@ -16,7 +16,7 @@ PyMODINIT_FUNC PyInit__ego_sample_cpu(void) { return NULL; }
 #endif
 
 // Returns `rowptr`, `col`, `n_id`, `e_id`, `ptr`, `root_n_id`
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
            torch::Tensor, torch::Tensor>
 ego_k_hop_sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
                      int64_t depth, int64_t num_neighbors, bool replace) {
diff --git a/csrc/extensions.h b/csrc/extensions.h
@@ -1,5 +1,6 @@
 
 #include <torch/torch.h>
+#include "sparse.h"
 
 // for getpid()
 #ifdef _WIN32
diff --git a/csrc/hgt_sample.cpp b/csrc/hgt_sample.cpp
@@ -16,7 +16,7 @@ PyMODINIT_FUNC PyInit__hgt_sample_cpu(void) { return NULL; }
 #endif
 
 // Returns 'output_node_dict', 'row_dict', 'col_dict', 'output_edge_dict'
-std::tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
+SPARSE_API std::tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
            c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
 hgt_sample(const c10::Dict<std::string, torch::Tensor> &colptr_dict,
            const c10::Dict<std::string, torch::Tensor> &row_dict,
diff --git a/csrc/metis.cpp b/csrc/metis.cpp
@@ -15,7 +15,7 @@ PyMODINIT_FUNC PyInit__metis_cpu(void) { return NULL; }
 #endif
 #endif
 
-torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
                         torch::optional<torch::Tensor> optional_value,
                         int64_t num_parts, bool recursive) {
   if (rowptr.device().is_cuda()) {
@@ -30,7 +30,7 @@ torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
   }
 }
 
-torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
                          torch::optional<torch::Tensor> optional_value,
                          torch::optional<torch::Tensor> optional_node_weight,
                          int64_t num_parts, bool recursive) {
@@ -46,7 +46,7 @@ torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
   }
 }
 
-torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
                            torch::optional<torch::Tensor> optional_value,
                            torch::optional<torch::Tensor> optional_node_weight,
                            int64_t num_parts, bool recursive,
diff --git a/csrc/neighbor_sample.cpp b/csrc/neighbor_sample.cpp
@@ -16,7 +16,7 @@ PyMODINIT_FUNC PyInit__neighbor_sample_cpu(void) { return NULL; }
 #endif
 
 // Returns 'output_node', 'row', 'col', 'output_edge'
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 neighbor_sample(const torch::Tensor &colptr, const torch::Tensor &row,
                 const torch::Tensor &input_node,
                 const std::vector<int64_t> num_neighbors, const bool replace,
@@ -25,7 +25,7 @@ neighbor_sample(const torch::Tensor &colptr, const torch::Tensor &row,
                              directed);
 }
 
-std::tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
+SPARSE_API std::tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
            c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
 hetero_neighbor_sample(
     const std::vector<node_t> &node_types,
diff --git a/csrc/relabel.cpp b/csrc/relabel.cpp
@@ -15,7 +15,7 @@ PyMODINIT_FUNC PyInit__relabel_cpu(void) { return NULL; }
 #endif
 #endif
 
-std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
                                                  torch::Tensor idx) {
   if (col.device().is_cuda()) {
 #ifdef WITH_CUDA
@@ -28,7 +28,7 @@ std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
   }
 }
 
-std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
            torch::Tensor>
 relabel_one_hop(torch::Tensor rowptr, torch::Tensor col,
                 torch::optional<torch::Tensor> optional_value,
diff --git a/csrc/rw.cpp b/csrc/rw.cpp
@@ -19,7 +19,7 @@ PyMODINIT_FUNC PyInit__rw_cpu(void) { return NULL; }
 #endif
 #endif
 
-torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
                           torch::Tensor start, int64_t walk_length) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
diff --git a/csrc/sample.cpp b/csrc/sample.cpp
@@ -15,7 +15,7 @@ PyMODINIT_FUNC PyInit__sample_cpu(void) { return NULL; }
 #endif
 #endif
 
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
            int64_t num_neighbors, bool replace) {
   if (rowptr.device().is_cuda()) {
diff --git a/csrc/sparse.h b/csrc/sparse.h
@@ -2,69 +2,82 @@
 
 #include <torch/library.h>
 
-int64_t cuda_version();
+#ifdef _WIN32
+#if defined(torchsparse_EXPORTS)
+#define SPARSE_API __declspec(dllexport)
+#else
+#define SPARSE_API __declspec(dllimport)
+#endif
+#else
+#define SPARSE_API
+#endif
 
-torch::Tensor ind2ptr(torch::Tensor ind, int64_t M);
-torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E);
 
-torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API int64_t cuda_version();
+
+SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M);
+SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E);
+
+SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
                         torch::optional<torch::Tensor> optional_value,
                         int64_t num_parts, bool recursive);
 
-torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
                          torch::optional<torch::Tensor> optional_value,
                          torch::optional<torch::Tensor> optional_node_weight,
                          int64_t num_parts, bool recursive);
 
-torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
                            torch::optional<torch::Tensor> optional_value,
-                           int64_t num_parts, bool recursive);
-
-std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
+                           torch::optional<torch::Tensor> optional_node_weight,
+                           int64_t num_parts, bool recursive,
+                           int64_t num_workers);
+ 
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
                                                  torch::Tensor idx);
 
-std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
            torch::Tensor>
 relabel_one_hop(torch::Tensor rowptr, torch::Tensor col,
                 torch::optional<torch::Tensor> optional_value,
-                torch::Tensor idx);
+                torch::Tensor idx, bool bipartite);
 
-torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
+SPARSE_API torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
                           torch::Tensor start, int64_t walk_length);
 
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
 subgraph(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row,
          torch::Tensor col);
 
-std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
-sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor rowcount,
-           torch::Tensor idx, int64_t num_neighbors, bool replace);
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
+           int64_t num_neighbors, bool replace);
 
-torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
                        torch::Tensor rowptr, torch::Tensor col,
                        torch::optional<torch::Tensor> opt_value,
                        torch::optional<torch::Tensor> opt_colptr,
                        torch::optional<torch::Tensor> opt_csr2csc,
                        torch::Tensor mat);
 
-torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
                         torch::Tensor rowptr, torch::Tensor col,
                         torch::optional<torch::Tensor> opt_value,
                         torch::optional<torch::Tensor> opt_rowcount,
                         torch::optional<torch::Tensor> opt_colptr,
                         torch::optional<torch::Tensor> opt_csr2csc,
                         torch::Tensor mat);
 
-std::tuple<torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_min(torch::Tensor rowptr, torch::Tensor col,
          torch::optional<torch::Tensor> opt_value, torch::Tensor mat);
 
-std::tuple<torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_max(torch::Tensor rowptr, torch::Tensor col,
          torch::optional<torch::Tensor> opt_value, torch::Tensor mat);
 
-std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
 spspmm_sum(torch::Tensor rowptrA, torch::Tensor colA,
            torch::optional<torch::Tensor> optional_valueA,
            torch::Tensor rowptrB, torch::Tensor colB,
-           torch::optional<torch::Tensor> optional_valueB, int64_t K);
+           torch::optional<torch::Tensor> optional_valueB, int64_t K); 
diff --git a/csrc/spmm.cpp b/csrc/spmm.cpp
@@ -302,7 +302,7 @@ class SPMMMax : public torch::autograd::Function<SPMMMax> {
   }
 };
 
-torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
                        torch::Tensor rowptr, torch::Tensor col,
                        torch::optional<torch::Tensor> opt_value,
                        torch::optional<torch::Tensor> opt_colptr,
@@ -313,7 +313,7 @@ torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
                         mat, opt_value.has_value())[0];
 }
 
-torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
                         torch::Tensor rowptr, torch::Tensor col,
                         torch::optional<torch::Tensor> opt_value,
                         torch::optional<torch::Tensor> opt_rowcount,
@@ -325,15 +325,15 @@ torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
                          opt_csr2csc, mat, opt_value.has_value())[0];
 }
 
-std::tuple<torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_min(torch::Tensor rowptr, torch::Tensor col,
          torch::optional<torch::Tensor> opt_value, torch::Tensor mat) {
   auto value = opt_value.has_value() ? opt_value.value() : col;
   auto result = SPMMMin::apply(rowptr, col, value, mat, opt_value.has_value());
   return std::make_tuple(result[0], result[1]);
 }
 
-std::tuple<torch::Tensor, torch::Tensor>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_max(torch::Tensor rowptr, torch::Tensor col,
          torch::optional<torch::Tensor> opt_value, torch::Tensor mat) {
   auto value = opt_value.has_value() ? opt_value.value() : col;
diff --git a/csrc/spspmm.cpp b/csrc/spspmm.cpp
@@ -19,7 +19,7 @@ PyMODINIT_FUNC PyInit__spspmm_cpu(void) { return NULL; }
 #endif
 #endif
 
-std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
 spspmm_sum(torch::Tensor rowptrA, torch::Tensor colA,
            torch::optional<torch::Tensor> optional_valueA,
            torch::Tensor rowptrB, torch::Tensor colB,
diff --git a/csrc/version.cpp b/csrc/version.cpp
@@ -2,6 +2,7 @@
 #include <Python.h>
 #endif
 #include <torch/script.h>
+#include "sparse.h"
 
 #ifdef WITH_CUDA
 #include <cuda.h>
@@ -17,7 +18,7 @@ PyMODINIT_FUNC PyInit__version_cpu(void) { return NULL; }
 #endif
 #endif
 
-int64_t cuda_version() {
+SPARSE_API int64_t cuda_version() {
 #ifdef WITH_CUDA
   return CUDA_VERSION;
 #else
diff --git a/setup.py b/setup.py
@@ -35,6 +35,10 @@ def get_extensions():
 
     for main, suffix in product(main_files, suffices):
         define_macros = [('WITH_PYTHON', None)]
+
+        if sys.platform == 'win32':
+            define_macros += [('torchsparse_EXPORTS', None)]
+
         libraries = []
         if WITH_METIS:
             define_macros += [('WITH_METIS', None)]

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ PyMODINIT_FUNC PyInit__metis_cpu(void) { return NULL; }`
`15`	`15`	`#endif`
`16`	`16`	`#endif`
`17`	`17`
`18`		`-torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,`
	`18`	`+SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,`
`19`	`19`	`torch::optional<torch::Tensor> optional_value,`
`20`	`20`	`int64_t num_parts, bool recursive) {`
`21`	`21`	`if (rowptr.device().is_cuda()) {`
`@@ -30,7 +30,7 @@ torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,`
`30`	`30`	`}`
`31`	`31`	`}`
`32`	`32`
`33`		`-torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,`
	`33`	`+SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,`
`34`	`34`	`torch::optional<torch::Tensor> optional_value,`
`35`	`35`	`torch::optional<torch::Tensor> optional_node_weight,`
`36`	`36`	`int64_t num_parts, bool recursive) {`
`@@ -46,7 +46,7 @@ torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,`
`46`	`46`	`}`
`47`	`47`	`}`
`48`	`48`
`49`		`-torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,`
	`49`	`+SPARSE_API torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,`
`50`	`50`	`torch::optional<torch::Tensor> optional_value,`
`51`	`51`	`torch::optional<torch::Tensor> optional_node_weight,`
`52`	`52`	`int64_t num_parts, bool recursive,`