Add support for limiting num_threads in tbb task

config/advanced.yaml

@@ -20,3 +20,4 @@ adaptive_threshold:
 registration:
   max_num_iterations: 500 # <- optional
   convergence_criterion: 0.0001 # <- optional
+  max_num_threads: 0 # <- optional, 0 means automatic

cpp/kiss_icp/core/Registration.cpp

@@ -23,6 +23,8 @@
 #include "Registration.hpp"
 
 #include <tbb/blocked_range.h>
+#include <tbb/global_control.h>
+#include <tbb/info.h>
 #include <tbb/parallel_reduce.h>
 
 #include <algorithm>
@@ -142,6 +144,17 @@ LinearSystem BuildLinearSystem(const Associations &associations, double kernel)
 
 namespace kiss_icp {
 
+Registration::Registration(int max_num_iteration, double convergence_criterion, int max_num_threads)
+    : max_num_iterations_(max_num_iteration),
+      convergence_criterion_(convergence_criterion),
+      // Only manipulate the number of threads if the user specifies something greater than 0
+      max_num_threads_(max_num_threads > 0 ? max_num_threads : tbb::info::default_concurrency()) {
+    // This global variable requires static duration storage to be able to manipulate the max
+    // concurrency from TBB across the entire class
+    static const auto tbb_control_settings = tbb::global_control(
+        tbb::global_control::max_allowed_parallelism, static_cast<size_t>(max_num_threads_));
+}
+
 Sophus::SE3d Registration::AlignPointsToMap(const std::vector<Eigen::Vector3d> &frame,
                                             const VoxelHashMap &voxel_map,
                                             const Sophus::SE3d &initial_guess,

cpp/kiss_icp/core/Registration.hpp

@@ -31,8 +31,7 @@
 namespace kiss_icp {
 
 struct Registration {
-    explicit Registration(int max_num_iteration, double convergence_criterion)
-        : max_num_iterations_(max_num_iteration), convergence_criterion_(convergence_criterion) {}
+    explicit Registration(int max_num_iteration, double convergence_criterion, int max_num_threads);
 
     Sophus::SE3d AlignPointsToMap(const std::vector<Eigen::Vector3d> &frame,
                                   const VoxelHashMap &voxel_map,
@@ -42,5 +41,6 @@ struct Registration {
 
     int max_num_iterations_;
     double convergence_criterion_;
+    int max_num_threads_;
 };
 }  // namespace kiss_icp

cpp/kiss_icp/pipeline/KissICP.hpp

@@ -46,6 +46,7 @@ struct KISSConfig {
     // registration params
     int max_num_iterations = 500;
     double convergence_criterion = 0.0001;
+    int max_num_threads = 0;
 
     // Motion compensation
     bool deskew = false;
@@ -59,7 +60,8 @@ class KissICP {
 public:
     explicit KissICP(const KISSConfig &config)
         : config_(config),
-          registration_(config.max_num_iterations, config.convergence_criterion),
+          registration_(
+              config.max_num_iterations, config.convergence_criterion, config.max_num_threads),
           local_map_(config.voxel_size, config.max_range, config.max_points_per_voxel),
           adaptive_threshold_(config.initial_threshold, config.min_motion_th, config.max_range) {}
 

python/kiss_icp/config/config.py

@@ -39,6 +39,7 @@ class MappingConfig(BaseModel):
 class RegistrationConfig(BaseModel):
     max_num_iterations: Optional[int] = 500
     convergence_criterion: Optional[float] = 0.0001
+    max_num_threads: Optional[int] = 0  # 0 means automatic
 
 
 class AdaptiveThresholdConfig(BaseModel):

python/kiss_icp/pybind/kiss_icp_pybind.cpp

@@ -75,7 +75,8 @@ PYBIND11_MODULE(kiss_icp_pybind, m) {
     // Point Cloud registration
     py::class_<Registration> internal_registration(m, "_Registration", "Don't use this");
     internal_registration
-        .def(py::init<int, double>(), "max_num_iterations"_a, "convergence_criterion"_a)
+        .def(py::init<int, double, int>(), "max_num_iterations"_a, "convergence_criterion"_a,
+             "max_num_threads"_a)
         .def(
             "_align_points_to_map",
             [](Registration &self, const std::vector<Eigen::Vector3d> &points,

python/kiss_icp/registration.py

@@ -31,14 +31,21 @@ def get_registration(config: KISSConfig):
     return Registration(
         max_num_iterations=config.registration.max_num_iterations,
         convergence_criterion=config.registration.convergence_criterion,
+        max_num_threads=config.registration.max_num_threads,
     )
 
 
 class Registration:
-    def __init__(self, max_num_iterations: int, convergence_criterion: float):
+    def __init__(
+        self,
+        max_num_iterations: int,
+        convergence_criterion: float,
+        max_num_threads: int = 0,
+    ):
         self._registration = kiss_icp_pybind._Registration(
             max_num_iterations=max_num_iterations,
             convergence_criterion=convergence_criterion,
+            max_num_threads=max_num_threads,
         )
 
     def align_points_to_map(