diff --git a/config/advanced.yaml b/config/advanced.yaml index a1650c10..102f2890 100644 --- 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 index 90a6473e..2634b0a8 100644 --- a/cpp/kiss_icp/core/Registration.cpp +++ b/cpp/kiss_icp/core/Registration.cpp @@ -23,6 +23,8 @@ #include "Registration.hpp" #include +#include +#include #include #include @@ -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(max_num_threads_)); +} + Sophus::SE3d Registration::AlignPointsToMap(const std::vector &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 index 25d3abf5..d878c2fc 100644 --- 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 &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 index 8605b44a..e43c6093 100644 --- 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 index 92dc8d98..cbb9753b 100644 --- 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 index c2ff3eb2..b514a6f5 100644 --- 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_ internal_registration(m, "_Registration", "Don't use this"); internal_registration - .def(py::init(), "max_num_iterations"_a, "convergence_criterion"_a) + .def(py::init(), "max_num_iterations"_a, "convergence_criterion"_a, + "max_num_threads"_a) .def( "_align_points_to_map", [](Registration &self, const std::vector &points, diff --git a/python/kiss_icp/registration.py b/python/kiss_icp/registration.py index 48a05d32..65282dfc 100644 --- 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(