Merge pull request kokkos#7264 from masterleinad/improve_kokkos_sort_…

…performance Use raw pointers for std::sort if possible
ndellingwood · Aug 27, 2024 · 03abc0a · 03abc0a
2 parents fad5e74 + 449d593
commit 03abc0a
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 13 deletions.
diff --git a/algorithms/src/sorting/Kokkos_SortPublicAPI.hpp b/algorithms/src/sorting/Kokkos_SortPublicAPI.hpp
@@ -53,9 +53,13 @@ void sort(const ExecutionSpace& exec,
 
   if constexpr (Impl::better_off_calling_std_sort_v<ExecutionSpace>) {
     exec.fence("Kokkos::sort without comparator use std::sort");
-    auto first = ::Kokkos::Experimental::begin(view);
-    auto last  = ::Kokkos::Experimental::end(view);
-    std::sort(first, last);
+    if (view.span_is_contiguous()) {
+      std::sort(view.data(), view.data() + view.size());
+    } else {
+      auto first = ::Kokkos::Experimental::begin(view);
+      auto last  = ::Kokkos::Experimental::end(view);
+      std::sort(first, last);
+    }
   } else {
     Impl::sort_device_view_without_comparator(exec, view);
   }
@@ -107,9 +111,13 @@ void sort(const ExecutionSpace& exec,
 
   if constexpr (Impl::better_off_calling_std_sort_v<ExecutionSpace>) {
     exec.fence("Kokkos::sort with comparator use std::sort");
-    auto first = ::Kokkos::Experimental::begin(view);
-    auto last  = ::Kokkos::Experimental::end(view);
-    std::sort(first, last, comparator);
+    if (view.span_is_contiguous()) {
+      std::sort(view.data(), view.data() + view.size(), comparator);
+    } else {
+      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);
   }

diff --git a/algorithms/src/sorting/impl/Kokkos_SortImpl.hpp b/algorithms/src/sorting/impl/Kokkos_SortImpl.hpp
@@ -268,19 +268,29 @@ void copy_to_host_run_stdsort_copy_back(
     KE::copy(exec, view, view_dc);
 
     // run sort on the mirror of view_dc
-    auto mv_h  = create_mirror_view_and_copy(Kokkos::HostSpace(), view_dc);
-    auto first = KE::begin(mv_h);
-    auto last  = KE::end(mv_h);
-    std::sort(first, last, std::forward<MaybeComparator>(maybeComparator)...);
+    auto mv_h = create_mirror_view_and_copy(Kokkos::HostSpace(), view_dc);
+    if (view.span_is_contiguous()) {
+      std::sort(mv_h.data(), mv_h.data() + mv_h.size(),
+                std::forward<MaybeComparator>(maybeComparator)...);
+    } else {
+      auto first = KE::begin(mv_h);
+      auto last  = KE::end(mv_h);
+      std::sort(first, last, std::forward<MaybeComparator>(maybeComparator)...);
+    }
     Kokkos::deep_copy(exec, view_dc, mv_h);
 
     // copy back to argument view
     KE::copy(exec, KE::cbegin(view_dc), KE::cend(view_dc), KE::begin(view));
   } else {
     auto view_h = create_mirror_view_and_copy(Kokkos::HostSpace(), view);
-    auto first  = KE::begin(view_h);
-    auto last   = KE::end(view_h);
-    std::sort(first, last, std::forward<MaybeComparator>(maybeComparator)...);
+    if (view.span_is_contiguous()) {
+      std::sort(view_h.data(), view_h.data() + view_h.size(),
+                std::forward<MaybeComparator>(maybeComparator)...);
+    } else {
+      auto first = KE::begin(view_h);
+      auto last  = KE::end(view_h);
+      std::sort(first, last, std::forward<MaybeComparator>(maybeComparator)...);
+    }
     Kokkos::deep_copy(exec, view, view_h);
   }
 }