nonlinear_factor_set_gpu.hpp

// SPDX-License-Identifier: MIT
// Copyright (c) 2021  Kenji Koide (k.koide@aist.go.jp)
 
#pragma once
 
#include <memory>
#include <vector>
 
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam_points/factors/nonlinear_factor_gpu.hpp>
#include <gtsam_points/optimizers/linearization_hook.hpp>
 
struct CUstream_st;
 
namespace gtsam_points {
 
class NonlinearFactorSetGPU : public NonlinearFactorSet {
public:
  NonlinearFactorSetGPU();
  ~NonlinearFactorSetGPU();
 
  int size() const override { return factors.size(); }
 
  void clear() override { factors.clear(); }
 
  void clear_counts() override;
 
  int linearization_count() const override { return num_linearizations; }
 
  int evaluation_count() const override { return num_evaluations; }
 
  bool add(gtsam::NonlinearFactor::shared_ptr factor) override;
 
  void add(const gtsam::NonlinearFactorGraph& factors) override;
 
  void linearize(const gtsam::Values& linearization_point) override;
 
  void error(const gtsam::Values& values) override;
 
  std::vector<gtsam::GaussianFactor::shared_ptr> calc_linear_factors(const gtsam::Values& linearization_point) override;
 
private:
  struct DeviceBuffer {
    DeviceBuffer();
    ~DeviceBuffer();
 
    DeviceBuffer(const DeviceBuffer&) = delete;
    DeviceBuffer& operator=(const DeviceBuffer&) = delete;
 
    void resize(size_t size, CUstream_st* stream);
    unsigned char* data() { return buffer; }
    const unsigned char* data() const { return buffer; }
 
    size_t size;
    unsigned char* buffer;
  };
 
private:
  CUstream_st* stream;
 
  int num_linearizations;
  int num_evaluations;
 
  std::vector<NonlinearFactorGPU::shared_ptr> factors;
 
  std::vector<unsigned char, Eigen::aligned_allocator<unsigned char>> linearization_input_buffer_cpu;
  std::vector<unsigned char, Eigen::aligned_allocator<unsigned char>> linearization_output_buffer_cpu;
  std::unique_ptr<DeviceBuffer> linearization_input_buffer_gpu;
  std::unique_ptr<DeviceBuffer> linearization_output_buffer_gpu;
 
  std::vector<unsigned char, Eigen::aligned_allocator<unsigned char>> evaluation_input_buffer_cpu;
  std::vector<unsigned char, Eigen::aligned_allocator<unsigned char>> evaluation_output_buffer_cpu;
  std::unique_ptr<DeviceBuffer> evaluation_input_buffer_gpu;
  std::unique_ptr<DeviceBuffer> evaluation_output_buffer_gpu;
};
 
}  // namespace gtsam_points