gaussian_voxelmap_gpu.hpp

// SPDX-License-Identifier: MIT
// Copyright (c) 2021  Kenji Koide (k.koide@aist.go.jp)
 
#pragma once
 
#include <memory>
#include <Eigen/Core>
 
#include <gtsam_points/types/offloadable.hpp>
#include <gtsam_points/types/gaussian_voxelmap.hpp>
 
// forward declaration
struct CUstream_st;
 
namespace gtsam_points {
 
struct VoxelMapInfo {
  int num_voxels;             
  int num_buckets;            
  int max_bucket_scan_count;  
  float voxel_resolution;     
};
 
struct VoxelBucket {
  Eigen::Vector3i first;
  int second;
};
 
class GaussianVoxelMapGPU : public GaussianVoxelMap, public OffloadableGPU {
public:
  using Ptr = std::shared_ptr<GaussianVoxelMapGPU>;
  using ConstPtr = std::shared_ptr<const GaussianVoxelMapGPU>;
 
  GaussianVoxelMapGPU(
    float resolution,
    int init_num_buckets = 8192 * 2,
    int max_bucket_scan_count = 10,
    double target_points_drop_rate = 1e-3,
    CUstream_st* stream = 0);
  virtual ~GaussianVoxelMapGPU();
 
  virtual double voxel_resolution() const override { return voxelmap_info.voxel_resolution; }
 
  virtual void insert(const PointCloud& frame) override;
 
  void save_compact(const std::string& path) const override;
 
  static GaussianVoxelMapGPU::Ptr load(const std::string& path);
 
  // GPU memory offloading
  size_t memory_usage_gpu() const override;
  bool loaded_on_gpu() const override;
  bool offload_gpu(CUstream_st* stream = 0) override;
  bool reload_gpu(CUstream_st* stream = 0) override;
 
private:
  void create_bucket_table(CUstream_st* stream, const PointCloud& frame);
 
public:
  CUstream_st* stream;
 
  const int init_num_buckets;            
  const double target_points_drop_rate;  
  VoxelMapInfo voxelmap_info;            
  VoxelMapInfo* voxelmap_info_ptr;       
 
  VoxelBucket* buckets;  
 
  // voxel data
  int* num_points;               
  Eigen::Vector3f* voxel_means;  
  Eigen::Matrix3f* voxel_covs;   
  float* voxel_intensities;      
 
  // GPU memory offloading
  std::vector<VoxelBucket> offloaded_buckets;          
  std::vector<int> offloaded_num_points;               
  std::vector<Eigen::Vector3f> offloaded_voxel_means;  
  std::vector<Eigen::Matrix3f> offloaded_voxel_covs;   
};
 
std::vector<VoxelBucket> download_buckets(const GaussianVoxelMapGPU& voxelmap, CUstream_st* stream = nullptr);
std::vector<int> download_voxel_num_points(const GaussianVoxelMapGPU& voxelmap, CUstream_st* stream = nullptr);
std::vector<Eigen::Vector3f> download_voxel_means(const GaussianVoxelMapGPU& voxelmap, CUstream_st* stream = nullptr);
std::vector<Eigen::Matrix3f> download_voxel_covs(const GaussianVoxelMapGPU& voxelmap, CUstream_st* stream = nullptr);
std::vector<float> download_voxel_intensities(const GaussianVoxelMapGPU& vm, CUstream_st* stream = nullptr);
}  // namespace gtsam_points