Refactor ST_Distance header only API

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

@@ -16,7 +16,7 @@
 
 #pragma once
 
-#include <cuspatial/detail/algorithm/linestring_distance.cuh>
+#include <cuspatial/detail/kernel/pairwise_distance.cuh>
 #include <cuspatial/error.hpp>
 #include <cuspatial/geometry/vec_2d.hpp>
 #include <cuspatial/traits.hpp>

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

@@ -20,92 +20,16 @@
 
 #include <cuspatial/cuda_utils.hpp>
 #include <cuspatial/detail/algorithm/is_point_in_polygon.cuh>
-#include <cuspatial/detail/functors.cuh>
-#include <cuspatial/detail/utility/device_atomics.cuh>
-#include <cuspatial/detail/utility/linestring.cuh>
-#include <cuspatial/detail/utility/zero_data.cuh>
-#include <cuspatial/geometry/vec_2d.hpp>
-#include <cuspatial/iterator_factory.cuh>
-#include <cuspatial/range/range.cuh>
+#include <cuspatial/detail/kernel/pairwise_distance.cuh>
 
 #include <thrust/fill.h>
-#include <thrust/functional.h>
-#include <thrust/iterator/discard_iterator.h>
-#include <thrust/iterator/permutation_iterator.h>
-#include <thrust/logical.h>
-#include <thrust/scan.h>
-#include <thrust/zip_function.h>
 
 #include <rmm/cuda_stream_view.hpp>
-#include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
 #include <limits>
 
 namespace cuspatial {
-namespace detail {
-
-/**
- * @brief Computes distances between the multilinestring and multipolygons
- *
- * @param multilinestrings Range to the multilinestring
- * @param multipolygons Range to the multipolygon
- * @param thread_bounds Range to the boundary of thread partitions
- * @param multilinestrings_segment_offsets Range to the indices where the first segment of each
- * multilinestring begins
- * @param multipolygons_segment_offsets Range to the indices where the first segment of each
- * multipolygon begins
- * @param intersects A uint8_t array that indicates if the corresponding pair of multipoint and
- * multipolygon intersects
- * @param distances Output range of distances, pre-filled with std::numerical_limits<T>::max()
- */
-template <typename MultiLinestringRange,
-          typename MultiPolygonRange,
-          typename IndexRange,
-          typename OutputIt>
-void __global__
-pairwise_linestring_polygon_distance_kernel(MultiLinestringRange multilinestrings,
-                                            MultiPolygonRange multipolygons,
-                                            IndexRange thread_bounds,
-                                            IndexRange multilinestrings_segment_offsets,
-                                            IndexRange multipolygons_segment_offsets,
-                                            uint8_t* intersects,
-                                            OutputIt* distances)
-{
-  using T       = typename MultiLinestringRange::element_t;
-  using index_t = iterator_value_type<typename MultiLinestringRange::geometry_it_t>;
-
-  auto num_threads = thread_bounds[thread_bounds.size() - 1];
-  for (auto idx = blockDim.x * blockIdx.x + threadIdx.x; idx < num_threads;
-       idx += blockDim.x * gridDim.x) {
-    auto it = thrust::prev(
-      thrust::upper_bound(thrust::seq, thread_bounds.begin(), thread_bounds.end(), idx));
-    auto geometry_id = thrust::distance(thread_bounds.begin(), it);
-    auto local_idx   = idx - *it;
-
-    if (intersects[geometry_id]) {
-      distances[geometry_id] = 0.0f;
-      continue;
-    }
-
-    // Retrieve the number of segments in multilinestrings[geometry_id]
-    auto num_segment_this_multilinestring =
-      multilinestrings.multilinestring_segment_count_begin()[geometry_id];
-    // The segment id from the multilinestring this thread is computing (local_id + global_offset)
-    auto multilinestring_segment_id =
-      local_idx % num_segment_this_multilinestring + multilinestrings_segment_offsets[geometry_id];
-    // The segment id from the multipolygon this thread is computing (local_id + global_offset)
-    auto multipolygon_segment_id =
-      local_idx / num_segment_this_multilinestring + multipolygons_segment_offsets[geometry_id];
-
-    auto [a, b] = multilinestrings.segment_begin()[multilinestring_segment_id];
-    auto [c, d] = multipolygons.segment_begin()[multipolygon_segment_id];
-
-    atomicMin(&distances[geometry_id], sqrt(squared_segment_distance(a, b, c, d)));
-  }
-};
-
-}  // namespace detail
 
 template <class MultiLinestringRange, class MultiPolygonRange, class OutputIt>
 OutputIt pairwise_linestring_polygon_distance(MultiLinestringRange multilinestrings,
@@ -119,70 +43,25 @@ OutputIt pairwise_linestring_polygon_distance(MultiLinestringRange multilinestri
   CUSPATIAL_EXPECTS(multilinestrings.size() == multipolygons.size(),
                     "Must have the same number of input rows.");
 
-  if (multilinestrings.size() == 0) return distances_first;
+  auto size = multilinestrings.size();
+
+  if (size == 0) return distances_first;
 
   // Create a multipoint range from multilinestrings, computes intersection
-  auto multipoints           = multilinestrings.as_multipoint_range();
-  auto multipoint_intersects = point_polygon_intersects(multipoints, multipolygons, stream);
-
-  // Compute the "boundary" of threads. Threads are partitioned based on the number of linestrings
-  // times the number of polygons in a multipoint-multipolygon pair.
-  auto segment_count_product_it = thrust::make_transform_iterator(
-    thrust::make_zip_iterator(multilinestrings.multilinestring_segment_count_begin(),
-                              multipolygons.multipolygon_segment_count_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>(multilinestrings.size() + 1, 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()));
-
-  // Compute offsets to the first segment of each multilinestring and multipolygon
-  auto multilinestring_segment_offsets =
-    rmm::device_uvector<index_t>(multilinestrings.num_multilinestrings() + 1, stream);
-  detail::zero_data_async(
-    multilinestring_segment_offsets.begin(), multilinestring_segment_offsets.end(), stream);
-
-  auto multipolygon_segment_offsets =
-    rmm::device_uvector<index_t>(multipolygons.num_multipolygons() + 1, stream);
-  detail::zero_data_async(
-    multipolygon_segment_offsets.begin(), multipolygon_segment_offsets.end(), stream);
-
-  thrust::inclusive_scan(rmm::exec_policy(stream),
-                         multilinestrings.multilinestring_segment_count_begin(),
-                         multilinestrings.multilinestring_segment_count_begin() +
-                           multilinestrings.num_multilinestrings(),
-                         thrust::next(multilinestring_segment_offsets.begin()));
-
-  thrust::inclusive_scan(
-    rmm::exec_policy(stream),
-    multipolygons.multipolygon_segment_count_begin(),
-    multipolygons.multipolygon_segment_count_begin() + multipolygons.num_multipolygons(),
-    thrust::next(multipolygon_segment_offsets.begin()));
-
-  // Initialize output range
+  auto multipoints = multilinestrings.as_multipoint_range();
+  auto intersects  = point_polygon_intersects(multipoints, multipolygons, stream);
+
+  auto polygons_as_linestrings = multipolygons.as_multilinestring_range();
+
   thrust::fill(rmm::exec_policy(stream),
                distances_first,
-               distances_first + multilinestrings.num_multilinestrings(),
+               distances_first + size,
                std::numeric_limits<T>::max());
 
-  auto num_threads       = thread_bounds.back_element(stream);
-  auto [tpb, num_blocks] = grid_1d(num_threads);
-
-  detail::pairwise_linestring_polygon_distance_kernel<<<num_blocks, tpb, 0, stream.value()>>>(
-    multilinestrings,
-    multipolygons,
-    range{thread_bounds.begin(), thread_bounds.end()},
-    range{multilinestring_segment_offsets.begin(), multilinestring_segment_offsets.end()},
-    range{multipolygon_segment_offsets.begin(), multipolygon_segment_offsets.end()},
-    multipoint_intersects.begin(),
-    distances_first);
+  auto [threads_per_block, num_blocks] = grid_1d(multilinestrings.num_points());
+
+  detail::linestring_distance<<<num_blocks, threads_per_block, 0, stream.value()>>>(
+    multilinestrings, polygons_as_linestrings, intersects.begin(), distances_first);
 
   return distances_first + multilinestrings.num_multilinestrings();
 }

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

@@ -16,68 +16,23 @@
 
 #pragma once
 
-#include <cuspatial/detail/utility/device_atomics.cuh>
-#include <cuspatial/detail/utility/linestring.cuh>
+#include <cuspatial/cuda_utils.hpp>
+#include <cuspatial/detail/kernel/pairwise_distance.cuh>
 #include <cuspatial/error.hpp>
 #include <cuspatial/geometry/vec_2d.hpp>
 #include <cuspatial/traits.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <thrust/advance.h>
-#include <thrust/binary_search.h>
-#include <thrust/distance.h>
-#include <thrust/execution_policy.h>
 #include <thrust/fill.h>
-#include <thrust/memory.h>
 
-#include <iterator>
 #include <limits>
-#include <memory>
 #include <type_traits>
 
 namespace cuspatial {
 namespace detail {
 
-/**
- * @brief Kernel to compute the distance between pairs of point and linestring.
- *
- * 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 MultiPointRange, class MultiLinestringRange, class OutputIterator>
-void __global__ pairwise_point_linestring_distance_kernel(MultiPointRange multipoints,
-                                                          MultiLinestringRange multilinestrings,
-                                                          OutputIterator distances)
-{
-  using T = typename MultiPointRange::element_t;
-
-  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 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 (!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_idx = multilinestrings.geometry_idx_from_part_idx(part_idx);
-
-    // Reduce the minimum distance between different parts of the multi-point.
-    auto [a, b]            = multilinestrings.segment(idx);
-    T min_distance_squared = std::numeric_limits<T>::max();
-
-    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        = min(distance_squared, min_distance_squared);
-    }
-    atomicMin(&distances[geometry_idx], static_cast<T>(sqrt(min_distance_squared)));
-  }
-}
-
 }  // namespace detail
 template <class MultiPointRange, class MultiLinestringRange, class OutputIt>
 OutputIt pairwise_point_linestring_distance(MultiPointRange multipoints,
@@ -96,20 +51,18 @@ OutputIt pairwise_point_linestring_distance(MultiPointRange multipoints,
 
   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,
                  multilinestrings.size(),
                  std::numeric_limits<T>::max());
 
-  std::size_t constexpr threads_per_block = 256;
-  std::size_t const num_blocks =
-    (multilinestrings.size() + threads_per_block - 1) / threads_per_block;
+  auto [threads_per_block, num_blocks] = grid_1d(multilinestrings.num_points());
 
-  detail::
-    pairwise_point_linestring_distance_kernel<<<num_blocks, threads_per_block, 0, stream.value()>>>(
-      multipoints, multilinestrings, distances_first);
+  detail::point_linestring_distance<<<num_blocks, threads_per_block, 0, stream.value()>>>(
+    multipoints, multilinestrings, thrust::nullopt, distances_first);
 
   CUSPATIAL_CUDA_TRY(cudaGetLastError());
 

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

@@ -19,74 +19,21 @@
 #include "distance_utils.cuh"
 
 #include <cuspatial/cuda_utils.hpp>
-#include <cuspatial/detail/algorithm/is_point_in_polygon.cuh>
-#include <cuspatial/detail/utility/device_atomics.cuh>
-#include <cuspatial/detail/utility/linestring.cuh>
-#include <cuspatial/detail/utility/zero_data.cuh>
+#include <cuspatial/detail/kernel/pairwise_distance.cuh>
 #include <cuspatial/error.hpp>
-#include <cuspatial/geometry/vec_2d.hpp>
-#include <cuspatial/iterator_factory.cuh>
-#include <cuspatial/range/range.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
-#include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <thrust/functional.h>
-#include <thrust/iterator/counting_iterator.h>
-#include <thrust/iterator/discard_iterator.h>
-#include <thrust/logical.h>
-#include <thrust/reduce.h>
-#include <thrust/tabulate.h>
-
-#include <cstdint>
 #include <limits>
 #include <type_traits>
 
 namespace cuspatial {
-namespace detail {
-
-/**
- * @brief Kernel to compute the distance between pairs of point and polygon.
- */
-template <class MultiPointRange,
-          class MultiPolygonRange,
-          class IntersectionRange,
-          class OutputIterator>
-void __global__ pairwise_point_polygon_distance_kernel(MultiPointRange multipoints,
-                                                       MultiPolygonRange multipolygons,
-                                                       IntersectionRange intersects,
-                                                       OutputIterator distances)
-{
-  using T = typename MultiPointRange::element_t;
-
-  for (auto idx = threadIdx.x + blockIdx.x * blockDim.x; idx < multipolygons.num_points();
-       idx += gridDim.x * blockDim.x) {
-    auto geometry_idx = multipolygons.geometry_idx_from_segment_idx(idx);
-    if (geometry_idx == MultiPolygonRange::INVALID_INDEX) continue;
-
-    if (intersects[geometry_idx]) {
-      distances[geometry_idx] = T{0.0};
-      continue;
-    }
-
-    auto [a, b] = multipolygons.get_segment(idx);
-
-    T dist_squared = std::numeric_limits<T>::max();
-    for (vec_2d<T> point : multipoints[geometry_idx]) {
-      dist_squared = min(dist_squared, point_to_segment_distance_squared(point, a, b));
-    }
-
-    atomicMin(&distances[geometry_idx], sqrt(dist_squared));
-  }
-}
-
-}  // namespace detail
 
 template <class MultiPointRange, class MultiPolygonRange, class OutputIt>
 OutputIt pairwise_point_polygon_distance(MultiPointRange multipoints,
                                          MultiPolygonRange multipolygons,
-                                         OutputIt distances_first,
+                                         OutputIt distances,
                                          rmm::cuda_stream_view stream)
 {
   using T       = typename MultiPointRange::element_t;
@@ -95,23 +42,25 @@ OutputIt pairwise_point_polygon_distance(MultiPointRange multipoints,
   CUSPATIAL_EXPECTS(multipoints.size() == multipolygons.size(),
                     "Must have the same number of input rows.");
 
-  if (multipoints.size() == 0) return distances_first;
+  if (multipoints.size() == 0) return distances;
+
+  auto intersects = point_polygon_intersects(multipoints, multipolygons, stream);
 
-  auto multipoint_intersects = point_polygon_intersects(multipoints, multipolygons, stream);
+  auto polygons_as_linestrings = multipolygons.as_multilinestring_range();
 
   thrust::fill(rmm::exec_policy(stream),
-               distances_first,
-               distances_first + multipoints.size(),
+               distances,
+               distances + multipoints.size(),
                std::numeric_limits<T>::max());
-  auto [threads_per_block, n_blocks] = grid_1d(multipolygons.num_points());
 
-  detail::
-    pairwise_point_polygon_distance_kernel<<<n_blocks, threads_per_block, 0, stream.value()>>>(
-      multipoints, multipolygons, multipoint_intersects.begin(), distances_first);
+  auto [threads_per_block, n_blocks] = grid_1d(polygons_as_linestrings.num_points());
+
+  detail::point_linestring_distance<<<n_blocks, threads_per_block, 0, stream.value()>>>(
+    multipoints, polygons_as_linestrings, intersects.begin(), distances);
 
   CUSPATIAL_CHECK_CUDA(stream.value());
 
-  return distances_first + multipoints.size();
+  return distances + multipoints.size();
 }
 
 }  // namespace cuspatial