diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 502e9306c..421d4af3a 100644 --- a/.gitlab-ci.yml +++ b/.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 diff --git a/.gitlab/ci/docker/ubuntu2004/kokkos-hip/Dockerfile b/.gitlab/ci/docker/ubuntu2004/kokkos-hip/Dockerfile index ebb8501e6..8cffc8b17 100644 --- a/.gitlab/ci/docker/ubuntu2004/kokkos-hip/Dockerfile +++ b/.gitlab/ci/docker/ubuntu2004/kokkos-hip/Dockerfile @@ -14,6 +14,7 @@ RUN apt update && \ ninja-build \ rsync \ ssh \ + rocthrust-dev \ && \ apt clean diff --git a/CMake/VTKmDeviceAdapters.cmake b/CMake/VTKmDeviceAdapters.cmake index bbecd18fc..fb13d0bf8 100644 --- a/CMake/VTKmDeviceAdapters.cmake +++ b/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) diff --git a/CMakeLists.txt b/CMakeLists.txt index 725199cc4..0065590b1 100644 --- a/CMakeLists.txt +++ b/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. diff --git a/vtkm/cont/kokkos/internal/DeviceAdapterAlgorithmKokkos.h b/vtkm/cont/kokkos/internal/DeviceAdapterAlgorithmKokkos.h index b5b6e6a26..606309d06 100644 --- a/vtkm/cont/kokkos/internal/DeviceAdapterAlgorithmKokkos.h +++ b/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 +#include +#endif namespace vtkm { @@ -771,6 +779,88 @@ public: SortImpl(values, comp, typename std::is_scalar::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 + VTKM_CONT static std::enable_if_t<(std::is_same::value || + std::is_same::value)> + SortByKeyImpl(vtkm::cont::ArrayHandle& keys, + vtkm::cont::ArrayHandle& 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 keys_view(keys_portal.GetArray(), + keys_portal.GetNumberOfValues()); + kokkos::internal::KokkosViewExec values_view(values_portal.GetArray(), + values_portal.GetNumberOfValues()); + + thrust::device_ptr keys_begin(keys_view.data()); + thrust::device_ptr keys_end(keys_view.data() + keys_view.size()); + thrust::device_ptr values_begin(values_view.data()); + + if (std::is_same::value) + { + thrust::sort_by_key(keys_begin, keys_end, values_begin, thrust::less()); + } + else + { + thrust::sort_by_key(keys_begin, keys_end, values_begin, thrust::greater()); + } + } + +#endif + + template + VTKM_CONT static void SortByKeyImpl(vtkm::cont::ArrayHandle& keys, + vtkm::cont::ArrayHandle& values, + BinaryCompare binary_compare, + ValidKeys, + ValidValues) + { + // Default to general algorithm + Superclass::SortByKey(keys, values, binary_compare); + } + +public: + template + VTKM_CONT static void SortByKey(vtkm::cont::ArrayHandle& keys, + vtkm::cont::ArrayHandle& 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 + VTKM_CONT static void SortByKey(vtkm::cont::ArrayHandle& keys, + vtkm::cont::ArrayHandle& 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::type{}, + typename std::is_scalar::type{}); + } + + //---------------------------------------------------------------------------- + VTKM_CONT static void Synchronize() { vtkm::cont::kokkos::internal::GetExecutionSpaceInstance().fence(); diff --git a/vtkm/internal/CMakeLists.txt b/vtkm/internal/CMakeLists.txt index 88f7622f8..8e6989c3a 100755 --- a/vtkm/internal/CMakeLists.txt +++ b/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}) diff --git a/vtkm/internal/Configure.h.in b/vtkm/internal/Configure.h.in index 39b578971..f2caf2a64 100644 --- a/vtkm/internal/Configure.h.in +++ b/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