Adds load balanced kernel for point-point, point-linestring and linestring-linestring

cpp/include/cuspatial/detail/algorithm/linestring_distance.cuh

@@ -21,8 +21,13 @@
 
 #include <rmm/device_uvector.hpp>
 
+#include <thrust/binary_search.h>
 #include <thrust/optional.h>
 
+#include <cub/cub.cuh>
+
+#include <cooperative_groups.h>
+
 namespace cuspatial {
 namespace detail {
 
@@ -72,5 +77,84 @@ __global__ void linestring_distance(MultiLinestringRange1 multilinestrings1,
   }
 }
 
+template<typename BoundsIterator, typename IndexType>
+auto __device__ compute_geometry_id(BoundsIterator bounds_begin, BoundsIterator bounds_end, IndexType idx)
+{
+    auto it = thrust::prev(
+      thrust::upper_bound(thrust::seq, bounds_begin, bounds_end, idx));
+    auto const geometry_id = thrust::distance(bounds_begin, it);
+    return thrust::make_pair(it, geometry_id);
+}
+
+/**
+ * @internal
+ * @brief The kernel to compute (multi)linestring to (multi)linestring distance
+ *
+ * Load balanced kernel to compute distances between one pair of segments from the multilinestring
+ * and multilinestring.
+ *
+ * `intersects` is an optional pointer to a boolean range where the `i`th element indicates the
+ * `i`th output should be set to 0 and bypass distance computation. This argument is optional, if
+ * set to nullopt, no distance computation will be bypassed.
+ */
+template <typename T, int block_size, class SegmentRange1, class SegmentRange2, class IndexRange, class OutputIt>
+__global__ void linestring_distance_load_balanced(SegmentRange1 multilinestrings1,
+                                                  SegmentRange2 multilinestrings2,
+                                                  IndexRange thread_bounds,
+                                                  thrust::optional<uint8_t*> intersects,
+                                                  OutputIt distances)
+{
+  using index_t = typename IndexRange::value_type;
+
+  auto num_segment_pairs = thread_bounds[thread_bounds.size() - 1];
+  for (auto idx = threadIdx.x + blockIdx.x * blockDim.x; idx < num_segment_pairs;
+       idx += gridDim.x * blockDim.x) {
+
+    auto [it, geometry_id] = compute_geometry_id(thread_bounds.begin(), thread_bounds.end(), idx);
+
+    auto first_thread_of_this_block = blockDim.x * blockIdx.x;
+    auto last_thread_of_block = blockDim.x * (blockIdx.x + 1) - 1;
+    auto first_thread_of_next_geometry = thread_bounds[geometry_id + 1];
+
+    bool split_block = first_thread_of_this_block < first_thread_of_next_geometry && first_thread_of_next_geometry <= last_thread_of_block;
+
+    if (intersects.has_value() && intersects.value()[geometry_id]) {
+      distances[geometry_id] = 0.0f;
+      continue;
+    }
+
+    auto const local_idx = idx - *it;
+    // Retrieve the number of segments in multilinestrings[geometry_id]
+    auto num_segment_this_multilinestring =
+      multilinestrings1.multigeometry_count_begin()[geometry_id];
+    // The segment id from multilinestring1 this thread is computing (local_id + global_offset)
+    auto multilinestring1_segment_id = local_idx % num_segment_this_multilinestring +
+                                       multilinestrings1.multigeometry_offset_begin()[geometry_id];
+
+    // The segment id from multilinestring2 this thread is computing (local_id + global_offset)
+    auto multilinestring2_segment_id = local_idx / num_segment_this_multilinestring +
+                                       multilinestrings2.multigeometry_offset_begin()[geometry_id];
+
+    auto [a, b] = multilinestrings1.begin()[multilinestring1_segment_id];
+    auto [c, d] = multilinestrings2.begin()[multilinestring2_segment_id];
+
+    auto partial = sqrt(squared_segment_distance(a, b, c, d));
+
+    if (split_block)
+      atomicMin(&distances[geometry_id], partial);
+    else
+    {
+      // block reduce
+      typedef cub::BlockReduce<T, block_size> BlockReduce;
+      __shared__ typename BlockReduce::TempStorage temp_storage;
+      auto aggregate = BlockReduce(temp_storage).Reduce(partial, cub::Min());
+
+      // atmomic with leading thread
+      if (cooperative_groups::this_thread_block().thread_rank() == 0)
+        atomicMin(&distances[geometry_id], aggregate);
+    }
+  }
+}
+
 }  // namespace detail
 }  // namespace cuspatial

cpp/include/cuspatial/detail/distance/distance_utils.cuh

@@ -29,6 +29,7 @@
 #include <thrust/iterator/discard_iterator.h>
 #include <thrust/logical.h>
 #include <thrust/tabulate.h>
+#include <thrust/zip_function.h>
 
 namespace cuspatial {
 namespace detail {
@@ -115,5 +116,45 @@ rmm::device_uvector<uint8_t> point_polygon_intersects(MultiPointRange multipoint
   return multipoint_intersects;
 }
 
+/**
+ * @brief Compute the thread bound between two ranges of partitions
+ *
+ * @tparam CountIterator1
+ * @tparam CountIterator2
+ * @param lhs
+ * @param rhs
+ * @param stream
+ * @return rmm::device_uvector<IndexType>
+ */
+template <typename CountIterator1,
+          typename CountIterator2,
+          typename index_t = iterator_value_type<CountIterator1>>
+rmm::device_uvector<index_t> compute_segment_thread_bounds(CountIterator1 lhs_begin,
+                                                           CountIterator1 lhs_end,
+                                                           CountIterator2 rhs_begin,
+                                                           rmm::cuda_stream_view stream)
+{
+  auto size = thrust::distance(lhs_begin, lhs_end) + 1;
+
+  // Compute the "boundary" of threads. Threads are partitioned based on the number of linestrings
+  // times the number of polygons in a multilinestring-multipolygon pair.
+  auto segment_count_product_it =
+    thrust::make_transform_iterator(thrust::make_zip_iterator(lhs_begin, rhs_begin),
+                                    thrust::make_zip_function(thrust::multiplies<index_t>{}));
+
+  // Computes the "thread boundary" of each pair. This array partitions the thread range by
+  // geometries. E.g. threadIdx within [thread_bounds[i], thread_bounds[i+1]) computes distances of
+  // the ith pair.
+  auto thread_bounds = rmm::device_uvector<index_t>(size, stream);
+  detail::zero_data_async(thread_bounds.begin(), thread_bounds.end(), stream);
+
+  thrust::inclusive_scan(rmm::exec_policy(stream),
+                         segment_count_product_it,
+                         segment_count_product_it + thread_bounds.size() - 1,
+                         thrust::next(thread_bounds.begin()));
+
+  return thread_bounds;
+}
+
 }  // namespace detail
 }  // namespace cuspatial

cpp/include/cuspatial/detail/distance/linestring_distance.cuh

@@ -17,9 +17,11 @@
 #pragma once
 
 #include <cuspatial/detail/algorithm/linestring_distance.cuh>
+#include <cuspatial/detail/distance/distance_utils.cuh>
 #include <cuspatial/error.hpp>
 #include <cuspatial/geometry/vec_2d.hpp>
 #include <cuspatial/traits.hpp>
+#include <cuspatial/range/range.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
@@ -33,8 +35,8 @@
 namespace cuspatial {
 
 template <class MultiLinestringRange1, class MultiLinestringRange2, class OutputIt>
-OutputIt pairwise_linestring_distance(MultiLinestringRange1 multilinestrings1,
-                                      MultiLinestringRange2 multilinestrings2,
+OutputIt pairwise_linestring_distance(MultiLinestringRange1 lhs,
+                                      MultiLinestringRange2 rhs,
                                       OutputIt distances_first,
                                       rmm::cuda_stream_view stream)
 {
@@ -48,25 +50,43 @@ OutputIt pairwise_linestring_distance(MultiLinestringRange1 multilinestrings1,
                         typename MultiLinestringRange2::point_t>(),
                 "All input types must be cuspatial::vec_2d with the same value type");
 
-  CUSPATIAL_EXPECTS(multilinestrings1.size() == multilinestrings2.size(),
-                    "Inputs must have the same number of rows.");
+  CUSPATIAL_EXPECTS(lhs.size() == rhs.size(), "Inputs must have the same number of rows.");
 
-  if (multilinestrings1.size() == 0) return distances_first;
+  if (lhs.size() == 0) return distances_first;
 
+  // Make views to the segments in the multilinestring
+  auto lhs_segments       = lhs._segments(stream);
+  auto lhs_segments_range = lhs_segments.view();
+
+  // Make views to the segments in the multilinestring
+  auto rhs_segments       = rhs._segments(stream);
+  auto rhs_segments_range = rhs_segments.view();
+
+  auto thread_bounds =
+    detail::compute_segment_thread_bounds(lhs_segments_range.multigeometry_count_begin(),
+                                          lhs_segments_range.multigeometry_count_end(),
+                                          rhs_segments_range.multigeometry_count_begin(),
+                                          stream);
+  // Initialize the output range
   thrust::fill(rmm::exec_policy(stream),
                distances_first,
-               distances_first + multilinestrings1.size(),
+               distances_first + lhs.size(),
                std::numeric_limits<T>::max());
 
   std::size_t constexpr threads_per_block = 256;
-  std::size_t const num_blocks =
-    (multilinestrings1.num_points() + threads_per_block - 1) / threads_per_block;
+  std::size_t num_threads                 = 1e8;
+  std::size_t const num_blocks = (num_threads + threads_per_block - 1) / threads_per_block;
 
-  detail::linestring_distance<<<num_blocks, threads_per_block, 0, stream.value()>>>(
-    multilinestrings1, multilinestrings2, thrust::nullopt, distances_first);
+  detail::linestring_distance_load_balanced<T, threads_per_block>
+    <<<num_blocks, threads_per_block, 0, stream.value()>>>(
+      lhs_segments_range,
+      rhs_segments_range,
+      range(thread_bounds.begin(), thread_bounds.end()),
+      thrust::nullopt,
+      distances_first);
 
   CUSPATIAL_CUDA_TRY(cudaGetLastError());
-  return distances_first + multilinestrings1.size();
+  return distances_first + lhs.size();
 }
 
 }  // namespace cuspatial