Offloading points and voxelmaps on GPU

config/config_global_mapping_gpu.json

@@ -22,6 +22,13 @@
   // use_isam2_dogleg               : If true, use dogleg optimizer (robust but slow)
   // isam2_relinearize_skip         : Relinearization is performed every $isam2_relinearize_skip optimization calls
   // isam2_relinearize_thresh       : Relinearization is performed only when linear delta gets larger than this
+  //
+  // --- Memory settings ---
+  // gpu_memory_offload_mb          : Threshold for GPU memory offloading in MB (0 = disable offloading).
+  //                                : If GPU memory usage exceeds this value, old submap points and voxelmaps are offloaded to CPU memory.
+  //                                : Note that the current implementation doesn't support offloading of factors and graphics data, and thus
+  //                                : the actual GPU memory usage will be much larger than this value.
+  //                                : (Setting this to a half or quarter of the total GPU memory would be a good choice).
   */
   "global_mapping": {
     "so_name": "libglobal_mapping.so",
@@ -46,6 +53,8 @@
     // Optimizer settings
     "use_isam2_dogleg": false,
     "isam2_relinearize_skip": 1,
-    "isam2_relinearize_thresh": 0.1
+    "isam2_relinearize_thresh": 0.1,
+    // Memory settings
+    "gpu_memory_offload_mb": 2048
   }
 }

include/glim/mapping/global_mapping.hpp

@@ -12,8 +12,9 @@ class NonlinearFactorGraph;
 
 namespace gtsam_points {
 class ISAM2Ext;
-class StreamTempBufferRoundRobin;
 struct ISAM2ResultExt;
+class StreamTempBufferRoundRobin;
+class OffloadableGPU;
 }  // namespace gtsam_points
 
 namespace glim {
@@ -35,6 +36,8 @@ struct GlobalMappingParams {
   bool enable_between_factors;
   std::string between_registration_type;
 
+  size_t gpu_memory_offload_mb;  // in MB
+
   std::string registration_error_factor_type;
   double submap_voxel_resolution;
   double submap_voxel_resolution_max;
@@ -84,6 +87,7 @@ class GlobalMapping : public GlobalMappingBase {
   std::shared_ptr<gtsam::NonlinearFactorGraph> create_matching_cost_factors(int current) const;
 
   void update_submaps();
+  void offload_gpu_memory();
   gtsam_points::ISAM2ResultExt update_isam2(const gtsam::NonlinearFactorGraph& new_factors, const gtsam::Values& new_values);
 
   void recover_graph() override;
@@ -96,7 +100,11 @@ class GlobalMapping : public GlobalMappingBase {
   std::mt19937 mt;
   int session_id;
 
+  size_t bytes_gpu;
+  std::vector<std::shared_ptr<gtsam_points::OffloadableGPU>> offloadables;
+
   std::unique_ptr<IMUIntegration> imu_integration;
+  std::any main_stream;
   std::any stream_buffer_roundrobin;
 
   std::vector<SubMap::Ptr> submaps;

include/glim/viewer/standard_viewer.hpp

@@ -22,12 +22,16 @@ namespace gtsam {
 class NonlinearFactor;
 }
 
+namespace gtsam_points {
+class OffloadableGPU;
+}
+
 namespace glim {
 
 class TrajectoryManager;
 struct EstimationFrame;
 
-struct SubMapMemoryStats;
+struct SubMap;
 struct FactorMemoryStats;
 
 class StandardViewer : public ExtensionModule {
@@ -89,11 +93,8 @@ class StandardViewer : public ExtensionModule {
   float min_overlap;
 
   bool show_memory_stats;
-  int submap_memstats_count;
-  std::vector<SubMapMemoryStats> submap_memstats;
-
-  int global_factor_stats_count;
-  std::vector<FactorMemoryStats> global_factor_memstats;
+  std::vector<std::shared_ptr<const SubMap>> submaps;
+  std::vector<std::shared_ptr<const gtsam_points::OffloadableGPU>> gpu_factors;
 
   size_t total_gl_bytes;
 

include/glim/viewer/standard_viewer_mem.hpp

@@ -19,9 +19,12 @@ struct SubMapMemoryStats {
   size_t voxelmap_cpu_bytes;  ///< CPU memory usage for voxel map
   size_t frame_gpu_bytes;     ///< GPU memory usage
   size_t voxelmap_gpu_bytes;  ///< GPU memory usage for voxel map
+  size_t frame_gpu_net_bytes;     ///< GPU memory usage for the frame
+  size_t voxelmap_gpu_net_bytes;  ///< GPU memory usage for the voxel map
 
-  size_t odom_cpu_bytes;  ///< Total CPU memory usage for all frames
-  size_t odom_gpu_bytes;  ///< Total GPU memory usage for all frames
+  size_t odom_cpu_bytes;      ///< Total CPU memory usage for all frames
+  size_t odom_gpu_bytes;      ///< Total GPU memory usage for all frames
+  size_t odom_gpu_net_bytes;  ///< Total GPU memory usage for all frames (networked)
 
   size_t num_custom_data;  ///< Number of custom data entries in the submap
 };

src/glim/mapping/global_mapping.cpp

@@ -27,6 +27,7 @@
 #include <gtsam_points/optimizers/isam2_ext.hpp>
 #include <gtsam_points/optimizers/isam2_ext_dummy.hpp>
 #include <gtsam_points/optimizers/levenberg_marquardt_ext.hpp>
+#include <gtsam_points/cuda/cuda_stream.hpp>
 #include <gtsam_points/cuda/stream_temp_buffer_roundrobin.hpp>
 
 #include <glim/util/config.hpp>
@@ -55,6 +56,8 @@ GlobalMappingParams::GlobalMappingParams() {
 
   enable_between_factors = config.param<bool>("global_mapping", "create_between_factors", false);
   between_registration_type = config.param<std::string>("global_mapping", "between_registration_type", "GICP");
+  gpu_memory_offload_mb = config.param<int>("global_mapping", "gpu_memory_offload_mb", 0);
+
   registration_error_factor_type = config.param<std::string>("global_mapping", "registration_error_factor_type", "VGICP");
   submap_voxel_resolution = config.param<double>("global_mapping", "submap_voxel_resolution", 1.0);
   submap_voxel_resolution_max = config.param<double>("global_mapping", "submap_voxel_resolution_max", submap_voxel_resolution);
@@ -87,6 +90,7 @@ GlobalMapping::GlobalMapping(const GlobalMappingParams& params) : params(params)
 #endif
 
   session_id = 0;
+  bytes_gpu = 0;
   imu_integration.reset(new IMUIntegration);
 
   new_values.reset(new gtsam::Values);
@@ -107,7 +111,8 @@ GlobalMapping::GlobalMapping(const GlobalMappingParams& params) : params(params)
   }
 
 #ifdef GTSAM_POINTS_USE_CUDA
-  stream_buffer_roundrobin = std::make_shared<gtsam_points::StreamTempBufferRoundRobin>(64);
+  main_stream = std::make_shared<gtsam_points::CUDAStream>();
+  stream_buffer_roundrobin = std::make_shared<gtsam_points::StreamTempBufferRoundRobin>(4);
 #endif
 
 #ifdef GTSAM_USE_TBB
@@ -229,6 +234,8 @@ void GlobalMapping::insert_submap(const SubMap::Ptr& submap) {
 
   update_submaps();
   Callbacks::on_update_submaps(submaps);
+
+  offload_gpu_memory();
 }
 
 void GlobalMapping::insert_submap(int current, const SubMap::Ptr& submap) {
@@ -265,6 +272,9 @@ void GlobalMapping::insert_submap(int current, const SubMap::Ptr& submap) {
       auto voxelmap = std::make_shared<gtsam_points::GaussianVoxelMapGPU>(resolution);
       voxelmap->insert(*submap->frame);
       submap->voxelmaps.push_back(voxelmap);
+
+      offloadables.emplace_back(voxelmap);
+      bytes_gpu += voxelmap->memory_usage_gpu();
     }
   }
 #endif
@@ -278,6 +288,14 @@ void GlobalMapping::insert_submap(int current, const SubMap::Ptr& submap) {
     }
   }
 
+#ifdef GTSAM_POINTS_USE_CUDA
+  auto points_gpu = std::dynamic_pointer_cast<gtsam_points::PointCloudGPU>(submap->frame);
+  if (points_gpu) {
+    offloadables.emplace_back(points_gpu);
+    bytes_gpu += points_gpu->memory_usage_gpu();
+  }
+#endif
+
   submaps.push_back(submap);
   subsampled_submaps.push_back(subsampled_submap);
 }
@@ -332,7 +350,9 @@ void GlobalMapping::find_overlapping_submaps(double min_overlap) {
         const auto& stream = stream_buffer.first;
         const auto& buffer = stream_buffer.second;
         for (const auto& voxelmap : submaps[i]->voxelmaps) {
-          new_factors->emplace_shared<gtsam_points::IntegratedVGICPFactorGPU>(X(i), X(j), voxelmap, subsampled_submaps[j], stream, buffer);
+          auto f = gtsam::make_shared<gtsam_points::IntegratedVGICPFactorGPU>(X(i), X(j), voxelmap, subsampled_submaps[j], stream, buffer);
+          f->set_enable_offloading(params.gpu_memory_offload_mb > 0);
+          new_factors->add(f);
         }
       }
 #endif
@@ -463,7 +483,9 @@ std::shared_ptr<gtsam::NonlinearFactorGraph> GlobalMapping::create_matching_cost
       const auto& stream = stream_buffer.first;
       const auto& buffer = stream_buffer.second;
       for (const auto& voxelmap : submaps[i]->voxelmaps) {
-        factors->emplace_shared<gtsam_points::IntegratedVGICPFactorGPU>(X(i), X(current), voxelmap, subsampled_submaps[current], stream, buffer);
+        auto f = gtsam::make_shared<gtsam_points::IntegratedVGICPFactorGPU>(X(i), X(current), voxelmap, subsampled_submaps[current], stream, buffer);
+        f->set_enable_offloading(params.gpu_memory_offload_mb > 0);
+        factors->add(f);
       }
     }
 #endif
@@ -489,6 +511,52 @@ void GlobalMapping::update_submaps() {
   }
 }
 
+void GlobalMapping::offload_gpu_memory() {
+#ifdef GTSAM_POINTS_USE_CUDA
+  const size_t thresh_hi = params.gpu_memory_offload_mb * (1024ull * 1024ull);
+  const size_t thresh_lo = thresh_hi * 0.9;
+
+  if (params.gpu_memory_offload_mb == 0 || bytes_gpu < thresh_hi) {
+    return;
+  }
+
+  logger->info("offload GPU memory ({:.2f}MB)", bytes_gpu / (1024.0 * 1024.0));
+
+  size_t new_bytes_gpu = 0;
+  for (const auto& offloadable : offloadables) {
+    if (offloadable->loaded_on_gpu()) {
+      new_bytes_gpu += offloadable->memory_usage_gpu();
+    }
+  }
+  logger->info("total GPU memory usage: {:.2f}MB", new_bytes_gpu / (1024.0 * 1024.0));
+
+  std::sort(offloadables.begin(), offloadables.end(), [](const gtsam_points::OffloadableGPU::Ptr& a, const gtsam_points::OffloadableGPU::Ptr& b) {
+    return a->last_accessed_time() < b->last_accessed_time();
+  });
+
+  auto stream = std::any_cast<std::shared_ptr<gtsam_points::CUDAStream>>(main_stream);
+  for (auto& offloadable : offloadables) {
+    if (!offloadable->loaded_on_gpu()) {
+      continue;
+    }
+
+    if (new_bytes_gpu < thresh_lo) {
+      break;
+    }
+
+    size_t bytes = offloadable->memory_usage_gpu();
+    offloadable->offload_gpu(*stream);
+    new_bytes_gpu -= bytes;
+    logger->debug("offloaded {} bytes", bytes);
+  }
+  stream->sync();
+
+  logger->info("offloaded GPU memory ({:.2f}MB)", new_bytes_gpu / (1024.0 * 1024.0));
+  bytes_gpu = new_bytes_gpu;
+
+#endif
+}
+
 gtsam_points::ISAM2ResultExt GlobalMapping::update_isam2(const gtsam::NonlinearFactorGraph& new_factors, const gtsam::Values& new_values) {
   gtsam_points::ISAM2ResultExt result;
 
@@ -540,6 +608,18 @@ gtsam_points::ISAM2ResultExt GlobalMapping::update_isam2(const gtsam::NonlinearF
     return update_isam2(factors, values);
   }
 
+  for (auto& factor : new_factors) {
+    if (!factor) {
+      continue;
+    }
+
+    auto offloadable = std::dynamic_pointer_cast<gtsam_points::OffloadableGPU>(factor);
+    if (offloadable) {
+      offloadables.emplace_back(offloadable);
+      bytes_gpu += offloadable->memory_usage_gpu();
+    }
+  }
+
   return result;
 }
 
@@ -822,7 +902,8 @@ bool GlobalMapping::load(const std::string& path) {
         const auto& buffer = stream_buffer.second;
 
         for (const auto& voxelmap : submaps[first]->voxelmaps) {
-          graph.emplace_shared<gtsam_points::IntegratedVGICPFactorGPU>(X(first), X(second), voxelmap, subsampled_submaps[second], stream, buffer);
+          auto f = gtsam::make_shared<gtsam_points::IntegratedVGICPFactorGPU>(X(first), X(second), voxelmap, subsampled_submaps[second], stream, buffer);
+          graph.add(f);
         }
 #else
         logger->warn("GPU is enabled but gtsam_points was built without CUDA!!");

src/glim/odometry/callbacks.cpp

@@ -1,6 +1,6 @@
 #include <glim/odometry/callbacks.hpp>
 
-#include <gtsam_unstable/nonlinear/FixedLagSmoother.h>
+#include <gtsam/nonlinear/FixedLagSmoother.h>
 
 namespace glim {