Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Kokkos::sort support custom comparator #6253

Merged
merged 12 commits into from
Jul 20, 2023

Conversation

fnrizzi
Copy link
Contributor

@fnrizzi fnrizzi commented Jul 3, 2023

Adds 2 overloads to Kokkos::sort to support a custom comparator.

template <class ExecutionSpace, class ComparatorType, class DataType, class... Properties>
void sort([[maybe_unused]] const ExecutionSpace& exec,
          const Kokkos::View<DataType, Properties...>& view,
          ComparatorType const& comparator)
{
   /* view must be rank-1 with layout{left,right,stride} */

  if constexpr (Impl::better_off_calling_std_sort_v<ExecutionSpace>){
    auto first = ::Kokkos::Experimental::begin(view);
    auto last  = ::Kokkos::Experimental::end(view);
    std::sort(first, last, comparator);
  } else {
    Impl::sort_device_view_with_comparator(exec, view, comparator);
  }
}

template <class ComparatorType, class DataType, class... Properties>
void sort(const Kokkos::View<DataType, Properties...>& view, ComparatorType const& comparator){
  // calls overload above using the view execution space
}

where:

Impl::sort_device_view_with_comparator(exec, view, comparator) 

calls cuda thrust for cuda, onedpl and the fallback case (for now) copies view to host, calls std::sort and copies back.
This solution is based on a discussion with @crtrott and @dalg24.

Note that this is potentially the place where we can later optimize for example implementing our own quicksort or what not.

@fnrizzi fnrizzi changed the title sort: support custom comparator Kokkos::sort: support custom comparator Jul 3, 2023
@fnrizzi fnrizzi changed the title Kokkos::sort: support custom comparator Kokkos::sort support custom comparator Jul 3, 2023
@fnrizzi

This comment was marked as outdated.

@fnrizzi
Copy link
Contributor Author

fnrizzi commented Jul 3, 2023

retest this please

algorithms/src/sorting/Kokkos_SortPublicAPI.hpp Outdated Show resolved Hide resolved
algorithms/src/sorting/impl/Kokkos_SortImpl.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/src/sorting/Kokkos_SortPublicAPI.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/unit_tests/TestSortCustomComp.hpp Outdated Show resolved Hide resolved
algorithms/src/sorting/impl/Kokkos_SortImpl.hpp Outdated Show resolved Hide resolved
@fnrizzi fnrizzi force-pushed the sort_with_custom_comparison branch from 23f7df8 to 5f494a2 Compare July 17, 2023 09:57
@fnrizzi
Copy link
Contributor Author

fnrizzi commented Jul 18, 2023

retest this please

1 similar comment
@fnrizzi
Copy link
Contributor Author

fnrizzi commented Jul 19, 2023

retest this please

@fnrizzi
Copy link
Contributor Author

fnrizzi commented Jul 19, 2023

Only CI failure is related to #6293

Copy link
Member

@dalg24 dalg24 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would approve once we clarify whether the fence is needed or not

algorithms/src/sorting/Kokkos_SortPublicAPI.hpp Outdated Show resolved Hide resolved
algorithms/src/sorting/impl/Kokkos_SortImpl.hpp Outdated Show resolved Hide resolved
@fnrizzi fnrizzi force-pushed the sort_with_custom_comparison branch from 80edadb to 6b39b4d Compare July 19, 2023 12:56
@dalg24
Copy link
Member

dalg24 commented Jul 20, 2023

The failure is unrelated (timeout during one of the HIP build)

Copy link
Member

@crtrott crtrott left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could consider later to keep a buffer around (say 500k or so) which would for smaller views allow us to avoid the costly alloc/dealloc in the host copy case. But we can do that later.

@crtrott
Copy link
Member

crtrott commented Jul 20, 2023

Failure looks unrelated.

@crtrott crtrott merged commit 78b856f into kokkos:develop Jul 20, 2023
27 of 28 checks passed
tcclevenger pushed a commit to tcclevenger/kokkos that referenced this pull request Jul 20, 2023
* sort: support custom comparator

* fix wrong call for sycl

* tests: add cases for layout left and right

* address reviews

* more changes to address reviews

* fix global namespace

* fix namespace

* slim function

* use west const

* remove fence not needed

* use west const

* enable format
@crtrott crtrott mentioned this pull request Aug 25, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants