Export symbols (#198)

* 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.

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);

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

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) {

csrc/extensions.h

 
 #include <torch/torch.h>
+#include "sparse.h"
 
 // for getpid()
 #ifdef _WIN32

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,

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,

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,

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,

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

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()) {

csrc/sparse.h

@@ -2,52 +2,65 @@
 
 #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,
@@ -55,16 +68,16 @@ torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
                         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); 

csrc/spmm.cpp

@@ -302,7 +302,7 @@ public:
   }
 };
 
-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,7 +325,7 @@ 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;
@@ -333,7 +333,7 @@ spmm_min(torch::Tensor rowptr, torch::Tensor col,
   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;

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,

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

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)]