Add OpenMP threading to search (#6284)

* Add OpenMP threading to search

* switch to ptrdiff_t

---------

Co-authored-by: jmackay2 <jmackay2>

Author: jmackay2

benchmarks/search/radius_search.cpp

@@ -46,10 +46,14 @@ static void
 BM_KdTreeAll(benchmark::State& state,
              const pcl::PointCloud<pcl::PointXYZ>::Ptr cloudIn,
              const double searchRadius,
-             const size_t neighborLimit)
+             const size_t neighborLimit,
+             const bool threaded)
 {
   pcl::search::KdTree<pcl::PointXYZ> kdtree(false);
   kdtree.setInputCloud(cloudIn);
+  if (threaded) {
+    kdtree.setNumberOfThreads(0);
+  }
 
   // Leaving indices empty to have it search through all points
   pcl::Indices indices;
@@ -123,7 +127,17 @@ main(int argc, char** argv)
       ->Unit(benchmark::kMicrosecond);
 
   benchmark::RegisterBenchmark(
-      "KdTreeAll", &BM_KdTreeAll, cloudFiltered, searchRadius, neighborLimit)
+      "KdTreeAll", &BM_KdTreeAll, cloudFiltered, searchRadius, neighborLimit, false)
+      ->Unit(benchmark::kMicrosecond)
+      ->UseManualTime()
+      ->Iterations(1);
+
+  benchmark::RegisterBenchmark("KdTreeAllThreaded",
+                               &BM_KdTreeAll,
+                               cloudFiltered,
+                               searchRadius,
+                               neighborLimit,
+                               true)
       ->Unit(benchmark::kMicrosecond)
       ->UseManualTime()
       ->Iterations(1);

search/CMakeLists.txt

@@ -3,7 +3,7 @@ set(SUBSYS_DESC "Point cloud generic search library")
 set(SUBSYS_DEPS common kdtree octree)
 
 PCL_SUBSYS_OPTION(build "${SUBSYS_NAME}" "${SUBSYS_DESC}" ON)
-PCL_SUBSYS_DEPEND(build NAME ${SUBSYS_NAME} DEPS ${SUBSYS_DEPS} EXT_DEPS flann)
+PCL_SUBSYS_DEPEND(build NAME ${SUBSYS_NAME} DEPS ${SUBSYS_DEPS} OPT_DEPS OpenMP EXT_DEPS flann)
 
 PCL_ADD_DOC("${SUBSYS_NAME}")
 

search/include/pcl/search/impl/search.hpp

@@ -120,14 +120,16 @@ pcl::search::Search<PointT>::nearestKSearch (
   {
     k_indices.resize (cloud.size ());
     k_sqr_distances.resize (cloud.size ());
-    for (std::size_t i = 0; i < cloud.size (); i++)
+    #pragma omp parallel for num_threads(num_threads_) default(none) shared(cloud, k, k_indices, k_sqr_distances)
+    for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(cloud.size ()); i++)
       nearestKSearch (cloud, static_cast<index_t> (i), k, k_indices[i], k_sqr_distances[i]);
   }
   else
   {
     k_indices.resize (indices.size ());
     k_sqr_distances.resize (indices.size ());
-    for (std::size_t i = 0; i < indices.size (); i++)
+    #pragma omp parallel for num_threads(num_threads_) default(none) shared(cloud, indices, k, k_indices, k_sqr_distances)
+    for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(indices.size ()); i++)
       nearestKSearch (cloud, indices[i], k, k_indices[i], k_sqr_distances[i]);
   }
 }
@@ -172,14 +174,16 @@ pcl::search::Search<PointT>::radiusSearch (
   {
     k_indices.resize (cloud.size ());
     k_sqr_distances.resize (cloud.size ());
-    for (std::size_t i = 0; i < cloud.size (); i++)
+    #pragma omp parallel for num_threads(num_threads_) default(none) shared(cloud, radius, k_indices, k_sqr_distances, max_nn)
+    for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(cloud.size ()); i++)
       radiusSearch (cloud, static_cast<index_t> (i), radius,k_indices[i], k_sqr_distances[i], max_nn);
   }
   else
   {
     k_indices.resize (indices.size ());
     k_sqr_distances.resize (indices.size ());
-    for (std::size_t i = 0; i < indices.size (); i++)
+    #pragma omp parallel for num_threads(num_threads_) default(none) shared(cloud, indices, radius, k_indices, k_sqr_distances, max_nn)
+    for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(indices.size ()); i++)
       radiusSearch (cloud,indices[i],radius,k_indices[i],k_sqr_distances[i], max_nn);
   }
 }

search/include/pcl/search/search.h

@@ -384,7 +384,7 @@ namespace pcl
             pc.resize (cloud.size ());
             for (std::size_t i = 0; i < cloud.size (); ++i)
               pcl::for_each_type <FieldList> (pcl::NdConcatenateFunctor <PointTDiff, PointT> (cloud[i], pc[i]));
-            radiusSearch (pc, Indices (), radius, k_indices, k_sqr_distances, max_nn);
+            radiusSearch (pc, Indices(), radius, k_indices, k_sqr_distances, max_nn);
           }
           else
           {
@@ -395,6 +395,19 @@ namespace pcl
           }
         }
 
+        /** \brief Set the number of threads to use for searching over multiple points or indices
+         * \param[in] nr_threads the number of threads to use (0 automatically sets the threads based on the hardware)
+         */
+        void setNumberOfThreads(unsigned int nr_threads) {
+          #ifdef _OPENMP
+          num_threads_ = nr_threads != 0 ? nr_threads : omp_get_num_procs();
+          #else
+          if (nr_threads != 1) {
+            PCL_WARN("OpenMP is not available. Keeping number of threads unchanged at 1\n");
+          }
+          #endif
+        }
+
       protected:
         void 
         sortResults (Indices& indices, std::vector<float>& distances) const;
@@ -403,6 +416,9 @@ namespace pcl
         IndicesConstPtr indices_;
         bool sorted_results_;
         std::string name_;
+
+        /** \brief The number of threads to use when searching over multiple points or indices */
+        unsigned int num_threads_{1};
 
       private:
         struct Compare