Merge topic 'thrust-sorting'

267ee49cb docker: update kokkos hip image
67bf9a966 docker: update kokkos hip image
50d4ab5cc CMake: VTKm_ENABLE_KOKKOS_THRUST to depend on Kokkos with CUDA or HIP enabled
0604f314a Fix for building SERIAL unit tests with KOKKOS_HIP/CUDA enabled
802bf80b0 CI: Add rocthrust installation step
fda475d5b Rework Thrust CMake options
e6f63a807 Adding CMake tweaks to turn off thrust algorithms if thrust is not detected.
5a72275ed Rewrite sorting specialization using std::enable_if_t
...

Acked-by: Kitware Robot <kwrobot@kitware.com>
Merge-request: !2926

.gitlab-ci.yml

@@ -58,7 +58,7 @@
     - .docker_image
 
 .ubuntu2004_hip_kokkos: &ubuntu2004_hip_kokkos
-  image: "kitware/vtkm:ci-ubuntu2004_hip_kokkos-20230125"
+  image: "kitware/vtkm:ci-ubuntu2004_hip_kokkos-20230220"
   extends:
     - .docker_image
 

.gitlab/ci/docker/ubuntu2004/kokkos-hip/Dockerfile

@@ -14,6 +14,7 @@ RUN apt update && \
     ninja-build \
     rsync \
     ssh \
+    rocthrust-dev \
     && \
     apt clean
 

CMake/VTKmDeviceAdapters.cmake

@@ -357,6 +357,14 @@ if(VTKm_ENABLE_KOKKOS AND NOT TARGET vtkm_kokkos)
     add_library(vtkm_kokkos_hip INTERFACE)
     set_property(TARGET vtkm_kokkos_hip PROPERTY EXPORT_NAME kokkos_hip)
     install(TARGETS vtkm_kokkos_hip EXPORT ${VTKm_EXPORT_NAME})
+
+    # Make sure rocthrust is available if requested
+    if(VTKm_ENABLE_KOKKOS_THRUST)
+      find_package(rocthrust)
+      if(NOT rocthrust_FOUND)
+        message(FATAL_ERROR "rocthrust not found. Please set VTKm_ENABLE_KOKKOS_THRUST to OFF.")
+      endif()
+    endif()
   endif()
 
   add_library(vtkm_kokkos INTERFACE IMPORTED GLOBAL)

CMakeLists.txt

@@ -227,6 +227,11 @@ include(VTKmBuildType)
 # Include the vtk-m wrappers
 include(VTKmWrappers)
 
+# By default: Set VTKm_ENABLE_KOKKOS_THRUST to ON if VTKm_ENABLE_KOKKOS is ON, otherwise
+# disable it (or if the user explicitly turns this option OFF)
+cmake_dependent_option(VTKm_ENABLE_KOKKOS_THRUST "Enable Kokkos thrust support (only valid with CUDA and HIP)"
+  ON "VTKm_ENABLE_KOKKOS;Kokkos_ENABLE_CUDA OR Kokkos_ENABLE_HIP" OFF)
+
 # Create vtkm_compiler_flags library. This is an interface library that
 # holds all the C++ compiler flags that are needed for consumers and
 # when building VTK-m.

vtkm/cont/kokkos/internal/DeviceAdapterAlgorithmKokkos.h

@@ -38,6 +38,14 @@ VTKM_THIRDPARTY_POST_INCLUDE
 #define VTKM_VOLATILE volatile
 #endif
 
+#if defined(VTKM_ENABLE_KOKKOS_THRUST) && (defined(__HIP__) || defined(__CUDA__))
+#define VTKM_USE_KOKKOS_THRUST
+#endif
+
+#if defined(VTKM_USE_KOKKOS_THRUST)
+#include <thrust/device_ptr.h>
+#include <thrust/sort.h>
+#endif
 
 namespace vtkm
 {
@@ -771,6 +779,88 @@ public:
     SortImpl(values, comp, typename std::is_scalar<T>::type{});
   }
 
+protected:
+  // Kokkos currently (11/10/2022) does not support a sort_by_key operator
+  // so instead we are using thrust if and only if HIP or CUDA are the backends for Kokkos
+#if defined(VTKM_USE_KOKKOS_THRUST)
+
+  template <typename T, typename U, typename BinaryCompare>
+  VTKM_CONT static std::enable_if_t<(std::is_same<BinaryCompare, vtkm::SortLess>::value ||
+                                     std::is_same<BinaryCompare, vtkm::SortGreater>::value)>
+  SortByKeyImpl(vtkm::cont::ArrayHandle<T>& keys,
+                vtkm::cont::ArrayHandle<U>& values,
+                BinaryCompare,
+                std::true_type,
+                std::true_type)
+  {
+    vtkm::cont::Token token;
+    auto keys_portal = keys.PrepareForInPlace(vtkm::cont::DeviceAdapterTagKokkos{}, token);
+    auto values_portal = values.PrepareForInPlace(vtkm::cont::DeviceAdapterTagKokkos{}, token);
+
+    kokkos::internal::KokkosViewExec<T> keys_view(keys_portal.GetArray(),
+                                                  keys_portal.GetNumberOfValues());
+    kokkos::internal::KokkosViewExec<U> values_view(values_portal.GetArray(),
+                                                    values_portal.GetNumberOfValues());
+
+    thrust::device_ptr<T> keys_begin(keys_view.data());
+    thrust::device_ptr<T> keys_end(keys_view.data() + keys_view.size());
+    thrust::device_ptr<U> values_begin(values_view.data());
+
+    if (std::is_same<BinaryCompare, vtkm::SortLess>::value)
+    {
+      thrust::sort_by_key(keys_begin, keys_end, values_begin, thrust::less<T>());
+    }
+    else
+    {
+      thrust::sort_by_key(keys_begin, keys_end, values_begin, thrust::greater<T>());
+    }
+  }
+
+#endif
+
+  template <typename T,
+            typename U,
+            class StorageT,
+            class StorageU,
+            class BinaryCompare,
+            typename ValidKeys,
+            typename ValidValues>
+  VTKM_CONT static void SortByKeyImpl(vtkm::cont::ArrayHandle<T, StorageT>& keys,
+                                      vtkm::cont::ArrayHandle<U, StorageU>& values,
+                                      BinaryCompare binary_compare,
+                                      ValidKeys,
+                                      ValidValues)
+  {
+    // Default to general algorithm
+    Superclass::SortByKey(keys, values, binary_compare);
+  }
+
+public:
+  template <typename T, typename U, class StorageT, class StorageU>
+  VTKM_CONT static void SortByKey(vtkm::cont::ArrayHandle<T, StorageT>& keys,
+                                  vtkm::cont::ArrayHandle<U, StorageU>& values)
+  {
+    // Make sure not to use the general algorithm here since
+    // it will use Sort algorithm instead of SortByKey
+    SortByKey(keys, values, internal::DefaultCompareFunctor());
+  }
+
+  template <typename T, typename U, class StorageT, class StorageU, class BinaryCompare>
+  VTKM_CONT static void SortByKey(vtkm::cont::ArrayHandle<T, StorageT>& keys,
+                                  vtkm::cont::ArrayHandle<U, StorageU>& values,
+                                  BinaryCompare binary_compare)
+  {
+    // If T or U are not scalar types, or the BinaryCompare is not supported
+    // then the general algorithm is called, otherwise we will run thrust
+    SortByKeyImpl(keys,
+                  values,
+                  binary_compare,
+                  typename std::is_scalar<T>::type{},
+                  typename std::is_scalar<U>::type{});
+  }
+
+  //----------------------------------------------------------------------------
+
   VTKM_CONT static void Synchronize()
   {
     vtkm::cont::kokkos::internal::GetExecutionSpaceInstance().fence();

vtkm/internal/CMakeLists.txt

@@ -21,6 +21,7 @@ set(VTKM_USE_64BIT_IDS ${VTKm_USE_64BIT_IDS})
 
 set(VTKM_ENABLE_CUDA ${VTKm_ENABLE_CUDA})
 set(VTKM_ENABLE_KOKKOS ${VTKm_ENABLE_KOKKOS})
+set(VTKM_ENABLE_KOKKOS_THRUST ${VTKm_ENABLE_KOKKOS_THRUST})
 set(VTKM_ENABLE_OPENMP ${VTKm_ENABLE_OPENMP})
 set(VTKM_ENABLE_TBB ${VTKm_ENABLE_TBB})
 

vtkm/internal/Configure.h.in

@@ -300,6 +300,10 @@
 #ifndef VTKM_KOKKOS_HIP
 #cmakedefine VTKM_KOKKOS_HIP
 #endif
+// Mark if Kokkos algorithms should use thrust
+#if defined(VTKM_KOKKOS_HIP) || defined(VTKM_KOKKOS_CUDA)
+#cmakedefine VTKM_ENABLE_KOKKOS_THRUST
+#endif
 
 //Mark if we are building with MPI enabled.
 #cmakedefine VTKM_ENABLE_MPI