reduce memory footprint in graph creation

rapidsai · Aug 30, 2024 · be504cc · be504cc
1 parent d5154d5
commit be504cc
Show file tree

Hide file tree

Showing 2 changed files with 240 additions and 88 deletions.
diff --git a/cpp/include/cugraph/partition_manager.hpp b/cpp/include/cugraph/partition_manager.hpp
@@ -42,7 +42,8 @@ class partition_manager {
   // partitioning along the major axis (major sub-communicator is responsible for this) and along
   // the minor axis (minor sub-communicator is responsible for this). This variable controls whether
   // to map the major sub-communicator to the GPU row communicator or the GPU column communicator.
-  static constexpr bool map_major_comm_to_gpu_row_comm = true;
+  static constexpr bool map_major_comm_to_gpu_row_comm =
+    false;  // FIXME: this is for benchmarking, reset to true before merging
 
 #ifdef __CUDACC__
   __host__ __device__
@@ -71,6 +72,30 @@ class partition_manager {
                                           : (major_comm_rank * minor_comm_size + minor_comm_rank);
   }
 
+#ifdef __CUDACC__
+  __host__ __device__
+#endif
+    static int
+    compute_major_comm_rank_from_global_comm_rank(int major_comm_size,
+                                                  int minor_comm_size,
+                                                  int comm_rank)
+  {
+    return map_major_comm_to_gpu_row_comm ? comm_rank % major_comm_size
+                                          : comm_rank / minor_comm_size;
+  }
+
+#ifdef __CUDACC__
+  __host__ __device__
+#endif
+    static int
+    compute_minor_comm_rank_from_global_comm_rank(int major_comm_size,
+                                                  int minor_comm_size,
+                                                  int comm_rank)
+  {
+    return map_major_comm_to_gpu_row_comm ? comm_rank / major_comm_size
+                                          : comm_rank % minor_comm_size;
+  }
+
 #ifdef __CUDACC__
   __host__ __device__
 #endif