diff --git a/src/gtsam_points/util/covariance_estimation.cpp b/src/gtsam_points/util/covariance_estimation.cpp index d2d0f21..e173e6e 100644 --- a/src/gtsam_points/util/covariance_estimation.cpp +++ b/src/gtsam_points/util/covariance_estimation.cpp @@ -6,6 +6,11 @@ #include #include #include +#include + +#ifdef GTSAM_POINTS_TBB +#include +#endif namespace gtsam_points { @@ -13,8 +18,7 @@ std::vector estimate_covariances(const Eigen::Vector4d* points, KdTree tree(points, num_points, params.num_threads); std::vector 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 k_indices(params.k_neighbors); std::vector 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]); @@ -22,7 +26,7 @@ std::vector estimate_covariances(const Eigen::Vector4d* points, 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 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(0, num_points, 8), [&](const tbb::blocked_range& 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 index 9ef00f3..cc6257d 100644 --- a/src/gtsam_points/util/normal_estimation.cpp +++ b/src/gtsam_points/util/normal_estimation.cpp @@ -5,14 +5,18 @@ #include #include +#include + +#ifdef GTSAM_POINTS_TBB +#include +#endif namespace gtsam_points { std::vector estimate_normals(const Eigen::Vector4d* points, const Eigen::Matrix4d* covs, int num_points, int num_threads) { std::vector 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 eig; eig.computeDirect(covs[i].block<3, 3>(0, 0)); normals[i].head<3>() = eig.eigenvectors().col(0); @@ -20,6 +24,24 @@ std::vector estimate_normals(const Eigen::Vector4d* points, con 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(0, num_points, 64), [&](const tbb::blocked_range& 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;