Skip to content

Commit

Permalink
Clean up thrust merge unit tests (#1819)
Browse files Browse the repository at this point in the history
  • Loading branch information
bernhardmgruber authored Jun 7, 2024
1 parent 2d011cc commit 998e020
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 76 deletions.
45 changes: 14 additions & 31 deletions thrust/testing/cuda/merge.cu
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,11 @@ __global__ void merge_kernel(
template <typename ExecutionPolicy>
void TestMergeDevice(ExecutionPolicy exec)
{
size_t n = 10000;
size_t sizes[] = {0, 1, n / 2, n, n + 1, 2 * n};
size_t num_sizes = sizeof(sizes) / sizeof(size_t);
const size_t n = 10000;
const size_t sizes[] = {0, 1, n / 2, n, n + 1, 2 * n};
const size_t num_sizes = sizeof(sizes) / sizeof(size_t);

thrust::host_vector<int> random =
const thrust::host_vector<int> random =
unittest::random_integers<unittest::int8_t>(n + *thrust::max_element(sizes, sizes + num_sizes));

thrust::host_vector<int> h_a(random.begin(), random.begin() + n);
Expand All @@ -36,30 +36,28 @@ void TestMergeDevice(ExecutionPolicy exec)
thrust::stable_sort(h_a.begin(), h_a.end());
thrust::stable_sort(h_b.begin(), h_b.end());

thrust::device_vector<int> d_a = h_a;
thrust::device_vector<int> d_b = h_b;
const thrust::device_vector<int> d_a = h_a;
const thrust::device_vector<int> d_b = h_b;

for (size_t i = 0; i < num_sizes; i++)
{
size_t size = sizes[i];
const size_t size = sizes[i];

thrust::host_vector<int> h_result(n + size);
thrust::device_vector<int> d_result(n + size);

typename thrust::host_vector<int>::iterator h_end;

typedef typename thrust::device_vector<int>::iterator iter_type;
using iter_type = typename thrust::device_vector<int>::iterator;
thrust::device_vector<iter_type> d_end(1);

h_end = thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.begin() + size, h_result.begin());
const auto h_end = thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.begin() + size, h_result.begin());
h_result.resize(h_end - h_result.begin());

merge_kernel<<<1, 1>>>(
exec, d_a.begin(), d_a.end(), d_b.begin(), d_b.begin() + size, d_result.begin(), d_end.begin());
cudaError_t const err = cudaDeviceSynchronize();
ASSERT_EQUAL(cudaSuccess, err);

d_result.resize((iter_type) d_end[0] - d_result.begin());
d_result.resize(static_cast<iter_type>(d_end[0]) - d_result.begin());

ASSERT_EQUAL(h_result, d_result);
}
Expand All @@ -80,31 +78,16 @@ DECLARE_UNITTEST(TestMergeDeviceDevice);

void TestMergeCudaStreams()
{
typedef thrust::device_vector<int> Vector;
typedef Vector::iterator Iterator;

Vector a(3), b(4);

// clang-format off
a[0] = 0; a[1] = 2; a[2] = 4;
b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;
// clang-format on

Vector ref(7);
ref[0] = 0;
ref[1] = 0;
ref[2] = 2;
ref[3] = 3;
ref[4] = 3;
ref[5] = 4;
ref[6] = 4;
using Vector = thrust::device_vector<int>;
const Vector a{0, 2, 4}, b{0, 3, 3, 4};
const Vector ref{0, 0, 2, 3, 3, 4, 4};

Vector result(7);

cudaStream_t s;
cudaStreamCreate(&s);

Iterator end = thrust::merge(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin(), b.end(), result.begin());
const auto end = thrust::merge(thrust::cuda::par.on(s), a.begin(), a.end(), b.begin(), b.end(), result.begin());

ASSERT_EQUAL_QUIET(result.end(), end);
ASSERT_EQUAL(ref, result);
Expand Down
68 changes: 23 additions & 45 deletions thrust/testing/merge.cu
Original file line number Diff line number Diff line change
Expand Up @@ -11,27 +11,11 @@
template <typename Vector>
void TestMergeSimple()
{
typedef typename Vector::iterator Iterator;

Vector a(3), b(4);

// clang-format off
a[0] = 0; a[1] = 2; a[2] = 4;
b[0] = 0; b[1] = 3; b[2] = 3; b[3] = 4;
// clang-format on

Vector ref(7);
ref[0] = 0;
ref[1] = 0;
ref[2] = 2;
ref[3] = 3;
ref[4] = 3;
ref[5] = 4;
ref[6] = 4;
const Vector a{0, 2, 4}, b{0, 3, 3, 4};
const Vector ref{0, 0, 2, 3, 3, 4, 4};

Vector result(7);

Iterator end = thrust::merge(a.begin(), a.end(), b.begin(), b.end(), result.begin());
const auto end = thrust::merge(a.begin(), a.end(), b.begin(), b.end(), result.begin());

ASSERT_EQUAL_QUIET(result.end(), end);
ASSERT_EQUAL(ref, result);
Expand Down Expand Up @@ -81,10 +65,10 @@ DECLARE_UNITTEST(TestMergeDispatchImplicit);
template <typename T>
void TestMerge(size_t n)
{
size_t sizes[] = {0, 1, n / 2, n, n + 1, 2 * n};
size_t num_sizes = sizeof(sizes) / sizeof(size_t);
const size_t sizes[] = {0, 1, n / 2, n, n + 1, 2 * n};
const size_t num_sizes = sizeof(sizes) / sizeof(size_t);

thrust::host_vector<T> random =
const thrust::host_vector<T> random =
unittest::random_integers<unittest::int8_t>(n + *thrust::max_element(sizes, sizes + num_sizes));

thrust::host_vector<T> h_a(random.begin(), random.begin() + n);
Expand All @@ -93,23 +77,20 @@ void TestMerge(size_t n)
thrust::stable_sort(h_a.begin(), h_a.end());
thrust::stable_sort(h_b.begin(), h_b.end());

thrust::device_vector<T> d_a = h_a;
thrust::device_vector<T> d_b = h_b;
const thrust::device_vector<T> d_a = h_a;
const thrust::device_vector<T> d_b = h_b;

for (size_t i = 0; i < num_sizes; i++)
{
size_t size = sizes[i];
const size_t size = sizes[i];

thrust::host_vector<T> h_result(n + size);
thrust::device_vector<T> d_result(n + size);

typename thrust::host_vector<T>::iterator h_end;
typename thrust::device_vector<T>::iterator d_end;
const auto h_end = thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.begin() + size, h_result.begin());
const auto d_end = thrust::merge(d_a.begin(), d_a.end(), d_b.begin(), d_b.begin() + size, d_result.begin());

h_end = thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.begin() + size, h_result.begin());
h_result.resize(h_end - h_result.begin());

d_end = thrust::merge(d_a.begin(), d_a.end(), d_b.begin(), d_b.begin() + size, d_result.begin());
d_result.resize(d_end - d_result.begin());

ASSERT_EQUAL(h_result, d_result);
Expand All @@ -126,14 +107,11 @@ void TestMergeToDiscardIterator(size_t n)
thrust::stable_sort(h_a.begin(), h_a.end());
thrust::stable_sort(h_b.begin(), h_b.end());

thrust::device_vector<T> d_a = h_a;
thrust::device_vector<T> d_b = h_b;

thrust::discard_iterator<> h_result =
thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.end(), thrust::make_discard_iterator());
const thrust::device_vector<T> d_a = h_a;
const thrust::device_vector<T> d_b = h_b;

thrust::discard_iterator<> d_result =
thrust::merge(d_a.begin(), d_a.end(), d_b.begin(), d_b.end(), thrust::make_discard_iterator());
const auto h_result = thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.end(), thrust::make_discard_iterator());
const auto d_result = thrust::merge(d_a.begin(), d_a.end(), d_b.begin(), d_b.end(), thrust::make_discard_iterator());

thrust::discard_iterator<> reference(2 * n);

Expand All @@ -151,19 +129,19 @@ void TestMergeDescending(size_t n)
thrust::stable_sort(h_a.begin(), h_a.end(), thrust::greater<T>());
thrust::stable_sort(h_b.begin(), h_b.end(), thrust::greater<T>());

thrust::device_vector<T> d_a = h_a;
thrust::device_vector<T> d_b = h_b;
const thrust::device_vector<T> d_a = h_a;
const thrust::device_vector<T> d_b = h_b;

thrust::host_vector<T> h_result(h_a.size() + h_b.size());
thrust::device_vector<T> d_result(d_a.size() + d_b.size());

typename thrust::host_vector<T>::iterator h_end;
typename thrust::device_vector<T>::iterator d_end;

h_end = thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.end(), h_result.begin(), thrust::greater<T>());

d_end = thrust::merge(d_a.begin(), d_a.end(), d_b.begin(), d_b.end(), d_result.begin(), thrust::greater<T>());
const auto h_end =
thrust::merge(h_a.begin(), h_a.end(), h_b.begin(), h_b.end(), h_result.begin(), thrust::greater<T>());
const auto d_end =
thrust::merge(d_a.begin(), d_a.end(), d_b.begin(), d_b.end(), d_result.begin(), thrust::greater<T>());

ASSERT_EQUAL(h_result, d_result);
ASSERT_EQUAL(h_end == h_result.end(), true);
ASSERT_EQUAL(d_end == d_result.end(), true);
}
DECLARE_VARIABLE_UNITTEST(TestMergeDescending);

0 comments on commit 998e020

Please sign in to comment.