Introduce multilinestring_range structure, simplifies point-linestring distance API

cpp/include/cuspatial/detail/utility/device_atomics.cuh

@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <thrust/device_ptr.h>
+
 #include <type_traits>
 
 namespace cuspatial {
@@ -146,6 +148,44 @@ __device__ inline float atomicMin(float* addr, float val)
   return atomicOp<float>(addr, val, min_<float>);
 }
 
+/**
+ * @internal
+ * @brief CUDA device atomic minimum for double
+ *
+ * Atomically computes the min of the value stored in `ptr` and `val` and stores it in `ptr`,
+ * returning the previous value of `*ptr`.
+ *
+ * @note based on https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
+ *
+ * @param ptr The thrust device pointer to the address to atomically compare and update with the
+ * minimum.
+ * @param val The value to compare
+ * @return The old value stored in `addr`.
+ */
+__device__ inline double atomicMin(thrust::device_ptr<double> ptr, double val)
+{
+  return atomicMin(thrust::raw_pointer_cast(ptr), val);
+}
+
+/**
+ * @internal
+ * @brief CUDA device atomic minimum for float
+ *
+ * Atomically computes the min of the value stored in `ptr` and `val` and stores it in `ptr`,
+ * returning the previous value of `*ptr`.
+ *
+ * @note based on https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
+ *
+ * @param ptr The thrust device pointer to address to atomically compare and update with the
+ * minimum.
+ * @param val The value to compare
+ * @return The old value stored in `addr`.
+ */
+__device__ inline float atomicMin(thrust::device_ptr<float> ptr, float val)
+{
+  return atomicMin(thrust::raw_pointer_cast(ptr), val);
+}
+
 /**
  * @internal
  * @brief CUDA device atomic maximum for double
@@ -182,5 +222,43 @@ __device__ inline float atomicMax(float* addr, float val)
   return atomicOp<float>(addr, val, max_<float>);
 }
 
+/**
+ * @internal
+ * @brief CUDA device atomic maximum for double
+ *
+ * Atomically computes the max of the value stored in `addr` and `val` and stores it in `ptr`,
+ * returning the previous value of `*ptr`.
+ *
+ * @note based on https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
+ *
+ * @param ptr The thrust device pointer to the address atomically compare and update with the
+ * maximum.
+ * @param val The value to compare
+ * @return The old value stored in `addr`.
+ */
+__device__ inline double atomicMax(thrust::device_ptr<double> ptr, double val)
+{
+  return atomicMax(thrust::raw_pointer_cast(ptr), val);
+}
+
+/**
+ * @internal
+ * @brief CUDA device atomic maximum for float
+ *
+ * Atomically computes the max of the value stored in `addr` and `val` and stores it in `addr`,
+ * returning the previous value of `*addr`.
+ *
+ * @note based on https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomic-functions
+ *
+ * @param addr The thrust device pointer to the address to atomically compare and update with the
+ * maximum.
+ * @param val The value to compare
+ * @return The old value stored in `addr`.
+ */
+__device__ inline float atomicMax(thrust::device_ptr<float> ptr, float val)
+{
+  return atomicMax(thrust::raw_pointer_cast(ptr), val);
+}
+
 }  // namespace detail
 }  // namespace cuspatial

cpp/include/cuspatial/experimental/detail/point_linestring_distance.cuh

@@ -41,176 +41,79 @@ namespace cuspatial {
 namespace detail {
 
 /**
- * @internal
- * @brief The kernel to compute multi-point to multi-linestring distance
+ * @brief Kernel to compute the distance between pairs of point and linestring.
  *
- * Each thread computes the distance between a line segment in the linestring and the
- * corresponding multi-point part in the pair. The shortest distance is computed in the
- * output array via an atomic operation.
- *
- * @tparam Cart2dItA Iterator to 2d cartesian coordinates. Must meet requirements of
- * [LegacyRandomAccessIterator][LinkLRAI] and be device-accessible.
- * @tparam Cart2dItB Iterator to 2d cartesian coordinates. Must meet requirements of
- * [LegacyRandomAccessIterator][LinkLRAI] and be device-accessible.
- * @tparam OffsetIteratorA Iterator to offsets. Must meet requirements of
- * [LegacyRandomAccessIterator][LinkLRAI] and be device-accessible.
- * @tparam OffsetIteratorB Iterator to offsets. Must meet requirements of
- * [LegacyRandomAccessIterator][LinkLRAI] and be device-accessible.
- * @tparam OffsetIteratorC Iterator to offsets. Must meet requirements of
- * [LegacyRandomAccessIterator][LinkLRAI] and be device-accessible.
- * @tparam OutputIterator Iterator to output distances. Must meet requirements of
- * [LegacyRandomAccessIterator][LinkLRAI] and be device-accessible and mutable.
- *
- * @param[in] point_geometry_offset_first Iterator to the beginning of the range of the multipoint
- * parts
- * @param[in] point_geometry_offset_last Iterator to the end of the range of the multipoint parts
- * @param[in] points_first Iterator to the beginning of the range of the points
- * @param[in] points_last  Iterator to the end of the range of the points
- * @param[in] linestring_geometry_offset_first Iterator to the beginning of the range of the
- * linestring parts
- * @param[in] linestring_geometry_offset_last Iterator to the end of the range of the linestring
- * parts
- * @param[in] linestring_part_offsets_first Iterator to the beginning of the range of the linestring
- * offsets
- * @param[in] linestring_part_offsets_last Iterator to the beginning of the range of the linestring
- * offsets
- * @param[in] linestring_points_first Iterator to the beginning of the range of the linestring
- * points
- * @param[in] linestring_points_last Iterator to the end of the range of the linestring points
- * @param[out] distances Iterator to the output range of shortest distances between pairs.
- *
- * [LinkLRAI]: https://en.cppreference.com/w/cpp/named_req/RandomAccessIterator
- * "LegacyRandomAccessIterator"
+ * The kernel is launched on one linestring point per thread. Each thread iterates on all points in
+ * the multipoint operand and use atomics to aggregate the shortest distance.
  */
-template <class Cart2dItA,
-          class Cart2dItB,
-          class OffsetIteratorA,
-          class OffsetIteratorB,
-          class OffsetIteratorC,
-          class OutputIterator>
-void __global__
-pairwise_point_linestring_distance_kernel(OffsetIteratorA point_geometry_offset_first,
-                                          OffsetIteratorA point_geometry_offset_last,
-                                          Cart2dItA points_first,
-                                          Cart2dItA points_last,
-                                          OffsetIteratorB linestring_geometry_offset_first,
-                                          OffsetIteratorB linestring_geometry_offset_last,
-                                          OffsetIteratorC linestring_part_offsets_first,
-                                          OffsetIteratorC linestring_part_offsets_last,
-                                          Cart2dItB linestring_points_first,
-                                          Cart2dItB linestring_points_last,
-                                          OutputIterator distances)
+template <class MultiPointRange, class MultiLinestringRange, class OutputIterator>
+void __global__ pairwise_point_linestring_distance_kernel(MultiPointRange multipoints,
+                                                          MultiLinestringRange multilinestrings,
+                                                          OutputIterator distances)
 {
-  using T = iterator_vec_base_type<Cart2dItA>;
+  using T = typename MultiPointRange::element_t;
 
-  for (auto idx = threadIdx.x + blockIdx.x * blockDim.x;
-       idx < std::distance(linestring_points_first, thrust::prev(linestring_points_last));
+  for (auto idx = threadIdx.x + blockIdx.x * blockDim.x; idx < multilinestrings.num_points();
        idx += gridDim.x * blockDim.x) {
     // Search from the part offsets array to determine the part idx of current linestring point
-    auto linestring_part_offsets_iter = thrust::upper_bound(
-      thrust::seq, linestring_part_offsets_first, linestring_part_offsets_last, idx);
-
+    auto part_idx = multilinestrings.part_idx_from_point_idx(idx);
     // Pointer to the last point in the linestring, skip iteration.
     // Note that the last point for the last linestring is guarded by the grid-stride loop.
-    if (linestring_part_offsets_iter != linestring_part_offsets_last &&
-        *linestring_part_offsets_iter - 1 == idx) {
-      continue;
-    }
-
-    auto part_offsets_idx =
-      thrust::distance(linestring_part_offsets_first, thrust::prev(linestring_part_offsets_iter));
+    if (!multilinestrings.is_valid_segment_id(idx, part_idx)) continue;
 
-    // Search from the linestring geometry offsets array to determine the geometry idx of current
-    // linestring point
-    auto geometry_offsets_iter = thrust::upper_bound(thrust::seq,
-                                                     linestring_geometry_offset_first,
-                                                     linestring_geometry_offset_last,
-                                                     part_offsets_idx);
-    // geometry_idx is also the index to corresponding multipoint in the pair
-    auto geometry_idx =
-      thrust::distance(linestring_geometry_offset_first, thrust::prev(geometry_offsets_iter));
+    // Search from the linestring geometry offsets array to determine the geometry idx of
+    // current linestring point
+    auto geometry_idx = multilinestrings.geometry_idx_from_part_idx(part_idx);
 
     // Reduce the minimum distance between different parts of the multi-point.
-    vec_2d<T> const a      = linestring_points_first[idx];
-    vec_2d<T> const b      = linestring_points_first[idx + 1];
+    auto [a, b]            = multilinestrings.segment(idx);
     T min_distance_squared = std::numeric_limits<T>::max();
 
-    for (auto point_idx = point_geometry_offset_first[geometry_idx];
-         point_idx < point_geometry_offset_first[geometry_idx + 1];
-         point_idx++) {
-      vec_2d<T> const c = points_first[point_idx];
-
+    for (vec_2d<T> const& c : multipoints[geometry_idx]) {
       // TODO: reduce redundant computation only related to `a`, `b` in this helper.
       auto const distance_squared = point_to_segment_distance_squared(c, a, b);
-      min_distance_squared        = std::min(distance_squared, min_distance_squared);
+      min_distance_squared        = min(distance_squared, min_distance_squared);
     }
-
-    atomicMin(&thrust::raw_reference_cast(*(distances + geometry_idx)),
-              static_cast<T>(std::sqrt(min_distance_squared)));
+    atomicMin(&distances[geometry_idx], static_cast<T>(sqrt(min_distance_squared)));
   }
 }
 
 }  // namespace detail
-
-template <class Cart2dItA,
-          class Cart2dItB,
-          class OffsetIteratorA,
-          class OffsetIteratorB,
-          class OffsetIteratorC,
-          class OutputIt>
-OutputIt pairwise_point_linestring_distance(OffsetIteratorA point_geometry_offset_first,
-                                            OffsetIteratorA point_geometry_offset_last,
-                                            Cart2dItA points_first,
-                                            Cart2dItA points_last,
-                                            OffsetIteratorB linestring_geometry_offset_first,
-                                            OffsetIteratorC linestring_part_offsets_first,
-                                            OffsetIteratorC linestring_part_offsets_last,
-                                            Cart2dItB linestring_points_first,
-                                            Cart2dItB linestring_points_last,
+template <class MultiPointRange, class MultiLinestringRange, class OutputIt>
+OutputIt pairwise_point_linestring_distance(MultiPointRange multipoints,
+                                            MultiLinestringRange multilinestrings,
                                             OutputIt distances_first,
                                             rmm::cuda_stream_view stream)
 {
-  using T = iterator_vec_base_type<Cart2dItA>;
+  using T = typename MultiPointRange::element_t;
 
-  static_assert(is_same_floating_point<T, iterator_vec_base_type<Cart2dItB>>(),
+  static_assert(is_same_floating_point<T, typename MultiLinestringRange::element_t>(),
                 "Inputs must have same floating point value type.");
 
   static_assert(
-    is_same<vec_2d<T>, iterator_value_type<Cart2dItA>, iterator_value_type<Cart2dItB>>(),
+    is_same<vec_2d<T>, typename MultiPointRange::point_t, typename MultiLinestringRange::point_t>(),
     "Inputs must be cuspatial::vec_2d");
 
-  auto const num_pairs =
-    thrust::distance(point_geometry_offset_first, point_geometry_offset_last) - 1;
-
-  if (num_pairs == 0) { return distances_first; }
-
-  auto const num_linestring_points =
-    thrust::distance(linestring_points_first, linestring_points_last);
+  CUSPATIAL_EXPECTS(multilinestrings.size() == multipoints.size(),
+                    "Input must have the same number of rows.");
+  if (multilinestrings.size() == 0) { return distances_first; }
 
-  thrust::fill_n(
-    rmm::exec_policy(stream), distances_first, num_pairs, std::numeric_limits<T>::max());
+  thrust::fill_n(rmm::exec_policy(stream),
+                 distances_first,
+                 multilinestrings.size(),
+                 std::numeric_limits<T>::max());
 
   std::size_t constexpr threads_per_block = 256;
   std::size_t const num_blocks =
-    (num_linestring_points + threads_per_block - 1) / threads_per_block;
+    (multilinestrings.size() + threads_per_block - 1) / threads_per_block;
 
   detail::
     pairwise_point_linestring_distance_kernel<<<num_blocks, threads_per_block, 0, stream.value()>>>(
-      point_geometry_offset_first,
-      point_geometry_offset_last,
-      points_first,
-      points_last,
-      linestring_geometry_offset_first,
-      linestring_geometry_offset_first + num_pairs + 1,
-      linestring_part_offsets_first,
-      linestring_part_offsets_last,
-      linestring_points_first,
-      linestring_points_last,
-      distances_first);
+      multipoints, multilinestrings, distances_first);
 
   CUSPATIAL_CUDA_TRY(cudaGetLastError());
 
-  return distances_first + num_pairs;
+  return distances_first + multilinestrings.size();
 }
 
 }  // namespace cuspatial