Add support for rocThrust in sort when using HIP #6793

Rombur · 2024-02-08T15:13:49Z

This PR add support for rocThrust. Unlike Thrust, rocThrust is not always installed with HIP. For instance, the docker images do not contain rocThrust unless you use the -complete tag. I will note that the 3.0 release of rocThrust is not backward compatible with the 2.X series. This is not an issue for us because sort is the same across version but it is problematic if we decide to require a minimum version of rocThrust because then the 3.0 version will be rejected. Since we are not affected by the break of backward compatibility, I've decided to forego the minimum version requirements. In practice, this means that the minimum version of rocThrust is the minimum version that is compatible with ROCm 5.2

Once this PR is merged, I'll open a PR to document the new option that enables/disables rocThrust.

dalg24 · 2024-02-08T16:26:04Z

cmake/KokkosCore_config.h.in

@@ -62,6 +62,7 @@
 #cmakedefine KOKKOS_ENABLE_LIBDL
 #cmakedefine KOKKOS_ENABLE_LIBQUADMATH
 #cmakedefine KOKKOS_ENABLE_ONEDPL
+#cmakedefine KOKKOS_ENABLE_ROCTHRUST


You might want to add that setting to HIP::print_configuration()

@masterleinad same comment about KOKKOS_ENABLE_ONEDPL and SYCL

Let me do that after #6758 is merged.

dalg24 · 2024-02-08T16:28:50Z

containers/src/CMakeLists.txt

@@ -28,6 +28,9 @@ KOKKOS_LIB_INCLUDE_DIRECTORIES(kokkoscontainers
  ${CMAKE_CURRENT_BINARY_DIR}
  ${CMAKE_CURRENT_SOURCE_DIR}
 )
+
+KOKKOS_LINK_TPL(kokkoscontainers PUBLIC ROCTHRUST)


Why link containers?

My bad, I meant to add this in algorithm.

dalg24 · 2024-02-08T16:29:30Z

cmake/Modules/FindTPLROCTHRUST.cmake

@@ -0,0 +1,2 @@
+FIND_PACKAGE(rocthrust REQUIRED)
+KOKKOS_CREATE_IMPORTED_TPL(ROCTHRUST INTERFACE LINK_LIBRARIES roc::rocthrust)


Do we still have the arborx/ArborX#652 issue where the exported rocThrust package potentially adds flags to compile for all GPU architecture it knows about?

At least for the version shipped with ROCm 5.7, it works fine. I see only one architecture.

We need to find out for sure because IIRC compile time blows up

I looked at another PR it looks like the compile time is indeed 2x more for both 5.2 and 5.6 so we need to handle this and set the GPU_TARGETS before attempting to find rocThrust.

masterleinad

Looks OK to me.

cmake/Modules/FindTPLROCTHRUST.cmake

aprokop · 2024-02-12T21:51:16Z

Just wanted to confirm: will we be able to write something like this downstream and will it work? Including linking to rocthrust as long as we link to Kokkos::kokkos.

if(Kokkos_ENABLE_HIP)
  if(Kokkos_VERSION VERSION_LESS 4.3)
    # Do something
  else
    set(ARBORX_ENABLE_ROCTHRUST ${Kokkos_ENABLE_ROCTHRUST})
  endif()
endif()

dalg24 · 2024-02-12T22:28:44Z

containers/src/CMakeLists.txt

@@ -28,6 +28,7 @@ KOKKOS_LIB_INCLUDE_DIRECTORIES(kokkoscontainers
  ${CMAKE_CURRENT_BINARY_DIR}
  ${CMAKE_CURRENT_SOURCE_DIR}
 )
+


Please undo

dalg24 · 2024-02-13T05:10:28Z

cmake/Modules/FindTPLROCTHRUST.cmake

@@ -0,0 +1,2 @@
+FIND_PACKAGE(rocthrust REQUIRED)
+KOKKOS_CREATE_IMPORTED_TPL(ROCTHRUST INTERFACE LINK_LIBRARIES roc::rocthrust)


We need to find out for sure because IIRC compile time blows up

stanmoore1 · 2024-02-16T20:07:04Z

This would be useful for LAMMPS; Kokkos::BinSort isn't cutting it at large scale on Frontier (lammps/lammps#4075) .

dalg24

Just realized my comment was pending. Sorry about that.
I had mentioned the issue to Bruno in private.

dalg24 · 2024-02-13T05:16:53Z

cmake/Modules/FindTPLROCTHRUST.cmake

@@ -0,0 +1,2 @@
+FIND_PACKAGE(rocthrust REQUIRED)
+KOKKOS_CREATE_IMPORTED_TPL(ROCTHRUST INTERFACE LINK_LIBRARIES roc::rocthrust)


I looked at another PR it looks like the compile time is indeed 2x more for both 5.2 and 5.6 so we need to handle this and set the GPU_TARGETS before attempting to find rocThrust.

stanmoore1 · 2024-02-28T23:39:33Z

This will be in 4.3 correct? Would be really nice if it supported sort_by_key too...

dalg24 · 2024-02-28T23:41:41Z

This will be in 4.3 correct? Would be really nice if it supported sort_by_key too...

This is the plan yes. Bruno is swamped with something else but will get back to it next week. #6793 (comment) needs to be addressed and then we will merge.

aprokop · 2024-02-29T05:54:06Z

Rebased on develop to allow integration of rocThrust's sort_by_key in this PR.

…re when enabling rocthrust

masterleinad · 2024-03-06T18:31:11Z

Make sure to take https://github.com/kokkos/kokkos/pull/6856/files into account and define

template <class Layout>
inline constexpr bool sort_on_device_v<Kokkos::HIP, Layout> = true;

edit: I see that sort_by_key isn't part of this pull request (anymore?) were you intending to do that in a follow-up instead?

Rombur · 2024-03-07T01:00:35Z

This PR is to add support for rocThrust. sort_by_key will be its own PR.

dalg24 reviewed Feb 8, 2024

View reviewed changes

masterleinad approved these changes Feb 8, 2024

View reviewed changes

dalg24 reviewed Feb 8, 2024

View reviewed changes

cmake/Modules/FindTPLROCTHRUST.cmake Show resolved Hide resolved

masterleinad mentioned this pull request Feb 8, 2024

SYCL: Improve print_configuration #6795

Merged

aprokop mentioned this pull request Feb 9, 2024

Implement sort_by_key #6801

Merged

Rombur force-pushed the roc_thrust branch from 275b57f to 800fb6e Compare February 12, 2024 15:41

dalg24 reviewed Feb 13, 2024

View reviewed changes

stanmoore1 mentioned this pull request Feb 14, 2024

[BUG] Major performance regression in Special.build() component of replicate at large scale (512 nodes+) lammps/lammps#4075

Closed

dalg24 reviewed Feb 16, 2024

View reviewed changes

ajpowelsnl mentioned this pull request Feb 28, 2024

Release Themes for 2024 #6804

Open

aprokop force-pushed the roc_thrust branch from 800fb6e to 4b495d7 Compare February 29, 2024 05:53

dalg24 mentioned this pull request Feb 29, 2024

Bug in Kokkos::sort() fallback implementation when using a compare functor #6852

Open

Rombur added 3 commits March 6, 2024 09:27

Use rocthrust in sort when using HIP

2d106d5

Fix reviewer's comments

af3ddca

Make sure that we don't compile Kokkos for every supported architectu…

da500f4

…re when enabling rocthrust

Rombur force-pushed the roc_thrust branch from 4b495d7 to da500f4 Compare March 6, 2024 18:19

masterleinad approved these changes Mar 7, 2024

View reviewed changes

Export ROCTHRUST as a Kokkos dependency

2f5ecf1

dalg24 approved these changes Mar 7, 2024

View reviewed changes

dalg24 merged commit 35ad698 into kokkos:develop Mar 11, 2024
33 checks passed

This was referenced Mar 11, 2024

CHANGELOG: 4.3.0 #6519

Closed

Add documentation about the new CMake keyword: rocthrust kokkos/kokkos-core-wiki#506

Merged

masterleinad mentioned this pull request Mar 14, 2024

Fix linking with rothrust in downstream applications #6871

Merged

Rombur deleted the roc_thrust branch May 23, 2024 15:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for rocThrust in sort when using HIP #6793

Add support for rocThrust in sort when using HIP #6793

Rombur commented Feb 8, 2024

dalg24 Feb 8, 2024

masterleinad Feb 8, 2024

dalg24 Feb 8, 2024

Rombur Feb 8, 2024

dalg24 Feb 8, 2024

Rombur Feb 8, 2024

dalg24 Feb 13, 2024

dalg24 Feb 13, 2024

masterleinad left a comment

aprokop commented Feb 12, 2024 •

edited

dalg24 Feb 12, 2024

dalg24 Feb 13, 2024

stanmoore1 commented Feb 16, 2024 •

edited

dalg24 left a comment

dalg24 Feb 13, 2024

stanmoore1 commented Feb 28, 2024

dalg24 commented Feb 28, 2024

aprokop commented Feb 29, 2024

masterleinad commented Mar 6, 2024 •

edited

Rombur commented Mar 7, 2024

		@@ -0,0 +1,2 @@
		FIND_PACKAGE(rocthrust REQUIRED)
		KOKKOS_CREATE_IMPORTED_TPL(ROCTHRUST INTERFACE LINK_LIBRARIES roc::rocthrust)

Add support for rocThrust in sort when using HIP #6793

Add support for rocThrust in sort when using HIP #6793

Conversation

Rombur commented Feb 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masterleinad left a comment

Choose a reason for hiding this comment

aprokop commented Feb 12, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stanmoore1 commented Feb 16, 2024 • edited

dalg24 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stanmoore1 commented Feb 28, 2024

dalg24 commented Feb 28, 2024

aprokop commented Feb 29, 2024

masterleinad commented Mar 6, 2024 • edited

Rombur commented Mar 7, 2024

aprokop commented Feb 12, 2024 •

edited

stanmoore1 commented Feb 16, 2024 •

edited

masterleinad commented Mar 6, 2024 •

edited