Fix scalar binary search when using in kernels, add tests
The old implementation (from generic/binary_search.inl) uses temporary_array
which requires device-side memory allocation, which is currently
under development in HIP.
Add tests for both scalar and vector functions.
MODIFY
test/test_binary_search.cpp
test/test_binary_search.cpp
*/
#include <thrust/binary_search.h>
+#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/retag.h>
#include <thrust/sequence.h>
#include <thrust/sort.h>
#include "test_header.hpp"
-TESTS_DEFINE(BinarySearchTests, FullTestsParams);
+TESTS_DEFINE(BinarySearchTestsInKernel, NumericalTestsParams);
+
+struct custom_less
+{
+ template<class T>
+ __device__ inline
+ bool operator()(T a, T b)
+ {
+ return a < b;
+ }
+};
+
+template<class T>
+__global__
+void lower_bound_kernel(size_t n,
+ T* input,
+ ptrdiff_t* output)
+{
+ output[0] = thrust::lower_bound(thrust::device, input, input + n, T(0), custom_less()) - input;
+ output[1] = thrust::lower_bound(thrust::device, input, input + n, T(1)) - input;
+ output[2] = thrust::lower_bound(thrust::device, input, input + n, T(2)) - input;
+ output[3] = thrust::lower_bound(thrust::device, input, input + n, T(3)) - input;
+ output[4] = thrust::lower_bound(thrust::device, input, input + n, T(4), custom_less()) - input;
+ output[5] = thrust::lower_bound(thrust::device, input, input + n, T(5)) - input;
+ output[6] = thrust::lower_bound(thrust::device, input, input + n, T(6)) - input;
+ output[7] = thrust::lower_bound(thrust::device, input, input + n, T(7)) - input;
+ output[8] = thrust::lower_bound(thrust::device, input, input + n, T(8)) - input;
+ output[9] = thrust::lower_bound(thrust::device, input, input + n, T(9), custom_less()) - input;
+}
-// TODO: Check the compiler error
-//#if THRUST_DEVICE_COMPILER == THRUST_DEVICE_COMPILER_HCC
-//__global__ void kern(int N,
-// float* inp,
-// float val)
-//{
-// auto low_bound = thrust::lower_bound(thrust::device, inp, inp + N, val);
-//}
-//
-//TEST(BinarySearchTests, TestLoweBoundInKernel)
-//{
-// float * input_device;
-// hipMalloc(&input_device, sizeof(float)*6);
-// float value = 5;
-//
-// hipLaunchKernelGGL(HIP_KERNEL_NAME(kern), dim3(1), dim3(1), 0,0,6, input_device, value);
-//
-// hipFree(input_device);
-//}
-//#endif // THRUST_DEVICE_COMPILER == THRUST_DEVICE_COMPILER_HCC
+TYPED_TEST(BinarySearchTestsInKernel, TestLowerBound)
+{
+ using T = typename TestFixture::input_type;
+
+ thrust::device_vector<T> d_input(5);
+ d_input[0] = 0;
+ d_input[1] = 2;
+ d_input[2] = 5;
+ d_input[3] = 7;
+ d_input[4] = 8;
+
+ thrust::device_vector<ptrdiff_t> d_output(10);
+
+ hipLaunchKernelGGL(
+ HIP_KERNEL_NAME(lower_bound_kernel),
+ dim3(1), dim3(1), 0, 0,
+ size_t(d_input.size()),
+ thrust::raw_pointer_cast(d_input.data()),
+ thrust::raw_pointer_cast(d_output.data())
+ );
+
+ thrust::host_vector<ptrdiff_t> output = d_output;
+ ASSERT_EQ(output[0], 0);
+ ASSERT_EQ(output[1], 1);
+ ASSERT_EQ(output[2], 1);
+ ASSERT_EQ(output[3], 2);
+ ASSERT_EQ(output[4], 2);
+ ASSERT_EQ(output[5], 2);
+ ASSERT_EQ(output[6], 3);
+ ASSERT_EQ(output[7], 3);
+ ASSERT_EQ(output[8], 4);
+ ASSERT_EQ(output[9], 5);
+}
+template<class T>
+__global__
+void upper_bound_kernel(size_t n,
+ T* input,
+ ptrdiff_t* output)
+{
+ output[0] = thrust::upper_bound(thrust::device, input, input + n, T(0)) - input;
+ output[1] = thrust::upper_bound(thrust::device, input, input + n, T(1)) - input;
+ output[2] = thrust::upper_bound(thrust::device, input, input + n, T(2)) - input;
+ output[3] = thrust::upper_bound(thrust::device, input, input + n, T(3)) - input;
+ output[4] = thrust::upper_bound(thrust::device, input, input + n, T(4)) - input;
+ output[5] = thrust::upper_bound(thrust::device, input, input + n, T(5)) - input;
+ output[6] = thrust::upper_bound(thrust::device, input, input + n, T(6)) - input;
+ output[7] = thrust::upper_bound(thrust::device, input, input + n, T(7)) - input;
+ output[8] = thrust::upper_bound(thrust::device, input, input + n, T(8)) - input;
+ output[9] = thrust::upper_bound(thrust::device, input, input + n, T(9)) - input;
+}
+
+TYPED_TEST(BinarySearchTestsInKernel, TestUpperBound)
+{
+ using T = typename TestFixture::input_type;
+
+ thrust::device_vector<T> d_input(5);
+ d_input[0] = 0;
+ d_input[1] = 2;
+ d_input[2] = 5;
+ d_input[3] = 7;
+ d_input[4] = 8;
+
+ thrust::device_vector<ptrdiff_t> d_output(10);
+
+ hipLaunchKernelGGL(
+ HIP_KERNEL_NAME(upper_bound_kernel),
+ dim3(1), dim3(1), 0, 0,
+ size_t(d_input.size()),
+ thrust::raw_pointer_cast(d_input.data()),
+ thrust::raw_pointer_cast(d_output.data())
+ );
+
+ thrust::host_vector<ptrdiff_t> output = d_output;
+ ASSERT_EQ(output[0], 1);
+ ASSERT_EQ(output[1], 1);
+ ASSERT_EQ(output[2], 2);
+ ASSERT_EQ(output[3], 2);
+ ASSERT_EQ(output[4], 2);
+ ASSERT_EQ(output[5], 3);
+ ASSERT_EQ(output[6], 3);
+ ASSERT_EQ(output[7], 4);
+ ASSERT_EQ(output[8], 5);
+ ASSERT_EQ(output[9], 5);
+}
+
+template<class T>
+__global__
+void binary_search_kernel(size_t n,
+ T* input,
+ bool* output)
+{
+ output[0] = thrust::binary_search(thrust::device, input, input + n, T(0));
+ output[1] = thrust::binary_search(thrust::device, input, input + n, T(1));
+ output[2] = thrust::binary_search(thrust::device, input, input + n, T(2));
+ output[3] = thrust::binary_search(thrust::device, input, input + n, T(3));
+ output[4] = thrust::binary_search(thrust::device, input, input + n, T(4));
+ output[5] = thrust::binary_search(thrust::device, input, input + n, T(5));
+ output[6] = thrust::binary_search(thrust::device, input, input + n, T(6));
+ output[7] = thrust::binary_search(thrust::device, input, input + n, T(7));
+ output[8] = thrust::binary_search(thrust::device, input, input + n, T(8));
+ output[9] = thrust::binary_search(thrust::device, input, input + n, T(9));
+}
+
+TYPED_TEST(BinarySearchTestsInKernel, TestBinarySearch)
+{
+ using T = typename TestFixture::input_type;
+
+ thrust::device_vector<T> d_input(5);
+ d_input[0] = 0;
+ d_input[1] = 2;
+ d_input[2] = 5;
+ d_input[3] = 7;
+ d_input[4] = 8;
+
+ thrust::device_vector<bool> d_output(10);
+
+ hipLaunchKernelGGL(
+ HIP_KERNEL_NAME(binary_search_kernel),
+ dim3(1), dim3(1), 0, 0,
+ size_t(d_input.size()),
+ thrust::raw_pointer_cast(d_input.data()),
+ thrust::raw_pointer_cast(d_output.data())
+ );
+
+ thrust::host_vector<bool> output = d_output;
+ ASSERT_EQ(output[0], true);
+ ASSERT_EQ(output[1], false);
+ ASSERT_EQ(output[2], true);
+ ASSERT_EQ(output[3], false);
+ ASSERT_EQ(output[4], false);
+ ASSERT_EQ(output[5], true);
+ ASSERT_EQ(output[6], false);
+ ASSERT_EQ(output[7], true);
+ ASSERT_EQ(output[8], true);
+ ASSERT_EQ(output[9], false);
+}
+
+TESTS_DEFINE(BinarySearchTests, FullTestsParams);
__THRUST_DISABLE_MSVC_POSSIBLE_LOSS_OF_DATA_WARNING_BEGIN

MODIFY
test/test_binary_search_vector.cpp
test/test_binary_search_vector.cpp
#include <thrust/binary_search.h>
#include <thrust/iterator/discard_iterator.h>
+#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/retag.h>
#include <thrust/sequence.h>
#include <thrust/sort.h>
#include "test_header.hpp"
+TESTS_DEFINE(BinarySearchVectorTestsInKernel, NumericalTestsParams);
+
+struct custom_less
+{
+ template<class T>
+ __device__ inline
+ bool operator()(T a, T b)
+ {
+ return a < b;
+ }
+};
+
+template<class T>
+__global__
+void lower_bound_kernel(size_t n,
+ T* input,
+ ptrdiff_t* output)
+{
+ thrust::counting_iterator<T> values(0);
+ thrust::lower_bound(
+ thrust::device, input, input + n, values, values + 10, output, custom_less()
+ );
+}
+
+TYPED_TEST(BinarySearchVectorTestsInKernel, TestLowerBound)
+{
+ using T = typename TestFixture::input_type;
+
+ thrust::device_vector<T> d_input(5);
+ d_input[0] = 0;
+ d_input[1] = 2;
+ d_input[2] = 5;
+ d_input[3] = 7;
+ d_input[4] = 8;
+
+ thrust::device_vector<ptrdiff_t> d_output(10);
+
+ hipLaunchKernelGGL(
+ HIP_KERNEL_NAME(lower_bound_kernel),
+ dim3(1), dim3(1), 0, 0,
+ size_t(d_input.size()),
+ thrust::raw_pointer_cast(d_input.data()),
+ thrust::raw_pointer_cast(d_output.data())
+ );
+
+ thrust::host_vector<ptrdiff_t> output = d_output;
+ ASSERT_EQ(output[0], 0);
+ ASSERT_EQ(output[1], 1);
+ ASSERT_EQ(output[2], 1);
+ ASSERT_EQ(output[3], 2);
+ ASSERT_EQ(output[4], 2);
+ ASSERT_EQ(output[5], 2);
+ ASSERT_EQ(output[6], 3);
+ ASSERT_EQ(output[7], 3);
+ ASSERT_EQ(output[8], 4);
+ ASSERT_EQ(output[9], 5);
+}
+
+template<class T>
+__global__
+void upper_bound_kernel(size_t n,
+ T* input,
+ ptrdiff_t* output)
+{
+ thrust::counting_iterator<T> values(0);
+ thrust::upper_bound(thrust::device, input, input + n, values, values + 10, output);
+}
+
+TYPED_TEST(BinarySearchVectorTestsInKernel, TestUpperBound)
+{
+ using T = typename TestFixture::input_type;
+
+ thrust::device_vector<T> d_input(5);
+ d_input[0] = 0;
+ d_input[1] = 2;
+ d_input[2] = 5;
+ d_input[3] = 7;
+ d_input[4] = 8;
+
+ thrust::device_vector<ptrdiff_t> d_output(10);
+
+ hipLaunchKernelGGL(
+ HIP_KERNEL_NAME(upper_bound_kernel),
+ dim3(1), dim3(1), 0, 0,
+ size_t(d_input.size()),
+ thrust::raw_pointer_cast(d_input.data()),
+ thrust::raw_pointer_cast(d_output.data())
+ );
+
+ thrust::host_vector<ptrdiff_t> output = d_output;
+ ASSERT_EQ(output[0], 1);
+ ASSERT_EQ(output[1], 1);
+ ASSERT_EQ(output[2], 2);
+ ASSERT_EQ(output[3], 2);
+ ASSERT_EQ(output[4], 2);
+ ASSERT_EQ(output[5], 3);
+ ASSERT_EQ(output[6], 3);
+ ASSERT_EQ(output[7], 4);
+ ASSERT_EQ(output[8], 5);
+ ASSERT_EQ(output[9], 5);
+}
+
+template<class T>
+__global__
+void binary_search_kernel(size_t n,
+ T* input,
+ bool* output)
+{
+ thrust::counting_iterator<T> values(0);
+ thrust::binary_search(thrust::device, input, input + n, values, values + 10, output);
+}
+
+TYPED_TEST(BinarySearchVectorTestsInKernel, TestBinarySearch)
+{
+ using T = typename TestFixture::input_type;
+
+ thrust::device_vector<T> d_input(5);
+ d_input[0] = 0;
+ d_input[1] = 2;
+ d_input[2] = 5;
+ d_input[3] = 7;
+ d_input[4] = 8;
+
+ thrust::device_vector<bool> d_output(10);
+
+ hipLaunchKernelGGL(
+ HIP_KERNEL_NAME(binary_search_kernel),
+ dim3(1), dim3(1), 0, 0,
+ size_t(d_input.size()),
+ thrust::raw_pointer_cast(d_input.data()),
+ thrust::raw_pointer_cast(d_output.data())
+ );
+
+ thrust::host_vector<bool> output = d_output;
+ ASSERT_EQ(output[0], true);
+ ASSERT_EQ(output[1], false);
+ ASSERT_EQ(output[2], true);
+ ASSERT_EQ(output[3], false);
+ ASSERT_EQ(output[4], false);
+ ASSERT_EQ(output[5], true);
+ ASSERT_EQ(output[6], false);
+ ASSERT_EQ(output[7], true);
+ ASSERT_EQ(output[8], true);
+ ASSERT_EQ(output[9], false);
+}
+
TESTS_DEFINE(BinarySearchVectorTests, FullTestsParams);
TESTS_DEFINE(BinarySearchVectorIntegerTests, SignedIntegerTestsParams);

MODIFY
thrust/system/hip/detail/binary_search.h
thrust/system/hip/detail/binary_search.h
@@ -226,6 +226,8 @@ namespace __binary_search
// Thrust API entry points
//-------------------------
+// Vector functions
+
template <class Derived, class HaystackIt, class NeedlesIt, class OutputIt, class CompareOp>
OutputIt THRUST_HIP_FUNCTION
lower_bound(execution_policy<Derived>& policy,
@@ -236,6 +238,9 @@ lower_bound(execution_policy<Derived>& policy,
OutputIt result,
CompareOp compare_op)
{
+ auto workaround_par = [=]() mutable
+ {
+#if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__binary_search::lower_bound<execution_policy<Derived>,
HaystackIt,
@@ -243,11 +248,17 @@ lower_bound(execution_policy<Derived>& policy,
OutputIt,
CompareOp>)
);
-#if __THRUST_HAS_HIPRT__
+ return first;
+#else
return __binary_search::lower_bound(
policy, first, last, values_first, values_last, result, compare_op
);
-#else
+#endif
+ };
+ (void)workaround_par;
+
+ auto workaround_seq = [=]() mutable
+ {
return thrust::lower_bound(cvt_to_seq(derived_cast(policy)),
first,
last,
@@ -255,6 +266,13 @@ lower_bound(execution_policy<Derived>& policy,
values_last,
result,
compare_op);
+ };
+ (void)workaround_seq;
+
+#if __THRUST_HAS_HIPRT__
+ return workaround_par();
+#else
+ return workaround_seq();
#endif
}
@@ -282,6 +300,9 @@ upper_bound(execution_policy<Derived>& policy,
OutputIt result,
CompareOp compare_op)
{
+ auto workaround_par = [=]() mutable
+ {
+#if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__binary_search::upper_bound<execution_policy<Derived>,
HaystackIt,
@@ -289,11 +310,17 @@ upper_bound(execution_policy<Derived>& policy,
OutputIt,
CompareOp>)
);
-#if __THRUST_HAS_HIPRT__
+ return first;
+#else
return __binary_search::upper_bound(
policy, first, last, values_first, values_last, result, compare_op
);
-#else
+#endif
+ };
+ (void)workaround_par;
+
+ auto workaround_seq = [=]() mutable
+ {
return thrust::upper_bound(cvt_to_seq(derived_cast(policy)),
first,
last,
@@ -301,6 +328,13 @@ upper_bound(execution_policy<Derived>& policy,
values_last,
result,
compare_op);
+ };
+ (void)workaround_seq;
+
+#if __THRUST_HAS_HIPRT__
+ return workaround_par();
+#else
+ return workaround_seq();
#endif
}
@@ -328,6 +362,9 @@ binary_search(execution_policy<Derived>& policy,
OutputIt result,
CompareOp compare_op)
{
+ auto workaround_par = [=]() mutable
+ {
+#if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__binary_search::binary_search<execution_policy<Derived>,
HaystackIt,
@@ -335,10 +372,17 @@ binary_search(execution_policy<Derived>& policy,
OutputIt,
CompareOp>)
);
-#if __THRUST_HAS_HIPRT__
- return __binary_search::binary_search(
- policy, first, last, values_first, values_last, result, compare_op);
+ return first;
#else
+ return __binary_search::binary_search(
+ policy, first, last, values_first, values_last, result, compare_op
+ );
+#endif
+ };
+ (void)workaround_par;
+
+ auto workaround_seq = [=]() mutable
+ {
return thrust::binary_search(cvt_to_seq(derived_cast(policy)),
first,
last,
@@ -346,6 +390,13 @@ binary_search(execution_policy<Derived>& policy,
values_last,
result,
compare_op);
+ };
+ (void)workaround_seq;
+
+#if __THRUST_HAS_HIPRT__
+ return workaround_par();
+#else
+ return workaround_seq();
#endif
}
@@ -363,6 +414,188 @@ binary_search(execution_policy<Derived>& policy,
);
}
+// Scalar functions
+
+// We use these custom implementations instead of thrust/system/detail/generic/binary_search.inl
+// because HIP support of device-side memory allocation is under development
+// (it is used in generic/binary_search.inl by thrust::detail::temporary_array).
+
+template<typename Derived, typename HaystackIt, typename T, typename CompareOp>
+THRUST_HIP_FUNCTION
+HaystackIt lower_bound(execution_policy<Derived>& policy,
+ HaystackIt first,
+ HaystackIt last,
+ const T& value,
+ CompareOp compare_op)
+{
+ using difference_type = typename thrust::iterator_traits<HaystackIt>::difference_type;
+ using values_type = typename thrust::detail::temporary_array<T, Derived>;
+ using results_type = typename thrust::detail::temporary_array<difference_type, Derived>;
+
+ auto workaround_par = [=]() mutable
+ {
+#if __HCC__ && __HIP_DEVICE_COMPILE__
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
+ (__binary_search::lower_bound<execution_policy<Derived>,
+ HaystackIt,
+ typename values_type::iterator,
+ typename results_type::iterator,
+ CompareOp>)
+ );
+ return first;
+#else
+ values_type values(policy, 1);
+ results_type result(policy, 1);
+
+ values[0] = value;
+
+ __binary_search::lower_bound(
+ policy, first, last, values.begin(), values.end(), result.begin(), compare_op
+ );
+
+ return first + result[0];
+#endif
+ };
+ (void)workaround_par;
+
+ auto workaround_seq = [=]() mutable
+ {
+ difference_type result;
+ thrust::lower_bound(cvt_to_seq(derived_cast(policy)),
+ first,
+ last,
+ &value,
+ &value + 1,
+ &result,
+ compare_op);
+ return first + result;
+ };
+ (void)workaround_seq;
+
+#if __THRUST_HAS_HIPRT__
+ return workaround_par();
+#else
+ return workaround_seq();
+#endif
+}
+
+template<typename Derived, typename HaystackIt, typename T, typename CompareOp>
+THRUST_HIP_FUNCTION
+HaystackIt upper_bound(execution_policy<Derived>& policy,
+ HaystackIt first,
+ HaystackIt last,
+ const T& value,
+ CompareOp compare_op)
+{
+ using difference_type = typename thrust::iterator_traits<HaystackIt>::difference_type;
+ using values_type = typename thrust::detail::temporary_array<T, Derived>;
+ using results_type = typename thrust::detail::temporary_array<difference_type, Derived>;
+
+ auto workaround_par = [=]() mutable
+ {
+#if __HCC__ && __HIP_DEVICE_COMPILE__
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
+ (__binary_search::upper_bound<execution_policy<Derived>,
+ HaystackIt,
+ typename values_type::iterator,
+ typename results_type::iterator,
+ CompareOp>)
+ );
+ return first;
+#else
+ values_type values(policy, 1);
+ results_type result(policy, 1);
+
+ values[0] = value;
+
+ __binary_search::upper_bound(
+ policy, first, last, values.begin(), values.end(), result.begin(), compare_op
+ );
+
+ return first + result[0];
+#endif
+ };
+ (void)workaround_par;
+
+ auto workaround_seq = [=]() mutable
+ {
+ difference_type result;
+ thrust::upper_bound(cvt_to_seq(derived_cast(policy)),
+ first,
+ last,
+ &value,
+ &value + 1,
+ &result,
+ compare_op);
+ return first + result;
+ };
+ (void)workaround_seq;
+
+#if __THRUST_HAS_HIPRT__
+ return workaround_par();
+#else
+ return workaround_seq();
+#endif
+}
+
+template<typename Derived, typename HaystackIt, typename T, typename CompareOp>
+THRUST_HIP_FUNCTION
+bool binary_search(execution_policy<Derived>& policy,
+ HaystackIt first,
+ HaystackIt last,
+ const T& value,
+ CompareOp compare_op)
+{
+ using values_type = typename thrust::detail::temporary_array<T, Derived>;
+ using results_type = typename thrust::detail::temporary_array<int, Derived>;
+
+ auto workaround_par = [=]() mutable
+ {
+#if __HCC__ && __HIP_DEVICE_COMPILE__
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
+ (__binary_search::binary_search<execution_policy<Derived>,
+ HaystackIt,
+ typename values_type::iterator,
+ typename results_type::iterator,
+ CompareOp>)
+ );
+ return first;
+#else
+ values_type values(policy, 1);
+ results_type result(policy, 1);
+
+ values[0] = value;
+
+ __binary_search::binary_search(
+ policy, first, last, values.begin(), values.end(), result.begin(), compare_op
+ );
+
+ return result[0] != 0;
+#endif
+ };
+ (void)workaround_par;
+
+ auto workaround_seq = [=]() mutable
+ {
+ bool result;
+ thrust::binary_search(cvt_to_seq(derived_cast(policy)),
+ first,
+ last,
+ &value,
+ &value + 1,
+ &result,
+ compare_op);
+ return result;
+ };
+ (void)workaround_seq;
+
+#if __THRUST_HAS_HIPRT__
+ return workaround_par();
+#else
+ return workaround_seq();
+#endif
+}
+
} // namespace hip_rocprim
END_NS_THRUST

Fix scalar binary search when using in kernels, add tests

The old implementation (from generic/binary_search.inl) uses temporary_array
which requires device-side memory allocation, which is currently
under development in HIP.
Add tests for both scalar and vector functions.