Add transform-reduce extension points for CUDA VGICP derivatives

CMakeLists.txt

@@ -280,7 +280,7 @@ if(BUILD_WITH_CUDA)
   )
   set_target_properties(gtsam_points_cuda PROPERTIES
     VERSION   ${PROJECT_VERSION}
-    SOVERSION 1
+    SOVERSION 2
   )
 
   target_link_libraries(gtsam_points
@@ -351,6 +351,14 @@ if(BUILD_TESTS)
     gtest_discover_tests(${test_name} WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}")
   endforeach()
 
+  if(BUILD_WITH_CUDA)
+    set(test_name test_vgicp_derivatives_transform_reduce)
+    add_executable(${test_name} src/test/test_vgicp_derivatives_transform_reduce.cu)
+    target_link_libraries(${test_name} gtsam_points gtest_main)
+    target_include_directories(${test_name} PRIVATE ${Boost_INCLUDE_DIRS} src/test/include)
+    gtest_discover_tests(${test_name} WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}")
+  endif()
+
   if(BUILD_TESTS_PCL)
     enable_language(C)
     find_package(PCL REQUIRED)

include/gtsam_points/factors/impl/integrated_vgicp_derivatives_impl.cuh

@@ -0,0 +1,131 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2021  Kenji Koide (k.koide@aist.go.jp)
+
+#pragma once
+
+#include <thrust/functional.h>
+#include <thrust/iterator/transform_iterator.h>
+
+#include <cub/device/device_reduce.cuh>
+
+#include <gtsam_points/cuda/kernels/lookup_voxels.cuh>
+#include <gtsam_points/cuda/kernels/vgicp_derivatives.cuh>
+#include <gtsam_points/cuda/stream_temp_buffer_roundrobin.hpp>
+
+namespace gtsam_points {
+namespace detail {
+
+template <typename Result, typename Transform>
+struct transform_vgicp_linearization {
+  transform_vgicp_linearization(vgicp_derivatives_kernel derivatives, Transform transform, const Result& identity)
+  : derivatives(derivatives), transform(transform), identity(identity) {}
+
+  __device__ Result operator()(const thrust::pair<int, int>& correspondence) const {
+    if (correspondence.first < 0 || correspondence.second < 0) {
+      return identity;
+    }
+
+    return transform(correspondence, derivatives(correspondence));
+  }
+
+  vgicp_derivatives_kernel derivatives;
+  Transform transform;
+  Result identity;
+};
+
+template <typename Result, typename Transform>
+struct transform_vgicp_error {
+  transform_vgicp_error(vgicp_error_kernel error, Transform transform, const Result& identity)
+  : error(error), transform(transform), identity(identity) {}
+
+  __device__ Result operator()(const thrust::pair<int, int>& correspondence) const {
+    if (correspondence.first < 0 || correspondence.second < 0) {
+      return identity;
+    }
+
+    return transform(correspondence, error(correspondence));
+  }
+
+  vgicp_error_kernel error;
+  Transform transform;
+  Result identity;
+};
+
+}  // namespace detail
+
+template <typename Result, typename Transform, typename Reduction>
+void IntegratedVGICPDerivatives::issue_linearize_transform_reduce(
+  const Eigen::Isometry3f* d_x,
+  Result* d_output,
+  Transform transform,
+  Reduction reduction,
+  const Result& identity) {
+  //
+  lookup_voxels_kernel correspondence_lookup(enable_surface_validation, *target, source->points_gpu, source->normals_gpu, d_x);
+  auto correspondences = thrust::make_transform_iterator(source_inliers, correspondence_lookup);
+
+  vgicp_derivatives_kernel derivatives(d_x, *target, source->points_gpu, source->covs_gpu);
+  detail::transform_vgicp_linearization<Result, Transform> transformed_derivatives(derivatives, transform, identity);
+  auto first = thrust::make_transform_iterator(correspondences, transformed_derivatives);
+
+  void* temp_storage = nullptr;
+  size_t temp_storage_bytes = 0;
+
+  cub::DeviceReduce::Reduce(temp_storage, temp_storage_bytes, first, d_output, num_inliers, reduction, identity, stream);
+
+  temp_storage = temp_buffer->get_buffer(temp_storage_bytes);
+  cub::DeviceReduce::Reduce(temp_storage, temp_storage_bytes, first, d_output, num_inliers, reduction, identity, stream);
+}
+
+template <typename Result, typename Transform>
+void IntegratedVGICPDerivatives::issue_linearize_transform_reduce(
+  const Eigen::Isometry3f* d_x,
+  Result* d_output,
+  Transform transform,
+  const Result& identity) {
+  //
+  issue_linearize_transform_reduce(d_x, d_output, transform, thrust::plus<Result>(), identity);
+}
+
+template <typename Result, typename Transform, typename Reduction>
+void IntegratedVGICPDerivatives::issue_compute_error_transform_reduce(
+  const Eigen::Isometry3f* d_xl,
+  const Eigen::Isometry3f* d_xe,
+  Result* d_output,
+  Transform transform,
+  Reduction reduction,
+  const Result& identity) {
+  //
+  lookup_voxels_kernel correspondence_lookup(
+    enable_surface_validation,
+    *target,
+    source->points_gpu,
+    source->normals_gpu,
+    d_xl);
+  auto correspondences = thrust::make_transform_iterator(source_inliers, correspondence_lookup);
+
+  vgicp_error_kernel error(d_xl, d_xe, *target, source->points_gpu, source->covs_gpu);
+  detail::transform_vgicp_error<Result, Transform> transformed_error(error, transform, identity);
+  auto first = thrust::make_transform_iterator(correspondences, transformed_error);
+
+  void* temp_storage = nullptr;
+  size_t temp_storage_bytes = 0;
+
+  cub::DeviceReduce::Reduce(temp_storage, temp_storage_bytes, first, d_output, num_inliers, reduction, identity, stream);
+
+  temp_storage = temp_buffer->get_buffer(temp_storage_bytes);
+  cub::DeviceReduce::Reduce(temp_storage, temp_storage_bytes, first, d_output, num_inliers, reduction, identity, stream);
+}
+
+template <typename Result, typename Transform>
+void IntegratedVGICPDerivatives::issue_compute_error_transform_reduce(
+  const Eigen::Isometry3f* d_xl,
+  const Eigen::Isometry3f* d_xe,
+  Result* d_output,
+  Transform transform,
+  const Result& identity) {
+  //
+  issue_compute_error_transform_reduce(d_xl, d_xe, d_output, transform, thrust::plus<Result>(), identity);
+}
+
+}  // namespace gtsam_points

include/gtsam_points/factors/integrated_vgicp_derivatives.cuh

@@ -25,7 +25,7 @@ public:
     const PointCloud::ConstPtr& source,
     CUstream_st* ext_stream,
     std::shared_ptr<TempBufferManager> temp_buffer);
-  ~IntegratedVGICPDerivatives();
+  virtual ~IntegratedVGICPDerivatives();
 
   void set_inlier_update_thresh(double trans, double angle) {
     inlier_update_thresh_trans = trans;
@@ -49,8 +49,71 @@ public:
 
   // async interface
   void sync_stream();
-  void issue_linearize(const Eigen::Isometry3f* d_x, LinearizedSystem6* d_output);
-  void issue_compute_error(const Eigen::Isometry3f* d_xl, const Eigen::Isometry3f* d_xe, float* d_output);
+  virtual void issue_linearize(const Eigen::Isometry3f* d_x, LinearizedSystem6* d_output);
+  virtual void issue_compute_error(const Eigen::Isometry3f* d_xl, const Eigen::Isometry3f* d_xe, float* d_output);
+
+protected:
+  /// @brief Get the CUDA stream used by this derivatives instance.
+  CUstream_st* cuda_stream() const { return stream; }
+
+  /**
+   * @brief Apply a custom device transform to each valid VGICP correspondence after standard linearization and reduce the results.
+   *
+   * Transform must be device-callable with the following signature:
+   * @code
+   * Result operator()(const thrust::pair<int, int>& correspondence, const LinearizedSystem6& linearized) const;
+   * @endcode
+   * Invalid correspondences are mapped directly to identity and are never passed to Transform.
+   * Reduction must be device-callable and identity must be its neutral element.
+   * Transform must preserve LinearizedSystem6::num_inliers when its result contains the system used by
+   * IntegratedVGICPFactorGPU, because the count is also used to maintain the inlier index buffer.
+   */
+  template <typename Result, typename Transform, typename Reduction>
+  void issue_linearize_transform_reduce(
+    const Eigen::Isometry3f* d_x,
+    Result* d_output,
+    Transform transform,
+    Reduction reduction,
+    const Result& identity);
+
+  /**
+   * @brief Additive convenience overload of issue_linearize_transform_reduce().
+   */
+  template <typename Result, typename Transform>
+  void issue_linearize_transform_reduce(
+    const Eigen::Isometry3f* d_x,
+    Result* d_output,
+    Transform transform,
+    const Result& identity);
+
+  /**
+   * @brief Apply a custom device transform to each valid VGICP correspondence error and reduce the results.
+   *
+   * Transform must be device-callable with the following signature:
+   * @code
+   * Result operator()(const thrust::pair<int, int>& correspondence, float error) const;
+   * @endcode
+   * Invalid correspondences are mapped directly to identity and are never passed to Transform.
+   */
+  template <typename Result, typename Transform, typename Reduction>
+  void issue_compute_error_transform_reduce(
+    const Eigen::Isometry3f* d_xl,
+    const Eigen::Isometry3f* d_xe,
+    Result* d_output,
+    Transform transform,
+    Reduction reduction,
+    const Result& identity);
+
+  /**
+   * @brief Additive convenience overload of issue_compute_error_transform_reduce().
+   */
+  template <typename Result, typename Transform>
+  void issue_compute_error_transform_reduce(
+    const Eigen::Isometry3f* d_xl,
+    const Eigen::Isometry3f* d_xe,
+    Result* d_output,
+    Transform transform,
+    const Result& identity);
 
 private:
   bool enable_offloading;
@@ -73,4 +136,8 @@ private:
   int* num_inliers_gpu;
   int* source_inliers;
 };
-}  // namespace gtsam_points
+}  // namespace gtsam_points
+
+#if defined(__CUDACC__)
+#include <gtsam_points/factors/impl/integrated_vgicp_derivatives_impl.cuh>
+#endif

include/gtsam_points/factors/integrated_vgicp_factor_gpu.hpp

@@ -135,6 +135,16 @@ class IntegratedVGICPFactorGPU : public gtsam_points::NonlinearFactorGPU {
 
   virtual void sync() override;
 
+protected:
+  /**
+   * @brief Replace the derivatives implementation used by this factor.
+   * @note This is intended for derived factors that install a custom implementation during construction.
+   * @warning A derived factor that calls this function must override clone() and reinstall an equivalent
+   *          derivatives implementation in the cloned factor. The inherited clone() rejects custom derivatives
+   *          to prevent silently reverting to the standard implementation.
+   */
+  void replace_derivatives(std::unique_ptr<IntegratedVGICPDerivatives> replacement);
+
 private:
   Eigen::Isometry3f calc_delta(const gtsam::Values& values) const;
 
@@ -154,4 +164,4 @@ class IntegratedVGICPFactorGPU : public gtsam_points::NonlinearFactorGPU {
   mutable std::unique_ptr<LinearizedSystem6> linearization_result;
 };
 
-}  // namespace gtsam_points
+}  // namespace gtsam_points

src/gtsam_points/factors/integrated_vgicp_derivatives_compute.cu

@@ -3,39 +3,20 @@
 
 #include <gtsam_points/factors/integrated_vgicp_derivatives.cuh>
 
-#include <iostream>
-#include <thrust/remove.h>
-#include <thrust/iterator/transform_iterator.h>
-
-#include <cub/device/device_reduce.cuh>
-
-#include <gtsam_points/cuda/kernels/pose.cuh>
-#include <gtsam_points/cuda/kernels/untie.cuh>
-#include <gtsam_points/cuda/kernels/lookup_voxels.cuh>
-#include <gtsam_points/cuda/kernels/linearized_system.cuh>
-#include <gtsam_points/cuda/kernels/vgicp_derivatives.cuh>
-#include <gtsam_points/cuda/stream_temp_buffer_roundrobin.hpp>
-
-#include <gtsam_points/types/gaussian_voxelmap_gpu.hpp>
+#include <thrust/pair.h>
 
 namespace gtsam_points {
 
-void IntegratedVGICPDerivatives::issue_compute_error(const Eigen::Isometry3f* d_xl, const Eigen::Isometry3f* d_xe, float* d_output) {
-  //
-  lookup_voxels_kernel corr_kernel(enable_surface_validation, *target, source->points_gpu, source->normals_gpu, d_xl);
-  auto corr_first = thrust::make_transform_iterator(source_inliers, corr_kernel);
+namespace {
 
-  vgicp_error_kernel error_kernel(d_xl, d_xe, *target, source->points_gpu, source->covs_gpu);
-  auto first = thrust::make_transform_iterator(corr_first, error_kernel);
+struct identity_error_transform {
+  __device__ float operator()(const thrust::pair<int, int>&, float error) const { return error; }
+};
 
-  void* temp_storage = nullptr;
-  size_t temp_storage_bytes = 0;
+}  // namespace
 
-  cub::DeviceReduce::Reduce(temp_storage, temp_storage_bytes, first, d_output, num_inliers, thrust::plus<float>(), 0.0f, stream);
-
-  temp_storage = temp_buffer->get_buffer(temp_storage_bytes);
-
-  cub::DeviceReduce::Reduce(temp_storage, temp_storage_bytes, first, d_output, num_inliers, thrust::plus<float>(), 0.0f, stream);
+void IntegratedVGICPDerivatives::issue_compute_error(const Eigen::Isometry3f* d_xl, const Eigen::Isometry3f* d_xe, float* d_output) {
+  issue_compute_error_transform_reduce(d_xl, d_xe, d_output, identity_error_transform(), 0.0f);
 }
 
-}  // namespace gtsam_points
+}  // namespace gtsam_points

src/gtsam_points/factors/integrated_vgicp_derivatives_linearize.cu

@@ -3,55 +3,25 @@
 
 #include <gtsam_points/factors/integrated_vgicp_derivatives.cuh>
 
-#include <iostream>
-#include <thrust/remove.h>
-#include <thrust/iterator/transform_iterator.h>
+#include <gtsam_points/cuda/kernels/linearized_system.cuh>
 
-#include <cub/device/device_reduce.cuh>
+namespace gtsam_points {
 
-#include <gtsam_points/cuda/kernels/pose.cuh>
-#include <gtsam_points/cuda/kernels/untie.cuh>
-#include <gtsam_points/cuda/kernels/lookup_voxels.cuh>
-#include <gtsam_points/cuda/kernels/linearized_system.cuh>
-#include <gtsam_points/cuda/kernels/vgicp_derivatives.cuh>
-#include <gtsam_points/cuda/stream_temp_buffer_roundrobin.hpp>
+namespace {
 
-#include <gtsam_points/types/gaussian_voxelmap_gpu.hpp>
+struct identity_linearization_transform {
+  __device__ LinearizedSystem6 operator()(
+    const thrust::pair<int, int>&,
+    const LinearizedSystem6& linearized) const {
+    //
+    return linearized;
+  }
+};
 
-namespace gtsam_points {
+}  // namespace
 
 void IntegratedVGICPDerivatives::issue_linearize(const Eigen::Isometry3f* d_x, LinearizedSystem6* d_output) {
-  //
-  lookup_voxels_kernel corr_kernel(enable_surface_validation, *target, source->points_gpu, source->normals_gpu, d_x);
-  auto corr_first = thrust::make_transform_iterator(source_inliers, corr_kernel);
-
-  vgicp_derivatives_kernel deriv_kernel(d_x, *target, source->points_gpu, source->covs_gpu);
-  auto first = thrust::make_transform_iterator(corr_first, deriv_kernel);
-
-  void* temp_storage = nullptr;
-  size_t temp_storage_bytes = 0;
-
-  cub::DeviceReduce::Reduce(
-    temp_storage,
-    temp_storage_bytes,
-    first,
-    d_output,
-    num_inliers,
-    thrust::plus<LinearizedSystem6>(),
-    LinearizedSystem6::zero(),
-    stream);
-
-  temp_storage = temp_buffer->get_buffer(temp_storage_bytes);
-
-  cub::DeviceReduce::Reduce(
-    temp_storage,
-    temp_storage_bytes,
-    first,
-    d_output,
-    num_inliers,
-    thrust::plus<LinearizedSystem6>(),
-    LinearizedSystem6::zero(),
-    stream);
+  issue_linearize_transform_reduce(d_x, d_output, identity_linearization_transform(), LinearizedSystem6::zero());
 }
 
-}  // namespace gtsam_points
+}  // namespace gtsam_points