Merge branch 'ldg' of http://10.75.75.11:7001/wjie98/starrygl

.gitignore

@@ -26,6 +26,10 @@ share/python-wheels/
 *.egg
 MANIFEST
 
+# IDE temporary files (generated by IDEs like CLion, etc.)
+.idea/
+cmake-build-*/
+
 # PyInstaller
 #  Usually these files are written by a python script from a template
 #  before PyInstaller builds the exe, so as to inject date/other infos into it.

.gitmodules

+[submodule "csrc/partition/neighbor_clustering"]
+	path = csrc/partition/neighbor_clustering
+	url = https://gitee.com/onlynagesha/graph-partition-v4

CMakeLists.txt

@@ -6,6 +6,7 @@ option(WITH_PYTHON "Link to Python when building" ON)
 option(WITH_CUDA "Link to CUDA when building" ON)
 option(WITH_METIS "Link to METIS when building" ON)
 option(WITH_MTMETIS "Link to multi-threaded METIS when building" ON)
+option(WITH_LDG "Link to (multi-threaded optionally) LDG when building" ON)
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
@@ -54,7 +55,8 @@ if(WITH_METIS)
 
     include_directories(${METIS_INCLUDE_DIRS})
 
-    add_library(metis SHARED "csrc/partition/metis.cpp")
+    add_library(metis SHARED "csrc/partition/metis.cpp"
+            csrc/partition/neighbor_clustering.cpp)
     target_link_libraries(metis PRIVATE ${TORCH_LIBRARIES})
     target_link_libraries(metis PRIVATE ${GKLIB_LIBRARIES})
     target_link_libraries(metis PRIVATE ${METIS_LIBRARIES})
@@ -78,6 +80,11 @@ if(WITH_MTMETIS)
     target_compile_definitions(mtmetis PRIVATE -DMTMETIS_64BIT_PARTITIONS)
 endif()
 
+if (WITH_LDG)
+    # Imports neighbor-clustering based (e.g. LDG algorithm) graph partitioning implementation
+    add_subdirectory(./csrc/partition/neighbor_clustering)
+endif ()
+
 include_directories("csrc/include")
 add_library(${PROJECT_NAME} SHARED csrc/export.cpp)
 
@@ -94,13 +101,19 @@ if (WITH_CUDA)
 endif()
 
 if (WITH_METIS)
+    message(STATUS "Current project '${PROJECT_NAME}' uses METIS graph partitioning algorithm.")
     target_link_libraries(${PROJECT_NAME} PRIVATE metis)
 endif()
 
 if (WITH_MTMETIS)
+    message(STATUS "Current project '${PROJECT_NAME}' uses multi-threaded METIS graph partitioning algorithm.")
     target_link_libraries(${PROJECT_NAME} PRIVATE mtmetis)
 endif()
 
+if (WITH_LDG)
+    message(STATUS "Current project '${PROJECT_NAME}' uses LDG graph partitioning algorithm.")
+    target_link_libraries(${PROJECT_NAME} PRIVATE ldg-vertex-partition)
+endif()
 
 # add libsampler.so
 set(SAMLPER_NAME "${PROJECT_NAME}_sampler")

csrc/export.cpp

@@ -12,6 +12,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 
     m.def("metis_partition", &metis_partition, "metis graph partition");
     m.def("mt_metis_partition", &mt_metis_partition, "multi-threaded metis graph partition");
+    // Note: the switch WITH_MULTITHREADING=ON shall be triggered during compilation
+    // to enable multi-threading functionality.
+    m.def("ldg_partition", &ldg_partition, "(multi-threaded optionally) LDG graph partition");
 
     py::enum_<cudaMemoryAdvise>(m, "cudaMemoryAdvise")
         .value("cudaMemAdviseSetAccessedBy", cudaMemoryAdvise::cudaMemAdviseSetAccessedBy)

csrc/include/partition.h

@@ -21,4 +21,12 @@ at::Tensor mt_metis_partition(
     int64_t num_parts,
     int64_t num_workers,
     bool recursive
-);
\ No newline at end of file
+);
+
+at::Tensor ldg_partition(
+    at::Tensor edges,
+    at::optional<at::Tensor> vertex_weights,
+    at::optional<at::Tensor> initial_partition,
+    int64_t num_parts,
+    int64_t num_workers
+);

neighbor_clustering @ 5e8dd2eb

+Subproject commit 5e8dd2ebda4b0e3b3dec8f43c8ced67785646eef

csrc/partition/neighbor_clustering.cpp

+#include "neighbor_clustering/vertex_partition/vertex_partition.h"
+#include "neighbor_clustering/vertex_partition/params.h"
+
+at::Tensor ldg_partition(at::Tensor edges,
+                         at::optional<at::Tensor> vertex_weights,
+                         at::optional<at::Tensor> initial_partition,
+                         int64_t num_parts,
+                         int64_t num_workers) {
+    AT_ASSERT(edges.dim() == 2);
+    auto edges_n_cols = edges.size(1);
+    AT_ASSERT(edges_n_cols >= 2 && edges_n_cols <= 3);
+    // Note: other checks are performed in the implementation of vp::ldg_partition function series.
+
+    auto n = edges.slice(1, 0, 2).max().item<int64_t>() + 1;
+    auto params = vp::LDGParams{.N = n, .K = num_parts, .openmp_n_threads = static_cast<int>(num_workers)};
+
+    auto edges_clone = edges.clone();
+    if (vertex_weights.has_value()) {
+        auto vertex_weights_clone = vertex_weights->clone();
+        if (initial_partition.has_value()) {
+            auto initial_partition_clone = initial_partition->clone();
+            return vp::ldg_partition_v_init(edges_clone, vertex_weights_clone, initial_partition_clone, params);
+        } else {
+            return vp::ldg_partition_v(edges_clone, vertex_weights_clone, params);
+        }
+    } else {
+        if (initial_partition.has_value()) {
+            auto initial_partition_clone = initial_partition->clone();
+            return vp::ldg_partition_init(edges_clone, initial_partition_clone, params);
+        } else {
+            return vp::ldg_partition(edges_clone, params);
+        }
+    }
+}