From c77ed131b667850edb12554d57d455ab60d24bb9 Mon Sep 17 00:00:00 2001 From: Gerico Vidanes <65466924+GericoVi@users.noreply.github.com> Date: Sat, 23 Jul 2022 05:33:09 +0100 Subject: [PATCH] add capability to optionally include python during compile (#136) --- CMakeLists.txt | 11 +++++++++-- csrc/cluster.h | 4 +--- csrc/extensions.h | 2 +- csrc/fps.cpp | 4 ++++ csrc/graclus.cpp | 4 ++++ csrc/grid.cpp | 4 ++++ csrc/knn.cpp | 4 ++++ csrc/nearest.cpp | 4 ++++ csrc/radius.cpp | 4 ++++ csrc/rw.cpp | 4 ++++ csrc/sampler.cpp | 4 ++++ csrc/version.cpp | 4 ++++ setup.py | 2 +- 13 files changed, 48 insertions(+), 7 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8875580..e9f8c7f 100644 --- 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 index def1655..ac8315b 100644 --- a/csrc/cluster.h +++ b/csrc/cluster.h @@ -1,8 +1,6 @@ #pragma once -#include - -#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 index 5644db0..91c4df1 100644 --- a/csrc/extensions.h +++ b/csrc/extensions.h @@ -1,2 +1,2 @@ #include "macros.h" -#include +#include diff --git a/csrc/fps.cpp b/csrc/fps.cpp index c732c21..db39533 100644 --- a/csrc/fps.cpp +++ b/csrc/fps.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 index 0a50a31..b986513 100644 --- a/csrc/graclus.cpp +++ b/csrc/graclus.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 optional_weight) { diff --git a/csrc/grid.cpp b/csrc/grid.cpp index 893d25c..21b6688 100644 --- a/csrc/grid.cpp +++ b/csrc/grid.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 optional_start, diff --git a/csrc/knn.cpp b/csrc/knn.cpp index cfa1e4b..9d0589b 100644 --- a/csrc/knn.cpp +++ b/csrc/knn.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 ptr_x, diff --git a/csrc/nearest.cpp b/csrc/nearest.cpp index 60c4add..83eb41e 100644 --- a/csrc/nearest.cpp +++ b/csrc/nearest.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 index 46d475f..b79bcc1 100644 --- a/csrc/radius.cpp +++ b/csrc/radius.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 ptr_x, diff --git a/csrc/rw.cpp b/csrc/rw.cpp index f45a29c..0f8de62 100644 --- a/csrc/rw.cpp +++ b/csrc/rw.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 random_walk(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start, diff --git a/csrc/sampler.cpp b/csrc/sampler.cpp index 9d1243a..d0b3be8 100644 --- a/csrc/sampler.cpp +++ b/csrc/sampler.cpp @@ -1,15 +1,19 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 index dc45a5f..ad9bbfe 100644 --- a/csrc/version.cpp +++ b/csrc/version.cpp @@ -1,4 +1,6 @@ +#ifdef WITH_PYTHON #include +#endif #include #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 index 8832f8e..c495759 100644 --- 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)]