lammps: GPU/Kokkos package updates (#35885)

2 files changed, 95 insertions, 1 deletions

diff --git a/var/spack/repos/builtin/packages/lammps/hip_cmake.patch b/var/spack/repos/builtin/packages/lammps/hip_cmake.patch
new file mode 100644
index 0000000000..cf2cbbea84
--- /dev/null
+++ b/var/spack/repos/builtin/packages/lammps/hip_cmake.patch
@@ -0,0 +1,58 @@
+From b11049ba1a5fa22ff575f4d4afa6579973425962 Mon Sep 17 00:00:00 2001
+From: Richard Berger <richard.berger@outlook.com>
+Date: Sun, 5 Mar 2023 19:03:38 -0700
+Subject: [PATCH] CMake: Use hip::host and hip::hipcub targets
+
+---
+ cmake/Modules/Packages/GPU.cmake | 13 +++----------
+ 1 file changed, 3 insertions(+), 10 deletions(-)
+
+diff --git a/cmake/Modules/Packages/GPU.cmake b/cmake/Modules/Packages/GPU.cmake
+index 8ac1decc86..21d046606f 100644
+--- a/cmake/Modules/Packages/GPU.cmake
++++ b/cmake/Modules/Packages/GPU.cmake
+@@ -412,7 +412,8 @@ elseif(GPU_API STREQUAL "HIP")
+       set_property(TARGET gpu PROPERTY CXX_STANDARD 14)
+     endif()
+     # add hipCUB
+-    target_include_directories(gpu PRIVATE ${HIP_ROOT_DIR}/../include)
++    find_package(hipcub REQUIRED)
++    target_link_libraries(gpu PRIVATE hip::hipcub)
+     target_compile_definitions(gpu PRIVATE -DUSE_HIP_DEVICE_SORT)
+ 
+     if(HIP_PLATFORM STREQUAL "nvcc")
+@@ -461,30 +462,22 @@ elseif(GPU_API STREQUAL "HIP")
+ 
+   add_executable(hip_get_devices ${LAMMPS_LIB_SOURCE_DIR}/gpu/geryon/ucl_get_devices.cpp)
+   target_compile_definitions(hip_get_devices PRIVATE -DUCL_HIP)
+-  target_link_libraries(hip_get_devices hip::host)
++  target_link_libraries(hip_get_devices PRIVATE hip::host)
+ 
+   if(HIP_PLATFORM STREQUAL "nvcc")
+     target_compile_definitions(gpu PRIVATE -D__HIP_PLATFORM_NVCC__)
+-    target_include_directories(gpu PRIVATE ${HIP_ROOT_DIR}/../include)
+     target_include_directories(gpu PRIVATE ${CUDA_INCLUDE_DIRS})
+     target_link_libraries(gpu PRIVATE ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY})
+ 
+     target_compile_definitions(hip_get_devices PRIVATE -D__HIP_PLATFORM_NVCC__)
+-    target_include_directories(hip_get_devices PRIVATE ${HIP_ROOT_DIR}/include)
+     target_include_directories(hip_get_devices PRIVATE ${CUDA_INCLUDE_DIRS})
+     target_link_libraries(hip_get_devices PRIVATE ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY})
+   elseif(HIP_PLATFORM STREQUAL "hcc")
+     target_compile_definitions(gpu PRIVATE -D__HIP_PLATFORM_HCC__)
+-    target_include_directories(gpu PRIVATE ${HIP_ROOT_DIR}/../include)
+-
+     target_compile_definitions(hip_get_devices PRIVATE -D__HIP_PLATFORM_HCC__)
+-    target_include_directories(hip_get_devices PRIVATE ${HIP_ROOT_DIR}/../include)
+   elseif(HIP_PLATFORM STREQUAL "amd")
+     target_compile_definitions(gpu PRIVATE -D__HIP_PLATFORM_AMD__)
+-    target_include_directories(gpu PRIVATE ${HIP_ROOT_DIR}/../include)
+-
+     target_compile_definitions(hip_get_devices PRIVATE -D__HIP_PLATFORM_AMD__)
+-    target_include_directories(hip_get_devices PRIVATE ${HIP_ROOT_DIR}/../include)
+   endif()
+ 
+   target_link_libraries(lammps PRIVATE gpu)
+-- 
+2.39.2
+
diff --git a/var/spack/repos/builtin/packages/lammps/package.py b/var/spack/repos/builtin/packages/lammps/package.py
index 8e2b991678..66681bc98e 100644
--- a/var/spack/repos/builtin/packages/lammps/package.py
+++ b/var/spack/repos/builtin/packages/lammps/package.py
@@ -9,7 +9,7 @@ import archspec
 from spack.package import *
 
 
-class Lammps(CMakePackage, CudaPackage):
+class Lammps(CMakePackage, CudaPackage, ROCmPackage):
     """LAMMPS stands for Large-scale Atomic/Molecular Massively
     Parallel Simulator.
     """
@@ -513,6 +513,14 @@ class Lammps(CMakePackage, CudaPackage):
         values=("single", "double"),
         multi=False,
     )
+    variant(
+        "gpu_precision",
+        default="mixed",
+        when="~kokkos",
+        description="Select GPU precision (used by GPU package)",
+        values=("double", "mixed", "single"),
+        multi=False,
+    )
 
     depends_on("mpi", when="+mpi")
     depends_on("mpi", when="+mpiio")
@@ -561,6 +569,17 @@ class Lammps(CMakePackage, CudaPackage):
     depends_on("n2p2+shared", when="+lib ^n2p2")
     depends_on("vtk", when="+user-vtk")
     depends_on("vtk", when="+vtk")
+    depends_on("hipcub", when="~kokkos +rocm")
+    depends_on("llvm-amdgpu +openmp", when="+rocm +openmp", type="build")
+
+    # propagate CUDA and ROCm architecture when +kokkos
+    for arch in CudaPackage.cuda_arch_values:
+        depends_on("kokkos+cuda cuda_arch=%s" % arch, when="+kokkos+cuda cuda_arch=%s" % arch)
+
+    for arch in ROCmPackage.amdgpu_targets:
+        depends_on(
+            "kokkos+rocm amdgpu_target=%s" % arch, when="+kokkos+rocm amdgpu_target=%s" % arch
+        )
 
     depends_on("googletest", type="test")
     depends_on("libyaml", type="test")
@@ -568,6 +587,7 @@ class Lammps(CMakePackage, CudaPackage):
     extends("python", when="+python")
 
     conflicts("+cuda", when="+opencl")
+    conflicts("+rocm", when="+opencl")
     conflicts("+body", when="+poems@:20180628")
     conflicts("+latte", when="@:20170921")
     conflicts("+python", when="~lib")
@@ -605,6 +625,11 @@ class Lammps(CMakePackage, CudaPackage):
         when="^adios2+mpi",
         msg="With +adios, mpi setting for adios2 and lammps must be the same",
     )
+    conflicts(
+        "~kokkos+rocm",
+        when="@:20220602",
+        msg="ROCm builds of the GPU package not maintained prior to version 20220623",
+    )
 
     patch("lib.patch", when="@20170901")
     patch("660.patch", when="@20170922")
@@ -614,6 +639,7 @@ class Lammps(CMakePackage, CudaPackage):
         sha256="e6f1b62bbfdc79d632f4cea98019202d0dd25aa4ae61a70df1164cb4f290df79",
         when="@20200721 +cuda",
     )
+    patch("hip_cmake.patch", when="@20220623:20221222 ~kokkos+rocm")
 
     root_cmakelists_dir = "cmake"
 
@@ -641,6 +667,7 @@ class Lammps(CMakePackage, CudaPackage):
             if "+cuda" in spec:
                 args.append(self.define("PKG_GPU", True))
                 args.append(self.define("GPU_API", "cuda"))
+                args.append(self.define_from_variant("GPU_PREC", "gpu_precision"))
                 cuda_arch = spec.variants["cuda_arch"].value
                 if cuda_arch != "none":
                     args.append(self.define("GPU_ARCH", "sm_{0}".format(cuda_arch[0])))
@@ -650,6 +677,12 @@ class Lammps(CMakePackage, CudaPackage):
                 args.append(self.define("USE_STATIC_OPENCL_LOADER", False))
                 args.append(self.define("PKG_GPU", True))
                 args.append(self.define("GPU_API", "opencl"))
+                args.append(self.define_from_variant("GPU_PREC", "gpu_precision"))
+            elif "+rocm" in spec:
+                args.append(self.define("PKG_GPU", True))
+                args.append(self.define("GPU_API", "hip"))
+                args.append(self.define_from_variant("GPU_PREC", "gpu_precision"))
+                args.append(self.define_from_variant("HIP_ARCH", "amdgpu_target"))
             else:
                 args.append(self.define("PKG_GPU", False))
 
@@ -712,6 +745,9 @@ class Lammps(CMakePackage, CudaPackage):
             args.append(self.define("DOWNLOAD_N2P2", False))
             args.append(self.define("N2P2_DIR", self.spec["n2p2"].prefix))
 
+        if "+rocm" in spec:
+            args.append(self.define("CMAKE_CXX_COMPILER", spec["hip"].hipcc))
+
         return args
 
     def setup_run_environment(self, env):