Update partition workarounds
MODIFY
thrust/system/hip/detail/partition.h
thrust/system/hip/detail/partition.h
@@ -381,9 +381,19 @@ partition_copy(execution_policy<Derived>& policy,
RejectedOutIt rejected_result,
Predicate predicate)
{
- pair<SelectedOutIt, RejectedOutIt> ret
- = thrust::make_pair(selected_result, rejected_result);
-
+ struct workaround
+ {
+ __host__
+ static pair<SelectedOutIt, RejectedOutIt> par(
+ execution_policy<Derived>& policy,
+ InputIt first,
+ InputIt last,
+ StencilIt stencil,
+ SelectedOutIt selected_result,
+ RejectedOutIt rejected_result,
+ Predicate predicate)
+ {
+ #if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__partition::partition<detail::false_type,
Derived,
@@ -393,19 +403,39 @@ partition_copy(execution_policy<Derived>& policy,
RejectedOutIt,
Predicate>)
);
-#if __THRUST_HAS_HIPRT__
- ret = __partition::partition<detail::false_type>(
- policy, first, last, stencil, selected_result, rejected_result, predicate);
-#else // __THRUST_HAS_HIPRT__
- ret = thrust::partition_copy(cvt_to_seq(derived_cast(policy)),
+ #else
+
+ return __partition::partition<detail::false_type>(
+ policy, first, last, stencil, selected_result, rejected_result, predicate
+ );
+ #endif
+ }
+ __device__
+ static pair<SelectedOutIt, RejectedOutIt> seq(
+ execution_policy<Derived>& policy,
+ InputIt first,
+ InputIt last,
+ StencilIt stencil,
+ SelectedOutIt selected_result,
+ RejectedOutIt rejected_result,
+ Predicate predicate)
+ {
+ return thrust::partition_copy(
+ cvt_to_seq(derived_cast(policy)),
first,
last,
stencil,
selected_result,
rejected_result,
- predicate);
-#endif // __THRUST_HAS_HIPRT__
- return ret;
+ predicate
+ );
+ }
+ };
+ #if __THRUST_HAS_HIPRT__
+ return workaround::par(policy, first, last, stencil, selected_result, rejected_result, predicate);
+ #else
+ return workaround::seq(policy, first, last, stencil, selected_result, rejected_result, predicate);
+ #endif
}
__thrust_exec_check_disable__ template <class Derived,
@@ -421,12 +451,11 @@ partition_copy(execution_policy<Derived>& policy,
RejectedOutIt rejected_result,
Predicate predicate)
{
-
-
struct workaround
{
__host__
- static pair<SelectedOutIt, RejectedOutIt> par(execution_policy<Derived>& policy,
+ static pair<SelectedOutIt, RejectedOutIt> par(
+ execution_policy<Derived>& policy,
InputIt first,
InputIt last,
SelectedOutIt selected_result,
@@ -440,15 +469,17 @@ partition_copy(execution_policy<Derived>& policy,
InputIt,
SelectedOutIt,
RejectedOutIt,
- Predicate>)
+ Predicate>
+ );
#else
-
- return __partition::partition<detail::false_type>( policy, first, last, selected_result, rejected_result, predicate);
-
+ return __partition::partition<detail::false_type>(
+ policy, first, last, selected_result, rejected_result, predicate
+ );
#endif
}
__device__
- static pair<SelectedOutIt, RejectedOutIt> seq(execution_policy<Derived>& policy,
+ static pair<SelectedOutIt, RejectedOutIt> seq(
+ execution_policy<Derived>& policy,
InputIt first,
InputIt last,
SelectedOutIt selected_result,
@@ -468,8 +499,6 @@ partition_copy(execution_policy<Derived>& policy,
#else
return workaround::seq(policy, first, last, selected_result, rejected_result, predicate);
#endif
-
-
}
__thrust_exec_check_disable__ template <class Derived,
@@ -485,8 +514,18 @@ stable_partition_copy(execution_policy<Derived>& policy,
RejectedOutIt rejected_result,
Predicate predicate)
{
- pair<SelectedOutIt, RejectedOutIt> ret
- = thrust::make_pair(selected_result, rejected_result);
+ struct workaround
+ {
+ __host__
+ static pair<SelectedOutIt, RejectedOutIt> par(
+ execution_policy<Derived>& policy,
+ InputIt first,
+ InputIt last,
+ SelectedOutIt selected_result,
+ RejectedOutIt rejected_result,
+ Predicate predicate)
+ {
+ #if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__partition::partition<detail::false_type,
Derived,
@@ -496,18 +535,35 @@ stable_partition_copy(execution_policy<Derived>& policy,
Predicate>)
);
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND((hip_rocprim::reverse<Derived, InputIt>));
-#if __THRUST_HAS_HIPRT__
- ret = __partition::partition<detail::false_type>(
+ #else
+ return __partition::partition<detail::false_type>(
policy, first, last, selected_result, rejected_result, predicate);
-#else // __THRUST_HAS_HIPRT__
- ret = thrust::stable_partition_copy(cvt_to_seq(derived_cast(policy)),
+ #endif
+ }
+ __device__
+ static pair<SelectedOutIt, RejectedOutIt> seq(
+ execution_policy<Derived>& policy,
+ InputIt first,
+ InputIt last,
+ SelectedOutIt selected_result,
+ RejectedOutIt rejected_result,
+ Predicate predicate)
+ {
+ return thrust::stable_partition_copy(
+ cvt_to_seq(derived_cast(policy)),
first,
last,
selected_result,
rejected_result,
- predicate);
-#endif // __THRUST_HAS_HIPRT__
- return ret;
+ predicate
+ );
+ }
+ };
+ #if __THRUST_HAS_HIPRT__
+ return workaround::par(policy, first, last, selected_result, rejected_result, predicate);
+ #else
+ return workaround::seq(policy, first, last, selected_result, rejected_result, predicate);
+ #endif
}
__thrust_exec_check_disable__ template <class Derived,
@@ -525,8 +581,19 @@ stable_partition_copy(execution_policy<Derived>& policy,
RejectedOutIt rejected_result,
Predicate predicate)
{
- pair<SelectedOutIt, RejectedOutIt> ret
- = thrust::make_pair(selected_result, rejected_result);
+ struct workaround
+ {
+ __host__
+ static pair<SelectedOutIt, RejectedOutIt> par(
+ execution_policy<Derived>& policy,
+ InputIt first,
+ InputIt last,
+ StencilIt stencil,
+ SelectedOutIt selected_result,
+ RejectedOutIt rejected_result,
+ Predicate predicate)
+ {
+ #if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__partition::partition<detail::false_type,
Derived,
@@ -537,19 +604,37 @@ stable_partition_copy(execution_policy<Derived>& policy,
Predicate>)
);
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND((hip_rocprim::reverse<Derived, InputIt>));
-#if __THRUST_HAS_HIPRT__
- ret = __partition::partition<detail::false_type>(
+ #else
+ return __partition::partition<detail::false_type>(
policy, first, last, stencil, selected_result, rejected_result, predicate);
-#else // __THRUST_HAS_HIPRT__
- ret = thrust::stable_partition_copy(cvt_to_seq(derived_cast(policy)),
+ #endif
+ }
+ __device__
+ static pair<SelectedOutIt, RejectedOutIt> seq(
+ execution_policy<Derived>& policy,
+ InputIt first,
+ InputIt last,
+ StencilIt stencil,
+ SelectedOutIt selected_result,
+ RejectedOutIt rejected_result,
+ Predicate predicate)
+ {
+ return thrust::stable_partition_copy(
+ cvt_to_seq(derived_cast(policy)),
first,
last,
stencil,
selected_result,
rejected_result,
- predicate);
-#endif // __THRUST_HAS_HIPRT__
- return ret;
+ predicate
+ );
+ }
+ };
+ #if __THRUST_HAS_HIPRT__
+ return workaround::par(policy, first, last, stencil, selected_result, rejected_result, predicate);
+ #else
+ return workaround::seq(policy, first, last, stencil, selected_result, rejected_result, predicate);
+ #endif
}
/// inplace
@@ -564,16 +649,55 @@ partition(execution_policy<Derived>& policy,
StencilIt stencil,
Predicate predicate)
{
- Iterator ret = first;
+ struct workaround
+ {
+ __host__
+ static Iterator par(execution_policy<Derived>& policy,
+ Iterator first,
+ Iterator last,
+ StencilIt stencil,
+ Predicate predicate)
+ {
+ #if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__partition::partition_inplace<Derived, Iterator, StencilIt, Predicate>)
);
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
+ (hip_rocprim::reverse<Derived, Iterator>)
+ );
+ #else
+ Iterator result = __partition::partition_inplace(
+ policy,
+ first,
+ last,
+ stencil,
+ predicate
+ );
+ hip_rocprim::reverse<Derived,Iterator>(policy, result, last);
+ return result;
+ #endif
+ }
+ __device__
+ static Iterator seq(execution_policy<Derived>& policy,
+ Iterator first,
+ Iterator last,
+ StencilIt stencil,
+ Predicate predicate)
+ {
+ return thrust::partition(
+ cvt_to_seq(derived_cast(policy)),
+ first,
+ last,
+ stencil,
+ predicate
+ );
+ }
+ };
#if __THRUST_HAS_HIPRT__
- ret = __partition::partition_inplace(policy, first, last, stencil, predicate);
-#else // __THRUST_HAS_HIPRT__
- ret = thrust::partition(cvt_to_seq(derived_cast(policy)), first, last, stencil, predicate);
-#endif // __THRUST_HAS_HIPRT__
- return ret;
+ return workaround::par(policy, first, last, stencil, predicate);
+ #else
+ return workaround::seq(policy, first, last, stencil, predicate);
+ #endif
}
__thrust_exec_check_disable__ template <class Derived, class Iterator, class Predicate>
@@ -597,8 +721,8 @@ partition(execution_policy<Derived>& policy,
(__partition::partition<Derived, Iterator, Iterator, StencilIt, Predicate>)
);
#else
-
- return __partition::partition_inplace(policy,
+ return __partition::partition_inplace(
+ policy,
first,
last,
predicate
@@ -638,22 +762,55 @@ stable_partition(execution_policy<Derived>& policy,
StencilIt stencil,
Predicate predicate)
{
- Iterator result = first;
+ struct workaround
+ {
+ __host__
+ static Iterator par(execution_policy<Derived>& policy,
+ Iterator first,
+ Iterator last,
+ StencilIt stencil,
+ Predicate predicate)
+ {
+ #if __HCC__ && __HIP_DEVICE_COMPILE__
THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
(__partition::partition_inplace<Derived, Iterator, StencilIt, Predicate>)
);
- THRUST_HIP_PRESERVE_KERNELS_WORKAROUND((hip_rocprim::reverse<Derived, Iterator>));
-#if __THRUST_HAS_HIPRT__
- result = __partition::partition_inplace(policy, first, last, stencil, predicate);
-
- // partition returns rejected values in reverse order
- // so reverse the rejected elements to make it stable
- hip_rocprim::reverse(policy, result, last);
-#else // __THRUST_HAS_HIPRT__
- result = thrust::stable_partition(
- cvt_to_seq(derived_cast(policy)), first, last, stencil, predicate);
-#endif // __THRUST_HAS_HIPRT__
+ THRUST_HIP_PRESERVE_KERNELS_WORKAROUND(
+ (hip_rocprim::reverse<Derived, Iterator>)
+ );
+ #else
+ Iterator result = __partition::partition_inplace(
+ policy,
+ first,
+ last,
+ stencil,
+ predicate
+ );
+ hip_rocprim::reverse<Derived,Iterator>(policy, result, last);
return result;
+ #endif
+ }
+ __device__
+ static Iterator seq(execution_policy<Derived>& policy,
+ Iterator first,
+ Iterator last,
+ StencilIt stencil,
+ Predicate predicate)
+ {
+ return thrust::stable_partition(
+ cvt_to_seq(derived_cast(policy)),
+ first,
+ last,
+ stencil,
+ predicate
+ );
+ }
+ };
+ #if __THRUST_HAS_HIPRT__
+ return workaround::par(policy, first, last, stencil, predicate);
+ #else
+ return workaround::seq(policy, first, last, stencil, predicate);
+ #endif
}
__thrust_exec_check_disable__ template <class Derived, class Iterator, class Predicate>
@@ -663,7 +820,6 @@ stable_partition(execution_policy<Derived>& policy,
Iterator last,
Predicate predicate)
{
-
struct workaround
{
__host__
@@ -680,8 +836,8 @@ stable_partition(execution_policy<Derived>& policy,
(hip_rocprim::reverse<Derived, Iterator>)
);
#else
-
- Iterator result = __partition::partition_inplace(policy,
+ Iterator result = __partition::partition_inplace(
+ policy,
first,
last,
predicate
@@ -709,8 +865,6 @@ stable_partition(execution_policy<Derived>& policy,
#else
return workaround::seq(policy, first, last, predicate);
#endif
-
-
}
template <class Derived, class ItemsIt, class Predicate>

Update partition workarounds