Adding CUDA and OPENMP to LAMMPS experiment class (#570)

* progress on lammps

* venado configs

* lammps cuda

* Fix super invocations

* Fix mpi

* cufft

* venado specs

* changing blas name, allowing no extra args

* adding updated dryruns

* removing duplicate dryrun

---------

Co-authored-by: august-knox <[email protected]>
Co-authored-by: Riyaz Haque <[email protected]>
Co-authored-by: Riyaz Haque <[email protected]>

.github/workflows/run.yml

@@ -720,3 +720,39 @@ jobs:
             --disable-progress-bar \
             --disable-logger \
             workspace setup --dry-run
+
+      - name: Dry run dynamic lammps/openmp with dynamic tioga
+        run: |
+          system_id=$(./bin/benchpark system id ./tioga-system1)
+          ./bin/benchpark experiment init --dest=lammps-openmp-tioga lammps+openmp
+          ./bin/benchpark setup ./lammps-openmp-tioga ./tioga-system1 workspace/
+          . workspace/setup.sh
+          ramble \
+            --workspace-dir workspace/lammps-openmp-tioga/$system_id/workspace \
+            --disable-progress-bar \
+            --disable-logger \
+            workspace setup --dry-run
+
+      - name: Dry run dynamic lammps/openmp with dynamic sierra
+        run: |
+          system_id=$(./bin/benchpark system id ./sierra-system1)
+          ./bin/benchpark experiment init --dest=lammps-openmp-sierra lammps+openmp
+          ./bin/benchpark setup ./lammps-openmp-sierra ./sierra-system1 workspace/
+          . workspace/setup.sh
+          ramble \
+            --workspace-dir workspace/lammps-openmp-sierra/$system_id/workspace \
+            --disable-progress-bar \
+            --disable-logger \
+            workspace setup --dry-run
+
+      - name: Dry run dynamic lammps/openmp with dynamic ruby
+        run: |
+          system_id=$(./bin/benchpark system id ./ruby-system)
+          ./bin/benchpark experiment init --dest=lammps-openmp-ruby lammps+openmp
+          ./bin/benchpark setup ./lammps-openmp-ruby ./ruby-system workspace/
+          . workspace/setup.sh
+          ramble \
+            --workspace-dir workspace/lammps-openmp-ruby/$system_id/workspace \
+            --disable-progress-bar \
+            --disable-logger \
+            workspace setup --dry-run

experiments/lammps/experiment.py

@@ -6,12 +6,14 @@
 from benchpark.directives import variant
 from benchpark.experiment import Experiment
 from benchpark.openmp import OpenMPExperiment
+from benchpark.cuda import CudaExperiment
 from benchpark.rocm import ROCmExperiment
 
 
 class Lammps(
     Experiment,
     OpenMPExperiment,
+    CudaExperiment,
     ROCmExperiment,
 ):
     variant(
@@ -27,17 +29,35 @@ class Lammps(
         description="app version",
     )
 
+    variant(
+        "gpu-aware-mpi",
+        default=True,
+        values=(True, False),
+        when=("+cuda" or "+rocm"),
+        description="Enable GPU-aware MPI",
+    )
+
     def compute_applications_section(self):
         if self.spec.satisfies("+openmp"):
             problem_sizes = {"x": 8, "y": 8, "z": 8}
             kokkos_mode = "t {n_threads_per_proc}"
             kokkos_gpu_aware = "off"
             kokkos_comm = "host"
-        elif self.spec.satisfies("+rocm"):
+        elif self.spec.satisfies("+rocm") or self.spec.satisfies("+cuda"):
             problem_sizes = {"x": 20, "y": 40, "z": 32}
             kokkos_mode = "g 1"
-            kokkos_gpu_aware = "on"
+            kokkos_gpu_aware = "on" if self.spec.satisfies("+rocm") else "off"
+            kokkos_comm = "device"
+        elif self.spec.satisfies("+cuda"):
+            problem_sizes = {"x": 20, "y": 20, "z": 16}
+            kokkos_mode = "g 1"
+            kokkos_gpu_aware = "on" if self.spec.satisfies("+cuda") else "off"
             kokkos_comm = "device"
+        else:
+            problem_sizes = {"x": 8, "y": 8, "z": 8}
+            kokkos_mode = "t {n_threads_per_proc}"
+            kokkos_gpu_aware = "off"
+            kokkos_comm = "host"
 
         for nk, nv in problem_sizes.items():
             self.add_experiment_variable(nk, nv, True)
@@ -52,6 +72,14 @@ def compute_applications_section(self):
             self.add_experiment_variable("n_nodes", 8, True)
             self.add_experiment_variable("n_ranks_per_node", 8, True)
             self.add_experiment_variable("n_gpus", 64, True)
+        elif self.spec.satisfies("+cuda"):
+            self.add_experiment_variable("n_nodes", 4, True)
+            self.add_experiment_variable("n_ranks_per_node", 4, True)
+            self.add_experiment_variable("n_gpus", 16, True)
+        else:
+            self.add_experiment_variable("n_nodes", 1, True)
+            self.add_experiment_variable("n_ranks_per_node", 36, True)
+            self.add_experiment_variable("n_threads_per_proc", 1, True)
 
         self.add_experiment_variable("timesteps", 100, False)
         self.add_experiment_variable("input_file", "{input_path}/in.reaxc.hns", False)
@@ -70,21 +98,25 @@ def compute_spack_section(self):
         system_specs = {}
         system_specs["compiler"] = "default-compiler"
         system_specs["mpi"] = "default-mpi"
-        if self.spec.satisfies("+rocm"):
-            system_specs["rocm_arch"] = "{rocm_arch}"
-            system_specs["blas"] = "blas-rocm"
+        system_specs["blas"] = "blas"
 
         # set package spack specs
-        if self.spec.satisfies("+rocm"):
-            # empty package_specs value implies external package
-            self.add_spack_spec(system_specs["blas"])
         # empty package_specs value implies external package
         self.add_spack_spec(system_specs["mpi"])
 
+        if self.spec.satisfies("+cuda"):
+            system_specs["cuda_version"] = "{default_cuda_version}"
+            system_specs["cuda_arch"] = "{cuda_arch}"
+        elif self.spec.satisfies("+rocm"):
+            system_specs["rocm_arch"] = "{rocm_arch}"
+
+        # empty package_specs value implies external package
+        self.add_spack_spec(system_specs["blas"])
+
         self.add_spack_spec(
             self.name,
             [
-                f"lammps@{app_version} +opt+manybody+molecule+kspace+rigid+kokkos+asphere+dpd-basic+dpd-meso+dpd-react+dpd-smooth+reaxff lammps_sizes=bigbig ",
+                f"lammps@{app_version} +mpi+opt+manybody+molecule+kspace+rigid+kokkos+asphere+dpd-basic+dpd-meso+dpd-react+dpd-smooth+reaxff lammps_sizes=bigbig ",
                 system_specs["compiler"],
             ],
         )

repo/cray-mpich/package.py

@@ -13,7 +13,7 @@ class CrayMpich(BuiltinCM):
 
     @property
     def libs(self):
-        libs = super(CrayMpich, self).libs
+        libs = super().libs
 
         if self.spec.satisfies("+gtl"):
             gtl_lib_prefix = self.spec.extra_attributes["gtl_lib_path"]
@@ -24,7 +24,7 @@ def libs(self):
 
     def setup_run_environment(self, env):
 
-        super(CrayMpich, self).setup_run_environment(env)
+        super().setup_run_environment(env)
 
         if self.spec.satisfies("+gtl"):
             env.set("MPICH_GPU_SUPPORT_ENABLED", "1")

repo/kokkos/package.py

@@ -0,0 +1,15 @@
+# Copyright 2023 Lawrence Livermore National Security, LLC and other
+# Benchpark Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from spack.package import *
+from spack.pkg.builtin.kokkos import Kokkos as BuiltinKokkos
+
+
+class Kokkos(BuiltinKokkos):
+  flag_handler = build_system_flags
+
+  def setup_build_environment(self, env):
+    if "+cuda" in self.spec:
+      env.set("NVCC_APPEND_FLAGS", "-allow-unsupported-compiler")

repo/lammps/package.py

@@ -9,13 +9,17 @@
 
 class Lammps(BuiltinLammps):
 
-  depends_on("kokkos+openmp", when="+openmp")
+  depends_on("kokkos+openmp cxxstd=17", when="+openmp")
   depends_on("kokkos+rocm", when="+rocm")
-  depends_on("kokkos+cuda", when="+cuda")
+  depends_on("[email protected] +cuda cxxstd=17", when="+cuda")
+
+  conflicts("+rocm", when="+cuda")
+  conflicts("+cuda", when="+rocm")
 
-  def setup_run_environment(self, env):
+  flag_handler = build_system_flags
 
-    super(BuiltinLammps, self).setup_run_environment(env)
+  def setup_run_environment(self, env):
+    super().setup_run_environment(env)
 
     if self.compiler.extra_rpaths:
       for rpath in self.compiler.extra_rpaths:
@@ -24,7 +28,13 @@ def setup_run_environment(self, env):
   def setup_build_environment(self, env):
     super().setup_build_environment(env)
 
-    spec = self.spec
-    if "+mpi" in spec:
-      if spec["mpi"].extra_attributes and "ldflags" in spec["mpi"].extra_attributes:
-        env.append_flags("LDFLAGS", spec["mpi"].extra_attributes["ldflags"])
+    if "+cuda" in self.spec:
+      env.set("NVCC_APPEND_FLAGS", "-allow-unsupported-compiler")
+
+  def cmake_args(self):
+    args = super().cmake_args()
+    args.append(f"-DMPI_CXX_LINK_FLAGS='{self.spec['mpi'].libs.ld_flags}'")
+    args.append(f"-DMPI_C_COMPILER='{self.spec['mpi'].mpicc}'")
+    args.append(f"-DMPI_CXX_COMPILER={self.spec['mpi'].mpicxx}")
+
+    return args