Skip to content

Commit

Permalink
Implement Impl::max_reduce since ArborX::max has been deprecated
Browse files Browse the repository at this point in the history
  • Loading branch information
dalg24 committed Jan 3, 2024
1 parent 973e7d0 commit 7fa71f2
Showing 1 changed file with 25 additions and 1 deletion.
26 changes: 25 additions & 1 deletion core/src/Cabana_Experimental_NeighborList.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

#include <Cabana_NeighborList.hpp>
#include <Cabana_Slice.hpp>
#include <Cabana_Types.hpp> // is_accessible_from

#include <ArborX.hpp>

Expand Down Expand Up @@ -69,6 +70,29 @@ auto makePredicates(
return Impl::SubsliceAndRadius<stdcxx20::remove_cvref_t<Slice>>{
std::forward<Slice>( slice ), first, last, radius };
}

template <typename ExecutionSpace, typename D, typename... P>
typename Kokkos::View<D, P...>::non_const_value_type
max_reduce( ExecutionSpace const& space, Kokkos::View<D, P...> const& v )
{
using V = Kokkos::View<D, P...>;
static_assert( V::rank() == 1 );
static_assert( Kokkos::is_execution_space_v<ExecutionSpace> );
static_assert(
is_accessible_from<typename V::memory_space, ExecutionSpace>::value );
using Ret = typename Kokkos::View<D, P...>::non_const_value_type;
Ret max_val;
Kokkos::parallel_reduce(
Kokkos::RangePolicy<ExecutionSpace>( space, 0, v.extent( 0 ) ),
KOKKOS_LAMBDA( int i, Ret& partial_max ) {
if ( v( i ) > partial_max )
{
partial_max = v( i );
}
},
Kokkos::Max<Ret>( max_val ) );
return max_val;
}
//! \endcond
} // namespace Impl
} // namespace Experimental
Expand Down Expand Up @@ -437,7 +461,7 @@ auto make2DNeighborList( ExecutionSpace space, Tag,
Tag>{ counts } );
}

auto const max_neighbors = ArborX::max( space, counts );
auto const max_neighbors = Impl::max_reduce( space, counts );
if ( max_neighbors <= buffer_size )
{
// NOTE We do not bother shrinking to eliminate the excess allocation.
Expand Down

0 comments on commit 7fa71f2

Please sign in to comment.