Improve Hausdorff perf and accept larger number of inputs.

cpp/src/spatial/detail/cartesian_product_group_index_iterator.cuh

@@ -63,17 +63,18 @@ struct cartesian_product_group_index_functor {
     auto const group_idx_a            = *(group_lookup_a + group_lookup_idx_a);
     auto const group_size_a           = *(group_sizes_a + group_idx_a);
     auto const group_offset_a         = *(group_offsets_a + group_idx_a);
-    auto const group_block_offset_a   = group_offset_a * num_elements_b;
+    auto const group_block_offset_a   = static_cast<uint64_t>(group_offset_a) * num_elements_b;
 
     // second dimension
     uint32_t const group_lookup_idx_b = (idx - group_block_offset_a) / group_size_a;
     auto const group_idx_b            = *(group_lookup_b + group_lookup_idx_b);
     auto const group_size_b           = *(group_sizes_b + group_idx_b);
     auto const group_offset_b         = *(group_offsets_b + group_idx_b);
-    auto const group_block_offset_b   = group_offset_b * group_size_a;
+    auto const group_block_offset_b   = static_cast<uint64_t>(group_offset_b) * group_size_a;
 
     // relative index
-    uint32_t const relative_idx       = idx - group_block_offset_a - group_block_offset_b;
+    auto const relative_idx =
+      static_cast<uint32_t>(idx - group_block_offset_a - group_block_offset_b);
     auto const relative_element_idx_a = relative_idx % group_size_a;
     auto const relative_element_idx_b = relative_idx / group_size_a;
 

cpp/src/spatial/hausdorff.cu

@@ -121,38 +121,41 @@ struct hausdorff_functor {
 
     auto scatter_out = make_scatter_output_iterator(result_temp.begin(), scatter_map);
 
-    auto gpc_key_iter = thrust::make_transform_iterator(
+    auto gcp_key_iter = thrust::make_transform_iterator(
       gcp_iter, [] __device__(cartesian_product_group_index const& idx) {
         return thrust::make_pair(idx.group_a.idx, idx.group_b.idx);
       });
 
     // the following output iterator and `inclusive_scan_by_key` could be replaced by a
-    // reduce_by_key, if it supported non-commutative operators.
+    // reduce_by_key, if it supported non-commutative operators and more inputs (int 64).
 
-    auto const num_cartesian =
-      static_cast<uint64_t>(num_points) * static_cast<uint64_t>(num_points);
+    // copy offsets to host for batching purposes.
 
-    //
-    // `thrust::inclusive_scan_by_key` causes an out-of-memory
-    // error on input sizes close to (but not exactly) INT_MAX.
-    //
-    // Doing the reduction in chunks is a temporary workaround until this is fixed.
-    //
-    // The magic number between OOM vs. no OOM is somewhere between:
-    //                                             (1uL << 31uL) - 2048uL;
-    auto const magic_inclusive_scan_by_key_limit = (1uL << 31uL) - 4096uL;
+    thrust::host_vector<uint32_t> h_space_offsets(space_offsets.size());
 
-    for (auto itr = gpc_key_iter, end = gpc_key_iter + num_cartesian; itr < end;) {
-      auto const len = static_cast<uint32_t>(std::min(
-        static_cast<uint64_t>(thrust::distance(itr, end)), magic_inclusive_scan_by_key_limit));
+    CUDA_TRY(cudaMemcpyAsync(h_space_offsets.data(),
+                             space_offsets.data<uint32_t>(),
+                             space_offsets.size() * sizeof(uint32_t),
+                             cudaMemcpyDeviceToHost,
+                             stream.value()));
+
+    stream.synchronize();
+
+    for (uint32_t i = 1; i <= h_space_offsets.size(); i++) {
+      uint64_t space_size =
+        (i < h_space_offsets.size() ? h_space_offsets[i] : xs.size()) - h_space_offsets[i - 1];
+      uint64_t elements_in_batch = xs.size() * space_size;
 
       thrust::inclusive_scan_by_key(rmm::exec_policy(stream),
-                                    itr,
-                                    itr + len,
+                                    gcp_key_iter,
+                                    gcp_key_iter + elements_in_batch,
                                     hausdorff_acc_iter,
                                     scatter_out,
-                                    thrust::equal_to<thrust::pair<uint32_t, uint32_t>>());
-      thrust::advance(itr, len);
+                                    thrust::equal_to<thrust::pair<int32_t, int32_t>>());
+
+      hausdorff_acc_iter += elements_in_batch;
+      gcp_key_iter += elements_in_batch;
+      scatter_out += elements_in_batch;
     }
 
     thrust::transform(rmm::exec_policy(stream),

python/cuspatial/cuspatial/core/gis.py

@@ -87,7 +87,7 @@ def directed_hausdorff_distance(xs, ys, space_offsets):
         return DataFrame()
     xs, ys = normalize_point_columns(as_column(xs), as_column(ys))
     result = cpp_directed_hausdorff_distance(
-        xs, ys, as_column(space_offsets, dtype="int32"),
+        xs, ys, as_column(space_offsets, dtype="uint32"),
     )
     result = result.data_array_view
     result = result.reshape(num_spaces, num_spaces)