Fix missing binary_pred in thrust::unique_by_key

Fixes: #1956
NVIDIA · Jul 9, 2024 · 5f9c52d · 5f9c52d
1 parent f34a8c1
commit 5f9c52d
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 0 deletions.
diff --git a/thrust/testing/unique_by_key.cu b/thrust/testing/unique_by_key.cu
@@ -470,3 +470,58 @@ struct TestUniqueCopyByKeyLargeOutCount
 SimpleUnitTest<TestUniqueCopyByKeyLargeOutCount, IntegralTypes> TestUniqueCopyByKeyLargeOutCountInstance;
 
 #endif // non-OpenMP backend
+
+// Based on GitHub issue: https://github.com/NVIDIA/cccl/issues/1956
+namespace
+{
+struct Entry
+{
+  std::int32_t a;
+  float b;
+};
+using Key = thrust::pair<std::int32_t, Entry>;
+
+struct Compare
+{
+  _CCCL_HOST_DEVICE bool operator()(Key const& lhs, Key const& rhs) const
+  {
+    return lhs.first == rhs.first;
+  }
+};
+} // namespace
+
+// This test fails only on GCC 6
+#if !defined(__GNUC__) || __GNUC__ != 6
+void TestKeysWithoutEqualityOperator()
+{
+  const auto k1 = Key{1, {}};
+  const auto k2 = Key{2, {}};
+  const thrust::device_vector<Key> keys{k1, k1, k1, k2, k2};
+  thrust::device_vector<Entry> data{{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}};
+
+  thrust::device_vector<Key> unique_keys(5);
+  thrust::device_vector<Entry> unique_data(5);
+
+  const auto result = thrust::unique_by_key_copy(
+    thrust::device, keys.cbegin(), keys.cend(), data.begin(), unique_keys.begin(), unique_data.begin(), Compare{});
+
+  unique_keys.erase(result.first, unique_keys.end());
+  unique_data.erase(result.second, unique_data.end());
+
+  auto unique_keys_h = thrust::host_vector<Key>(unique_keys);
+  auto unique_data_h = thrust::host_vector<Entry>(unique_data);
+
+  ASSERT_EQUAL(unique_keys_h[0].first, k1.first);
+  ASSERT_EQUAL(unique_keys_h[0].second.a, k1.second.a);
+  ASSERT_EQUAL(unique_keys_h[0].second.b, k1.second.b);
+  ASSERT_EQUAL(unique_keys_h[1].first, k2.first);
+  ASSERT_EQUAL(unique_keys_h[1].second.a, k2.second.a);
+  ASSERT_EQUAL(unique_keys_h[1].second.b, k2.second.b);
+
+  ASSERT_EQUAL(unique_data_h[0].a, 0);
+  ASSERT_EQUAL(unique_data_h[0].b, 0);
+  ASSERT_EQUAL(unique_data_h[1].a, 3);
+  ASSERT_EQUAL(unique_data_h[1].b, 3);
+}
+DECLARE_UNITTEST(TestKeysWithoutEqualityOperator);
+#endif // !defined(__GNUC__) || __GNUC__ != 6
diff --git a/thrust/thrust/system/cuda/detail/unique_by_key.h b/thrust/thrust/system/cuda/detail/unique_by_key.h
@@ -121,6 +121,7 @@ struct DispatchUniqueByKey
       values_out,
       static_cast<OffsetT*>(nullptr),
       num_items,
+      binary_pred,
       stream);
     CUDA_CUB_RET_IF_FAIL(status);