[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/thrust/testing/shuffle.cu

@@ -0,0 +1,406 @@
+#include <thrust/detail/config.h>
+
+#include <thrust/gather.h>
+#include <thrust/iterator/shuffle_iterator.h>
+#include <thrust/random.h>
+#include <thrust/scatter.h>
+#include <thrust/sequence.h>
+#include <thrust/shuffle.h>
+#include <thrust/sort.h>
+
+#include <cuda/std/numbers>
+#include <cuda/std/random>
+
+#include <algorithm>
+#include <limits>
+#include <map>
+
+#include <unittest/unittest.h>
+
+struct iterator_shuffle_copy
+{
+  template <typename Iterator, typename ResultIterator>
+  void operator()(Iterator first, Iterator last, ResultIterator result, thrust::default_random_engine& g)
+  {
+    auto shuffle_iter = thrust::make_shuffle_iterator(static_cast<uint64_t>(last - first), g);
+    thrust::gather(shuffle_iter, shuffle_iter + (last - first), first, result);
+  }
+};
+
+struct iterator_shuffle
+{
+  template <typename Iterator>
+  void operator()(Iterator first, Iterator last, thrust::default_random_engine& g)
+  {
+    using thrust::system::detail::generic::select_system;
+    using InputType = typename thrust::detail::it_value_t<Iterator>;
+    using System    = typename thrust::iterator_system<Iterator>::type;
+    System system;
+    auto policy = select_system(system);
+    thrust::detail::temporary_array<InputType, System> temp(policy, first, last);
+    iterator_shuffle_copy{}(temp.begin(), temp.end(), first, g);
+  }
+};
+
+struct thrust_shuffle
+{
+  template <typename Iterator>
+  void operator()(Iterator first, Iterator last, thrust::default_random_engine& g)
+  {
+    thrust::shuffle(first, last, g);
+  }
+};
+
+struct thrust_shuffle_copy
+{
+  template <typename Iterator>
+  void operator()(Iterator first, Iterator last, Iterator result, thrust::default_random_engine& g)
+  {
+    thrust::shuffle_copy(first, last, result, g);
+  }
+};
+
+template <class ShuffleFunc, typename Vector>
+void TestShuffleSimpleBase()
+{
+  Vector data{0, 1, 2, 3, 4};
+  Vector shuffled(data.begin(), data.end());
+  thrust::default_random_engine g(2);
+  ShuffleFunc{}(shuffled.begin(), shuffled.end(), g);
+  thrust::sort(shuffled.begin(), shuffled.end());
+  // Check all of our data is present
+  // This only tests for strange conditions like duplicated elements
+  ASSERT_EQUAL(shuffled, data);
+}
+template <typename Vector>
+void TestShuffleSimple()
+{
+  TestShuffleSimpleBase<thrust_shuffle, Vector>();
+}
+template <typename Vector>
+void TestShuffleSimpleIterator()
+{
+  TestShuffleSimpleBase<iterator_shuffle, Vector>();
+}
+DECLARE_VECTOR_UNITTEST(TestShuffleSimple);
+DECLARE_VECTOR_UNITTEST(TestShuffleSimpleIterator);
+
+template <typename ShuffleFunc, typename ShuffleCopyFunc, typename Vector>
+void TestShuffleCopySimpleBase()
+{
+  Vector data{0, 1, 2, 3, 4};
+  Vector shuffled(5);
+  thrust::default_random_engine g(2);
+  ShuffleCopyFunc{}(data.begin(), data.end(), shuffled.begin(), g);
+  g.seed(2);
+  ShuffleFunc{}(data.begin(), data.end(), g);
+  ASSERT_EQUAL(shuffled, data);
+}
+template <typename Vector>
+void TestShuffleCopySimple()
+{
+  TestShuffleCopySimpleBase<thrust_shuffle, thrust_shuffle_copy, Vector>();
+}
+template <typename Vector>
+void TestShuffleCopySimpleIterator()
+{
+  TestShuffleCopySimpleBase<iterator_shuffle, iterator_shuffle_copy, Vector>();
+}
+DECLARE_VECTOR_UNITTEST(TestShuffleCopySimple);
+DECLARE_VECTOR_UNITTEST(TestShuffleCopySimpleIterator);
+
+template <typename Vector>
+void TestShuffleCudaStdPhilox()
+{
+  Vector data{0, 1, 2, 3, 4};
+  Vector shuffled(data.begin(), data.end());
+  cuda::std::philox4x32 g(2);
+  thrust::shuffle(shuffled.begin(), shuffled.end(), g);
+  thrust::sort(shuffled.begin(), shuffled.end());
+
+  ASSERT_EQUAL(shuffled, data);
+}
+DECLARE_VECTOR_UNITTEST(TestShuffleCudaStdPhilox);
+
+template <typename Vector>
+void TestShuffleCopyCudaStdPhilox()
+{
+  Vector data{0, 1, 2, 3, 4};
+  Vector shuffled(5);
+  Vector in_place(data.begin(), data.end());
+  cuda::std::philox4x32 g(2);
+
+  thrust::shuffle_copy(data.begin(), data.end(), shuffled.begin(), g);
+  g.seed(2);
+  thrust::shuffle(in_place.begin(), in_place.end(), g);
+
+  ASSERT_EQUAL(shuffled, in_place);
+}
+DECLARE_VECTOR_UNITTEST(TestShuffleCopyCudaStdPhilox);
+
+template <typename ShuffleFunc, typename T>
+void TestHostDeviceIdenticalBase(size_t m)
+{
+  thrust::host_vector<T> host_result(m);
+  thrust::device_vector<T> device_result(m);
+  thrust::sequence(host_result.begin(), host_result.end(), T{});
+  thrust::sequence(device_result.begin(), device_result.end(), T{});
+
+  thrust::default_random_engine host_g(183);
+  thrust::default_random_engine device_g(183);
+
+  ShuffleFunc{}(host_result.begin(), host_result.end(), host_g);
+  ShuffleFunc{}(device_result.begin(), device_result.end(), device_g);
+
+  ASSERT_EQUAL(device_result, host_result);
+}
+template <typename T>
+void TestHostDeviceIdentical(size_t m)
+{
+  TestHostDeviceIdenticalBase<thrust_shuffle, T>(m);
+}
+template <typename T>
+void TestHostDeviceIdenticalIterator(size_t m)
+{
+  TestHostDeviceIdenticalBase<iterator_shuffle, T>(m);
+}
+DECLARE_VARIABLE_UNITTEST(TestHostDeviceIdentical);
+DECLARE_VARIABLE_UNITTEST(TestHostDeviceIdenticalIterator);
+
+template <typename BijectionFunc, typename T>
+void TestFunctionIsBijectionBase(size_t m)
+{
+  thrust::default_random_engine device_g(0xD5);
+  BijectionFunc device_f(m, device_g);
+
+  const size_t total_length = device_f.size();
+  if (static_cast<double>(total_length) >= static_cast<double>(std::numeric_limits<T>::max()) || m == 0)
+  {
+    return;
+  }
+  ASSERT_LEQUAL(total_length, (std::max) (m * 2, size_t(256))); // Check the rounded up size is at most double the input
+
+  auto device_result_it = thrust::make_transform_iterator(thrust::make_counting_iterator(T(0)), device_f);
+
+  thrust::device_vector<T> unpermuted(total_length, T(0));
+
+  // Run a scatter, this should copy each value to the index matching is value, the result should be in ascending order
+  thrust::scatter(device_result_it,
+                  device_result_it + static_cast<T>(total_length), // total_length is guaranteed to fit T
+                  device_result_it,
+                  unpermuted.begin());
+
+  // Check every index is in the result, if any are missing then the function was not a bijection over [0,m)
+  ASSERT_EQUAL(true, thrust::equal(unpermuted.begin(), unpermuted.end(), thrust::make_counting_iterator(T(0))));
+}
+template <typename T>
+void TestFunctionIsBijection(size_t m)
+{
+  TestFunctionIsBijectionBase<thrust::detail::feistel_bijection, T>(m);
+}
+template <typename T>
+void TestFunctionIsBijectionIterator(size_t m)
+{
+  TestFunctionIsBijectionBase<thrust::detail::random_bijection<uint64_t>, T>(m);
+}
+DECLARE_INTEGRAL_VARIABLE_UNITTEST(TestFunctionIsBijection);
+DECLARE_INTEGRAL_VARIABLE_UNITTEST(TestFunctionIsBijectionIterator);
+
+void TestFeistelBijectionLength()
+{
+  thrust::default_random_engine g(0xD5);
+
+  uint64_t m = 345;
+  thrust::detail::feistel_bijection f(m, g);
+  ASSERT_EQUAL(f.size(), uint64_t(512));
+
+  m = 256;
+  f = thrust::detail::feistel_bijection(m, g);
+  ASSERT_EQUAL(f.size(), uint64_t(256));
+
+  m = 1;
+  f = thrust::detail::feistel_bijection(m, g);
+  ASSERT_EQUAL(f.size(), uint64_t(256));
+}
+DECLARE_UNITTEST(TestFeistelBijectionLength);
+
+void TestShuffleIteratorConstructibleFromBijection()
+{
+  thrust::default_random_engine g(0xD5);
+
+  thrust::detail::feistel_bijection f(32, g);
+  thrust::shuffle_iterator<uint64_t, decltype(f)> it(f);
+
+  g.seed(0xD5);
+  thrust::detail::random_bijection<uint64_t> f2(f.size(), g);
+  thrust::shuffle_iterator<uint64_t, decltype(f2)> it2(f2);
+
+  g.seed(0xD5);
+  thrust::shuffle_iterator<uint64_t> it3(f.size(), g);
+
+  ASSERT_EQUAL(true, thrust::equal(thrust::device, it, it + f.size(), it2));
+  ASSERT_EQUAL(true, thrust::equal(thrust::device, it, it + f.size(), it3));
+}
+DECLARE_UNITTEST(TestShuffleIteratorConstructibleFromBijection);
+
+void TestShuffleAndPermutationIterator()
+{
+  thrust::default_random_engine g(0xD5);
+
+  auto it = thrust::make_shuffle_iterator(32, g);
+
+  thrust::device_vector<uint64_t> data(32);
+  thrust::sequence(data.begin(), data.end(), 0);
+
+  auto permute_it = thrust::make_permutation_iterator(data.begin(), it);
+
+  thrust::device_vector<uint64_t> premute_vec(32);
+  thrust::gather(it, it + 32, data.begin(), premute_vec.begin());
+
+  ASSERT_EQUAL(true, thrust::equal(permute_it, permute_it + 32, premute_vec.begin()));
+}
+DECLARE_UNITTEST(TestShuffleAndPermutationIterator);
+
+void TestShuffleIteratorStateless()
+{
+  thrust::default_random_engine g(0xD5);
+
+  auto it = thrust::make_shuffle_iterator(32, g);
+
+  ASSERT_EQUAL(*it, *it);
+  ASSERT_EQUAL(*(it + 1), *(it + 1));
+  ++it;
+  ASSERT_EQUAL(*(it - 1), *(it - 1));
+}
+DECLARE_UNITTEST(TestShuffleIteratorStateless);
+
+double inverse_erf(double x)
+{
+  double tt1, tt2, lnx, sgn;
+  sgn = (x < 0) ? -1.0 : 1.0;
+
+  x   = (1 - x) * (1 + x);
+  lnx = cuda::std::log(x);
+
+  tt1 = 2 / (cuda::std::__numbers<double>::__pi() * 0.147) + 0.5f * lnx;
+  tt2 = 1 / (0.147) * lnx;
+
+  return (sgn * cuda::std::sqrt(-tt1 + cuda::std::sqrt(tt1 * tt1 - tt2)));
+}
+
+// Individual input keys should be permuted to output locations with uniform
+// probability. Perform chi-squared test with confidence 95%.
+template <typename ShuffleFunc, typename Vector>
+void TestShuffleKeyPositionBase()
+{
+  using T               = typename Vector::value_type;
+  const int num_samples = 1000;
+  const int n           = 20;
+  thrust::default_random_engine g(0xD5);
+  thrust::host_vector<double> expected_value(n);
+  thrust::host_vector<T> sequence(n);
+  thrust::sequence(sequence.begin(), sequence.end(), T(0));
+  for (size_t i = 0; i < num_samples; ++i)
+  {
+    Vector shuffled(sequence.begin(), sequence.end());
+    ShuffleFunc{}(shuffled.begin(), shuffled.end(), g);
+    thrust::host_vector<T> tmp(shuffled.begin(), shuffled.end());
+    for (size_t j = 0; j < n; ++j)
+    {
+      expected_value[j] += static_cast<double>(tmp[j]);
+    }
+  }
+
+  double mu    = (n - 1) / 2.0;
+  double sigma = cuda::std::sqrt((n * n - 1) / 12.0);
+  double zmax  = 0.0;
+  for (size_t i = 0; i < n; ++i)
+  {
+    double mean = expected_value[i] / double(num_samples);
+    double z    = cuda::std::abs(mean - mu) / (sigma / cuda::std::sqrt(double(num_samples)));
+    zmax        = cuda::std::max(zmax, cuda::std::abs(z));
+  }
+
+  double alpha = 0.05;
+  double zcrit = inverse_erf(1.0 - alpha / (2.0 * n)) * cuda::std::sqrt(2.0);
+  ASSERT_LESS(zmax, zcrit);
+}
+template <typename Vector>
+void TestShuffleKeyPosition()
+{
+  TestShuffleKeyPositionBase<thrust_shuffle, Vector>();
+}
+template <typename Vector>
+void TestShuffleKeyPositionIterator()
+{
+  TestShuffleKeyPositionBase<iterator_shuffle, Vector>();
+}
+DECLARE_INTEGRAL_VECTOR_UNITTEST(TestShuffleKeyPosition);
+DECLARE_INTEGRAL_VECTOR_UNITTEST(TestShuffleKeyPositionIterator);
+
+struct vector_compare
+{
+  template <typename VectorT>
+  bool operator()(const VectorT& a, const VectorT& b) const
+  {
+    for (auto i = 0ull; i < a.size(); i++)
+    {
+      if (a[i] < b[i])
+      {
+        return true;
+      }
+      if (a[i] > b[i])
+      {
+        return false;
+      }
+    }
+    return false;
+  }
+};
+
+// Brute force check permutations are uniformly distributed on small input
+// Uses a chi-squared test indicating 99% confidence the output is uniformly
+// random
+template <typename ShuffleFunc, typename Vector>
+void TestShuffleUniformPermutationBase()
+{
+  using T = typename Vector::value_type;
+
+  size_t m                  = 5;
+  size_t num_samples        = 1000;
+  size_t total_permutations = 1 * 2 * 3 * 4 * 5;
+  std::map<thrust::host_vector<T>, size_t, vector_compare> permutation_counts;
+  Vector sequence(m);
+  thrust::sequence(sequence.begin(), sequence.end(), T(0));
+  thrust::default_random_engine g(0xD5);
+  for (auto i = 0ull; i < num_samples; i++)
+  {
+    ShuffleFunc{}(sequence.begin(), sequence.end(), g);
+    thrust::host_vector<T> tmp(sequence.begin(), sequence.end());
+    permutation_counts[tmp]++;
+  }
+
+  ASSERT_EQUAL(permutation_counts.size(), total_permutations);
+
+  double chi_squared    = 0.0;
+  double expected_count = static_cast<double>(num_samples) / static_cast<double>(total_permutations);
+  for (const auto& kv : permutation_counts)
+  {
+    chi_squared += std::pow(expected_count - kv.second, 2) / expected_count;
+  }
+  // 119 degrees of freedom, 95% confidence
+  const double critical_value = 145.461;
+  ASSERT_LESS(chi_squared, critical_value);
+}
+template <typename Vector>
+void TestShuffleUniformPermutation()
+{
+  TestShuffleUniformPermutationBase<thrust_shuffle, Vector>();
+}
+template <typename Vector>
+void TestShuffleUniformPermutationIterator()
+{
+  TestShuffleUniformPermutationBase<iterator_shuffle, Vector>();
+}
+DECLARE_VECTOR_UNITTEST(TestShuffleUniformPermutation);
+DECLARE_VECTOR_UNITTEST(TestShuffleUniformPermutationIterator);