Add support for limiting num_threads in tbb task (#252)

* Add support for limiting num_threads in tbb task * Make num_threads a paramter of the voxel hash map For backwards compatibility, use a default value of "automatic" * This got a bit more hacky but is to avoid hardocding a praamter * Propagate max_threads to Python config * First draft on core library * Update python API * Rearange * Remove type alias * Fix build * Wrap constants into a configuration struct * Split the watters * It's all about drafts * Update python API I'm alredy regreting about this * Some renaming just because * Changing names trying to auto-convince myself... * Fix c++ build * rename function * renaming variables * renaming, should be one neighbor only * Tizianified a little bit * Draft on voxelhashmap * Rename * Rename Correspondences -> Associations Otherwise too much mix a match on my opinion * Move stuff around only * Remove redunant name * They are not there, we need to find them! * Shrink * Tiziano shows to guys -> with for_each * Tiziano shows to guys -> with transform_reduce....sexy * Consistent naming * Bring comments for readbilty * rename variable * Sacrifice name for one-liner * AlignCloudToMap -> AlignPointsToMap * Make rename like a book on ProbRob * Revert "Make rename like a book on ProbRob" This reverts commit 7a45c9d. * estimation_threshold -> convergence_criterion Requested by benedikt * Rename threshold also here * no typos no nacho * Remove comment * reduce diff * Small improvement. max_num_threads_ always represents what it says Assuming 10 logical cores: Before Registration::max_num_threads_ == 0, and now Registration::max_num_threads_ == 10. This way we hide less what's going under the hood * Remove single letter variable --------- Co-authored-by: Ignacio Vizzo <[email protected]> Co-authored-by: Benedikt Mersch <[email protected]> Co-authored-by: tizianoGuadagnino <[email protected]>
PRBonn · Mar 18, 2024 · 994d232 · 994d232
1 parent c2ea460
commit 994d232
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 5 deletions.
diff --git a/config/advanced.yaml b/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
diff --git a/cpp/kiss_icp/core/Registration.cpp b/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,

diff --git a/cpp/kiss_icp/core/Registration.hpp b/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
diff --git a/cpp/kiss_icp/pipeline/KissICP.hpp b/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) {}
 

diff --git a/python/kiss_icp/config/config.py b/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):

diff --git a/python/kiss_icp/pybind/kiss_icp_pybind.cpp b/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,

diff --git a/python/kiss_icp/registration.py b/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(