Add support for limiting num_threads in tbb task

cpp/kiss_icp/core/Registration.cpp

@@ -22,8 +22,10 @@
 // SOFTWARE.
 #include "Registration.hpp"
 
+#include <oneapi/tbb/task_arena.h>
 #include <tbb/blocked_range.h>
 #include <tbb/parallel_reduce.h>
+#include <tbb/task_arena.h>
 
 #include <algorithm>
 #include <cmath>
@@ -68,7 +70,8 @@ constexpr double ESTIMATION_THRESHOLD_ = 0.0001;
 std::tuple<Eigen::Matrix6d, Eigen::Vector6d> BuildLinearSystem(
     const std::vector<Eigen::Vector3d> &source,
     const std::vector<Eigen::Vector3d> &target,
-    double kernel) {
+    double kernel,
+    int max_threads = tbb::task_arena::automatic) {
     auto compute_jacobian_and_residual = [&](auto i) {
         const Eigen::Vector3d residual = source[i] - target[i];
         Eigen::Matrix3_6d J_r;
@@ -77,29 +80,33 @@ std::tuple<Eigen::Matrix6d, Eigen::Vector6d> BuildLinearSystem(
         return std::make_tuple(J_r, residual);
     };
 
-    const auto &[JTJ, JTr] = tbb::parallel_reduce(
-        // Range
-        tbb::blocked_range<size_t>{0, source.size()},
-        // Identity
-        ResultTuple(),
-        // 1st Lambda: Parallel computation
-        [&](const tbb::blocked_range<size_t> &r, ResultTuple J) -> ResultTuple {
-            auto Weight = [&](double residual2) {
-                return square(kernel) / square(kernel + residual2);
-            };
-            auto &[JTJ_private, JTr_private] = J;
-            for (auto i = r.begin(); i < r.end(); ++i) {
-                const auto &[J_r, residual] = compute_jacobian_and_residual(i);
-                const double w = Weight(residual.squaredNorm());
-                JTJ_private.noalias() += J_r.transpose() * w * J_r;
-                JTr_private.noalias() += J_r.transpose() * w * residual;
-            }
-            return J;
-        },
-        // 2nd Lambda: Parallel reduction of the private Jacboians
-        [&](ResultTuple a, const ResultTuple &b) -> ResultTuple { return a + b; });
-
-    return std::make_tuple(JTJ, JTr);
+    ResultTuple jacobian;
+    tbb::task_arena limited_arena(max_threads);
+    limited_arena.execute([&]() -> void {
+        jacobian = tbb::parallel_reduce(
+            // Range
+            tbb::blocked_range<size_t>{0, source.size()},
+            // Identity
+            ResultTuple(),
+            // 1st Lambda: Parallel computation
+            [&](const tbb::blocked_range<size_t> &r, ResultTuple J) -> ResultTuple {
+                auto Weight = [&](double residual2) {
+                    return square(kernel) / square(kernel + residual2);
+                };
+                auto &[JTJ_private, JTr_private] = J;
+                for (auto i = r.begin(); i < r.end(); ++i) {
+                    const auto &[J_r, residual] = compute_jacobian_and_residual(i);
+                    const double w = Weight(residual.squaredNorm());
+                    JTJ_private.noalias() += J_r.transpose() * w * J_r;
+                    JTr_private.noalias() += J_r.transpose() * w * residual;
+                }
+                return J;
+            },
+            // 2nd Lambda: Parallel reduction of the private Jacboians
+            [&](ResultTuple a, const ResultTuple &b) -> ResultTuple { return a + b; });
+    });
+
+    return std::make_tuple(jacobian.JTJ, jacobian.JTr);
 }
 }  // namespace
 
@@ -122,7 +129,7 @@ Sophus::SE3d RegisterFrame(const std::vector<Eigen::Vector3d> &frame,
         // Equation (10)
         const auto &[src, tgt] = voxel_map.GetCorrespondences(source, max_correspondence_distance);
         // Equation (11)
-        const auto &[JTJ, JTr] = BuildLinearSystem(src, tgt, kernel);
+        const auto &[JTJ, JTr] = BuildLinearSystem(src, tgt, kernel, voxel_map.max_threads_);
         const Eigen::Vector6d dx = JTJ.ldlt().solve(-JTr);
         const Sophus::SE3d estimation = Sophus::SE3d::exp(dx);
         // Equation (12)

cpp/kiss_icp/core/VoxelHashMap.cpp

@@ -24,6 +24,7 @@
 
 #include <tbb/blocked_range.h>
 #include <tbb/parallel_reduce.h>
+#include <tbb/task_arena.h>
 
 #include <Eigen/Core>
 #include <algorithm>
@@ -42,6 +43,7 @@ struct ResultTuple {
     std::vector<Eigen::Vector3d> source;
     std::vector<Eigen::Vector3d> target;
 };
+
 }  // namespace
 
 namespace kiss_icp {
@@ -90,39 +92,46 @@ VoxelHashMap::Vector3dVectorTuple VoxelHashMap::GetCorrespondences(
 
         return closest_neighbor;
     };
+
     using points_iterator = std::vector<Eigen::Vector3d>::const_iterator;
-    const auto [source, target] = tbb::parallel_reduce(
-        // Range
-        tbb::blocked_range<points_iterator>{points.cbegin(), points.cend()},
-        // Identity
-        ResultTuple(points.size()),
-        // 1st lambda: Parallel computation
-        [max_correspondance_distance, &GetClosestNeighboor](
-            const tbb::blocked_range<points_iterator> &r, ResultTuple res) -> ResultTuple {
-            auto &[src, tgt] = res;
-            src.reserve(r.size());
-            tgt.reserve(r.size());
-            for (const auto &point : r) {
-                Eigen::Vector3d closest_neighboors = GetClosestNeighboor(point);
-                if ((closest_neighboors - point).norm() < max_correspondance_distance) {
-                    src.emplace_back(point);
-                    tgt.emplace_back(closest_neighboors);
+    ResultTuple correspondences(points.size());
+    tbb::task_arena limited_arena(max_threads_);
+    limited_arena.execute([&]() -> void {
+        correspondences = tbb::parallel_reduce(
+            // Range
+            tbb::blocked_range<points_iterator>{points.cbegin(), points.cend()},
+            // Identity
+            ResultTuple(points.size()),
+            // 1st lambda: Parallel computation
+            [max_correspondance_distance, &GetClosestNeighboor](
+                const tbb::blocked_range<points_iterator> &r, ResultTuple res) -> ResultTuple {
+                auto &[src, tgt] = res;
+                src.reserve(r.size());
+                tgt.reserve(r.size());
+                for (const auto &point : r) {
+                    Eigen::Vector3d closest_neighboors = GetClosestNeighboor(point);
+                    if ((closest_neighboors - point).norm() < max_correspondance_distance) {
+                        src.emplace_back(point);
+                        tgt.emplace_back(closest_neighboors);
+                    }
                 }
-            }
-            return res;
-        },
-        // 2nd lambda: Parallel reduction
-        [](ResultTuple a, const ResultTuple &b) -> ResultTuple {
-            auto &[src, tgt] = a;
-            const auto &[srcp, tgtp] = b;
-            src.insert(src.end(),  //
-                       std::make_move_iterator(srcp.begin()), std::make_move_iterator(srcp.end()));
-            tgt.insert(tgt.end(),  //
-                       std::make_move_iterator(tgtp.begin()), std::make_move_iterator(tgtp.end()));
-            return a;
-        });
+                return res;
+            },
+            // 2nd lambda: Parallel reduction
+            [](ResultTuple a, const ResultTuple &b) -> ResultTuple {
+                auto &[src, tgt] = a;
+                const auto &[srcp, tgtp] = b;
+                src.insert(src.end(),  //
+                           std::make_move_iterator(srcp.begin()),
+                           std::make_move_iterator(srcp.end()));
+                tgt.insert(tgt.end(),  //
+                           std::make_move_iterator(tgtp.begin()),
+                           std::make_move_iterator(tgtp.end()));
+                return a;
+            });
+    });
 
-    return std::make_tuple(source, target);
+    return std::make_tuple(correspondences.source, correspondences.target);
 }
 
 std::vector<Eigen::Vector3d> VoxelHashMap::Pointcloud() const {

cpp/kiss_icp/core/VoxelHashMap.hpp

@@ -26,6 +26,7 @@
 // should at least acknoowledge the work from CT-ICP by giving a star on GitHub
 #pragma once
 
+#include <oneapi/tbb/task_arena.h>
 #include <tsl/robin_map.h>
 
 #include <Eigen/Core>
@@ -52,10 +53,14 @@ struct VoxelHashMap {
         }
     };
 
-    explicit VoxelHashMap(double voxel_size, double max_distance, int max_points_per_voxel)
+    explicit VoxelHashMap(double voxel_size,
+                          double max_distance,
+                          int max_points_per_voxel,
+                          int max_threads = tbb::task_arena::automatic)
         : voxel_size_(voxel_size),
           max_distance_(max_distance),
-          max_points_per_voxel_(max_points_per_voxel) {}
+          max_points_per_voxel_(max_points_per_voxel),
+          max_threads_(max_threads) {}
 
     Vector3dVectorTuple GetCorrespondences(const Vector3dVector &points,
                                            double max_correspondance_distance) const;
@@ -70,6 +75,7 @@ struct VoxelHashMap {
     double voxel_size_;
     double max_distance_;
     int max_points_per_voxel_;
+    int max_threads_;
     tsl::robin_map<Voxel, VoxelBlock, VoxelHash> map_;
 };
 }  // namespace kiss_icp