add capability to optionally include python during compile (#136)

rusty1s · Jul 23, 2022 · c77ed13 · c77ed13
1 parent 113e8a6
commit c77ed13
Show file tree

Hide file tree

Showing 13 changed files with 48 additions and 7 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -4,6 +4,7 @@ set(CMAKE_CXX_STANDARD 14)
 set(TORCHCLUSTER_VERSION 1.6.0)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
+option(WITH_PYTHON "Link to Python when building" ON)
 
 if(WITH_CUDA)
   enable_language(CUDA)
@@ -12,7 +13,10 @@ if(WITH_CUDA)
   set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
 endif()
 
-find_package(Python3 COMPONENTS Development)
+if (WITH_PYTHON)
+  add_definitions(-DWITH_PYTHON)
+  find_package(Python3 COMPONENTS Development)
+endif()
 find_package(Torch REQUIRED)
 
 file(GLOB HEADERS csrc/*.h)
@@ -22,7 +26,10 @@ if(WITH_CUDA)
 endif()
 
 add_library(${PROJECT_NAME} SHARED ${OPERATOR_SOURCES})
-target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python)
+target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES})
+if (WITH_PYTHON)
+  target_link_libraries(${PROJECT_NAME} PRIVATE Python3::Python)
+endif()
 set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchCluster)
 
 target_include_directories(${PROJECT_NAME} INTERFACE

diff --git a/csrc/cluster.h b/csrc/cluster.h
@@ -1,8 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
-
-#include "macros.h"
+#include "extensions.h"
 
 namespace cluster {
 CLUSTER_API int64_t cuda_version() noexcept;

diff --git a/csrc/extensions.h b/csrc/extensions.h
@@ -1,2 +1,2 @@
 #include "macros.h"
-#include <torch/extension.h>
+#include <torch/torch.h>
diff --git a/csrc/fps.cpp b/csrc/fps.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/fps_cpu.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__fps_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__fps_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, torch::Tensor ratio,
                   bool random_start) {

diff --git a/csrc/graclus.cpp b/csrc/graclus.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/graclus_cpu.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__graclus_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__graclus_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
                       torch::optional<torch::Tensor> optional_weight) {

diff --git a/csrc/grid.cpp b/csrc/grid.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/grid_cpu.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__grid_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__grid_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
                    torch::optional<torch::Tensor> optional_start,

diff --git a/csrc/knn.cpp b/csrc/knn.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/knn_cpu.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__knn_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__knn_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
                   torch::optional<torch::Tensor> ptr_x,

diff --git a/csrc/nearest.cpp b/csrc/nearest.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "extensions.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__nearest_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__nearest_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
                       torch::Tensor ptr_y) {

diff --git a/csrc/radius.cpp b/csrc/radius.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/radius_cpu.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__radius_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__radius_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y,
                      torch::optional<torch::Tensor> ptr_x,

diff --git a/csrc/rw.cpp b/csrc/rw.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/rw_cpu.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__rw_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__rw_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API std::tuple<torch::Tensor, torch::Tensor>
 random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,

diff --git a/csrc/sampler.cpp b/csrc/sampler.cpp
@@ -1,15 +1,19 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 
 #include "cpu/sampler_cpu.h"
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__sampler_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__sampler_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 CLUSTER_API torch::Tensor neighbor_sampler(torch::Tensor start, torch::Tensor rowptr,
                                int64_t count, double factor) {

diff --git a/csrc/version.cpp b/csrc/version.cpp
@@ -1,4 +1,6 @@
+#ifdef WITH_PYTHON
 #include <Python.h>
+#endif
 #include <torch/script.h>
 #include "cluster.h"
 #include "macros.h"
@@ -8,12 +10,14 @@
 #endif
 
 #ifdef _WIN32
+#ifdef WITH_PYTHON
 #ifdef WITH_CUDA
 PyMODINIT_FUNC PyInit__version_cuda(void) { return NULL; }
 #else
 PyMODINIT_FUNC PyInit__version_cpu(void) { return NULL; }
 #endif
 #endif
+#endif
 
 
 namespace cluster {

diff --git a/setup.py b/setup.py
@@ -33,7 +33,7 @@ def get_extensions():
     main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
 
     for main, suffix in product(main_files, suffices):
-        define_macros = []
+        define_macros = [('WITH_PYTHON', None)]
 
         if sys.platform == 'win32':
             define_macros += [('torchcluster_EXPORTS', None)]