tbb for normal/covariance estimation

koide3 · Sep 22, 2024 · 06ff50b · 06ff50b
1 parent 05da20f
commit 06ff50b
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 5 deletions.
diff --git a/src/gtsam_points/util/covariance_estimation.cpp b/src/gtsam_points/util/covariance_estimation.cpp
@@ -6,23 +6,27 @@
 #include <Eigen/Eigen>
 #include <Eigen/Geometry>
 #include <gtsam_points/ann/kdtree.hpp>
+#include <gtsam_points/util/parallelism.hpp>
+
+#ifdef GTSAM_POINTS_TBB
+#include <tbb/parallel_for.h>
+#endif
 
 namespace gtsam_points {
 
 std::vector<Eigen::Matrix4d> estimate_covariances(const Eigen::Vector4d* points, int num_points, const CovarianceEstimationParams& params) {
   KdTree tree(points, num_points, params.num_threads);
   std::vector<Eigen::Matrix4d> covs(num_points);
 
-#pragma omp parallel for num_threads(params.num_threads) schedule(guided, 8)
-  for (int i = 0; i < num_points; i++) {
+  const auto perpoint_task = [&](int i) {
     std::vector<size_t> k_indices(params.k_neighbors);
     std::vector<double> k_sq_dists(params.k_neighbors);
     size_t num_found = tree.knn_search(points[i].data(), params.k_neighbors, &k_indices[0], &k_sq_dists[0]);
 
     if (num_found < params.k_neighbors) {
       std::cerr << "warning: fewer than k neighbors found for point " << i << std::endl;
       covs[i].setIdentity();
-      continue;
+      return;
     }
 
     Eigen::Vector4d sum_points = Eigen::Vector4d::Zero();
@@ -48,6 +52,24 @@ std::vector<Eigen::Matrix4d> estimate_covariances(const Eigen::Vector4d* points,
         covs[i].block<3, 3>(0, 0) = eig.eigenvectors() * params.eigen_values.asDiagonal() * eig.eigenvectors().inverse();
       } break;
     }
+  };
+
+  if (is_omp_default()) {
+#pragma omp parallel for num_threads(params.num_threads) schedule(guided, 8)
+    for (int i = 0; i < num_points; i++) {
+      perpoint_task(i);
+    }
+  } else {
+#ifdef GTSAM_POINTS_TBB
+    tbb::parallel_for(tbb::blocked_range<int>(0, num_points, 8), [&](const tbb::blocked_range<int>& range) {
+      for (int i = range.begin(); i < range.end(); i++) {
+        perpoint_task(i);
+      }
+    });
+#else
+    std::cerr << "error: TBB is not available" << std::endl;
+    abort();
+#endif
   }
 
   return covs;

diff --git a/src/gtsam_points/util/normal_estimation.cpp b/src/gtsam_points/util/normal_estimation.cpp
@@ -5,21 +5,43 @@
 
 #include <Eigen/Eigen>
 #include <gtsam_points/util/covariance_estimation.hpp>
+#include <gtsam_points/util/parallelism.hpp>
+
+#ifdef GTSAM_POINTS_TBB
+#include <tbb/parallel_for.h>
+#endif
 
 namespace gtsam_points {
 
 std::vector<Eigen::Vector4d> estimate_normals(const Eigen::Vector4d* points, const Eigen::Matrix4d* covs, int num_points, int num_threads) {
   std::vector<Eigen::Vector4d> normals(num_points, Eigen::Vector4d::Zero());
 
-#pragma omp parallel for num_threads(num_threads)
-  for (int i = 0; i < num_points; i++) {
+  const auto perpoint_task = [&](int i) {
     Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> eig;
     eig.computeDirect(covs[i].block<3, 3>(0, 0));
     normals[i].head<3>() = eig.eigenvectors().col(0);
 
     if (points[i].dot(normals[i]) > 1.0) {
       normals[i] = -normals[i];
     }
+  };
+
+  if (is_omp_default()) {
+#pragma omp parallel for num_threads(num_threads)
+    for (int i = 0; i < num_points; i++) {
+      perpoint_task(i);
+    }
+  } else {
+#ifdef GTSAM_POINTS_TBB
+    tbb::parallel_for(tbb::blocked_range<int>(0, num_points, 64), [&](const tbb::blocked_range<int>& range) {
+      for (int i = range.begin(); i < range.end(); i++) {
+        perpoint_task(i);
+      }
+    });
+#else
+    std::cerr << "error : TBB is not enabled" << std::endl;
+    abort();
+#endif
   }
 
   return normals;