Fixed min_element, max_element and min_max_element
MODIFY
test/test_max_element.cpp
test/test_max_element.cpp
@@ -89,10 +89,10 @@ TYPED_TEST(MaxElementTests, TestMaxElementSimple)
data[5] = 1;
ASSERT_EQ( *thrust::max_element(data.begin(), data.end()), 5);
- ASSERT_EQ( thrust::max_element(data.begin(), data.end()) - data.begin(), 2);
+ ASSERT_EQ( thrust::max_element(data.begin(), data.end()) - data.begin(), 1);
- ASSERT_EQ( *thrust::max_element(data.begin(), data.end(), thrust::less<T>()), 1);
- ASSERT_EQ( thrust::max_element(data.begin(), data.end(), thrust::less<T>()) - data.begin(), -2);
+ ASSERT_EQ( *thrust::max_element(data.begin(), data.end(), thrust::greater<T>()), 1);
+ ASSERT_EQ( thrust::max_element(data.begin(), data.end(), thrust::greater<T>()) - data.begin(), 2);
}
TYPED_TEST(MaxElementTests, TestMaxElementWithTransform)
@@ -114,8 +114,7 @@ TYPED_TEST(MaxElementTests, TestMaxElementWithTransform)
ASSERT_EQ( *thrust::max_element(
thrust::make_transform_iterator(data.begin(), thrust::negate<T>()),
thrust::make_transform_iterator(data.end(), thrust::negate<T>()),
- thrust::less<T>()), -5);
-
+ thrust::greater<T>()), -5);
}
TYPED_TEST(MaxElementTests, TestMaxElement)

MODIFY
thrust/system/hip/detail/extrema.h
thrust/system/hip/detail/extrema.h
#include <thrust/extrema.h>
#include <thrust/pair.h>
#include <thrust/distance.h>
+#include <thrust/system/hip/detail/get_value.h>
+
+// rocprim include
+#include <rocprim/rocprim.hpp>
BEGIN_NS_THRUST
+namespace hip_rocprim {
namespace __extrema {
@@ -141,29 +146,60 @@ namespace __extrema {
template <class Derived,
class InputIt,
+ class Size,
class BinaryOp,
class T>
- T THRUST_HIP_FUNCTION
+ T THRUST_HIP_RUNTIME_FUNCTION
extrema(execution_policy<Derived> &policy,
InputIt first,
- InputIt last,
+ Size num_items,
BinaryOp binary_op,
T *)
{
- T result;
- T init = *first;
+ hipStream_t stream = hip_rocprim::stream(policy);
+
+
+ // Determine temporary device storage requirements.
+ T* ret_ptr = NULL;
+ size_t tmp_size = 0;
+ hip_rocprim::throw_on_error(
+ rocprim::reduce(nullptr, tmp_size,
+ first, ret_ptr, num_items, binary_op,
+ stream, THRUST_HIP_DEBUG_SYNC_FLAG),
+ "after reduction step 1");
- result = reduce(policy, first, last, init, binary_op);
+ // Allocate temporary storage.
- return result;
+ detail::temporary_array<detail::uint8_t, Derived>
+ tmp(policy, sizeof(T) + tmp_size);
+
+ // Run reduction.
+
+ // `tmp.begin()` yields a `normal_iterator`, which dereferences to a
+ // `reference`, which has an `operator&` that returns a `pointer`, which
+ // has a `.get` method that returns a raw pointer, which we can (finally)
+ // `static_cast` to `void*`.
+ //
+ // The array was dynamically allocated, so we assume that it's suitably
+ // aligned for any type of data. `malloc`/`cudaMalloc`/`new`/`std::allocator`
+ // make this guarantee.
+ ret_ptr = detail::aligned_reinterpret_cast<T*>((&*tmp.begin()).get());
+ void* tmp_ptr = static_cast<void*>((&*(tmp.begin() + sizeof(T))).get());
+ hip_rocprim::throw_on_error(
+ rocprim::reduce(tmp_ptr, tmp_size,
+ first, ret_ptr, num_items, binary_op,
+ stream, THRUST_HIP_DEBUG_SYNC_FLAG),
+ "after reduction step 2");
+
+ return hip_rocprim::get_value(policy,ret_ptr);
}
template <template <class, class, class> class ArgFunctor,
class Derived,
class ItemsIt,
class BinaryPred>
- ItemsIt THRUST_HIP_FUNCTION
+ ItemsIt THRUST_HIP_RUNTIME_FUNCTION
element(execution_policy<Derived> &policy,
ItemsIt first,
ItemsIt last,
@@ -177,21 +213,20 @@ namespace __extrema {
IndexType num_items = static_cast<IndexType>(thrust::distance(first, last));
- typedef tuple<ItemsIt, hip_rocprim::counting_iterator_t<IndexType> > iterator_tuple;
+ typedef tuple<ItemsIt, counting_iterator_t<IndexType> > iterator_tuple;
typedef zip_iterator<iterator_tuple> zip_iterator;
- iterator_tuple iter_tuple_begin = make_tuple(first, hip_rocprim::counting_iterator_t<IndexType>(0));
- iterator_tuple iter_tuple_end = make_tuple(last, hip_rocprim::counting_iterator_t<IndexType>(num_items));
+ iterator_tuple iter_tuple = make_tuple(first, counting_iterator_t<IndexType>(0));
+
typedef ArgFunctor<InputType, IndexType, BinaryPred> arg_min_t;
typedef tuple<InputType, IndexType> T;
- zip_iterator begin = make_zip_iterator(iter_tuple_begin);
- zip_iterator end = make_zip_iterator(iter_tuple_end);
+ zip_iterator begin = make_zip_iterator(iter_tuple);
T result = extrema(policy,
begin,
- end,
+ num_items,
arg_min_t(binary_pred),
(T *)(NULL));
return first + thrust::get<1>(result);
@@ -201,10 +236,7 @@ namespace __extrema {
} // namespace __extrema
-namespace hip_rocprim {
-
/// min element
-
__thrust_exec_check_disable__
template <class Derived,
class ItemsIt,
@@ -216,6 +248,10 @@ min_element(execution_policy<Derived> &policy,
BinaryPred binary_pred)
{
ItemsIt ret = first;
+
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND((
+ __extrema::element<__extrema::arg_min_f,Derived,ItemsIt,BinaryPred>
+ ));
#if __THRUST_HAS_HIPRT__
ret = __extrema::element<__extrema::arg_min_f>(policy,
first,
@@ -256,6 +292,10 @@ max_element(execution_policy<Derived> &policy,
BinaryPred binary_pred)
{
ItemsIt ret = first;
+
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND((
+ __extrema::element<__extrema::arg_max_f,Derived,ItemsIt,BinaryPred>
+ ));
#if __THRUST_HAS_HIPRT__
ret = __extrema::element<__extrema::arg_max_f>(policy,
first,
@@ -295,22 +335,14 @@ minmax_element(execution_policy<Derived> &policy,
{
pair<ItemsIt, ItemsIt> ret = thrust::make_pair(first, first);
-#if __THRUST_HAS_HIPRT__
- if (first == last)
- return thrust::make_pair(last, last);
-
typedef typename iterator_traits<ItemsIt>::value_type InputType;
typedef typename iterator_traits<ItemsIt>::difference_type IndexType;
- IndexType num_items = static_cast<IndexType>(thrust::distance(first, last));
-
typedef tuple<ItemsIt, hip_rocprim::counting_iterator_t<IndexType> > iterator_tuple;
typedef zip_iterator<iterator_tuple> zip_iterator;
- iterator_tuple iter_tuple_begin = make_tuple(first, hip_rocprim::counting_iterator_t<IndexType>(0));
- iterator_tuple iter_tuple_end = make_tuple(last, hip_rocprim::counting_iterator_t<IndexType>(num_items));
-
typedef __extrema::arg_minmax_f<InputType, IndexType, BinaryPred> arg_minmax_t;
+
typedef typename arg_minmax_t::two_pairs_type two_pairs_type;
typedef typename arg_minmax_t::duplicate_tuple duplicate_t;
typedef transform_input_iterator_t<two_pairs_type,
@@ -318,11 +350,21 @@ minmax_element(execution_policy<Derived> &policy,
duplicate_t>
transform_t;
- zip_iterator begin = make_zip_iterator(iter_tuple_begin);
- zip_iterator end = make_zip_iterator(iter_tuple_end);
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND((
+ __extrema::element<Derived,transform_t,IndexType,arg_minmax_t,two_pairs_type>
+ ));
+#if __THRUST_HAS_HIPRT__
+ if (first == last)
+ return thrust::make_pair(last, last);
+
+ IndexType num_items = static_cast<IndexType>(thrust::distance(first, last));
+
+ iterator_tuple iter_tuple = make_tuple(first, counting_iterator_t<IndexType>(0));
+
+ zip_iterator begin = make_zip_iterator(iter_tuple);
two_pairs_type result = __extrema::extrema(policy,
transform_t(begin, duplicate_t()),
- transform_t(end, duplicate_t()),
+ num_items,
arg_minmax_t(binary_pred),
(two_pairs_type *)(NULL));
ret = thrust::make_pair(first + get<1>(get<0>(result)),

Fixed min_element, max_element and min_max_element