Feature/gromacs (#476)

* gromacs rocm * gromacs rocm * gromacs openmp * gromacs cuda * lint * lint * lint * system changes * Remove dead code * Remove rocm consistency package, update llvlm version * Add dry runs * workflows * workflows * workflows * workflows --------- Co-authored-by: Riyaz Haque <[email protected]> Co-authored-by: pearce8 <[email protected]>
LLNL · Dec 15, 2024 · 7a7380a · 7a7380a
1 parent 9fe414b
commit 7a7380a
Show file tree

Hide file tree

Showing 15 changed files with 359 additions and 79 deletions.
diff --git a/.github/workflows/run.yml b/.github/workflows/run.yml
@@ -222,6 +222,45 @@ jobs:
             --disable-logger \
             workspace setup --dry-run
 
+      - name: Dry run dynamic gromacs/openmp with dynamic Ruby
+        run: |
+          ./bin/benchpark system init --dest=cluster-system llnl-cluster compiler=gcc lapack=intel-oneapi-mkl blas=intel-oneapi-mkl
+          system_id=$(./bin/benchpark system id ./cluster-system)
+          ./bin/benchpark experiment init --dest=gromacs-openmp gromacs +openmp~cuda~rocm gpu-aware-mpi=off
+          ./bin/benchpark setup ./gromacs-openmp ./cluster-system workspace/
+          . workspace/setup.sh
+          ramble \
+            --workspace-dir "workspace/gromacs-openmp/$system_id/workspace" \
+            --disable-progress-bar \
+            --disable-logger \
+            workspace setup --dry-run
+
+      - name: Dry run dynamic gromacs/cuda with dynamic Sierra
+        run: |
+          ./bin/benchpark system init --dest=sierra-system1 llnl-sierra cuda=11-8-0 compiler=clang-ibm lapack=cusolver blas=cublas
+          system_id=$(./bin/benchpark system id ./sierra-system1)
+          ./bin/benchpark experiment init --dest=gromacs-cuda gromacs +openmp+cuda~rocm gpu-aware-mpi=on
+          ./bin/benchpark setup ./gromacs-cuda ./sierra-system1 workspace/
+          . workspace/setup.sh
+          ramble \
+            --workspace-dir "workspace/gromacs-cuda/$system_id/workspace" \
+            --disable-progress-bar \
+            --disable-logger \
+            workspace setup --dry-run
+
+      - name: Dry run dynamic gromacs/rocm with dynamic Tioga
+        run: |
+          ./bin/benchpark system init --dest=tioga-system1 llnl-elcapitan rocm=5.5.1 compiler=cce +gtl lapack=intel-oneapi-mkl blas=intel-oneapi-mkl
+          system_id=$(./bin/benchpark system id ./tioga-system1)
+          ./bin/benchpark experiment init --dest=gromacs-rocm gromacs +openmp+rocm~cuda gpu-aware-mpi=force
+          ./bin/benchpark setup ./gromacs-rocm ./tioga-system1 workspace/
+          . workspace/setup.sh
+          ramble \
+            --workspace-dir "workspace/gromacs-rocm/$system_id/workspace" \
+            --disable-progress-bar \
+            --disable-logger \
+            workspace setup --dry-run
+
       - name: Dry run dynamic laghos/mpi-only on LLNL-Magma-Penguin-icelake-OmniPath with allocation modifier
         run: |
           ./bin/benchpark experiment init --dest=laghos-mpi-only laghos

diff --git a/experiments/gromacs/experiment.py b/experiments/gromacs/experiment.py
@@ -0,0 +1,119 @@
+# 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 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 Gromacs(
+    Experiment,
+    OpenMPExperiment,
+    CudaExperiment,
+    ROCmExperiment,
+):
+    variant(
+        "workload",
+        default="water_gmx50_adac",
+        description="workload name",
+    )
+
+    variant(
+        "version",
+        default="2024",
+        values=("2024", "2023.3"),
+        description="app version",
+    )
+
+    # off: turn off GPU-aware MPI
+    # on: turn on, but allow groamcs to disable it if GPU-aware MPI is not supported
+    # force: turn on and force gromacs to use GPU-aware MPI. May result in error if unsupported
+    variant(
+        "gpu-aware-mpi",
+        default="on",
+        values=("on", "off", "force"),
+        description="Use GPU-aware MPI",
+    )
+
+    def compute_applications_section(self):
+        if self.spec.satisfies("+openmp"):
+            self.set_environment_variable("OMP_PROC_BIND", "close")
+            self.set_environment_variable("OMP_PLACES", "cores")
+            self.add_experiment_variable("n_threads_per_proc", 8, True)
+            self.add_experiment_variable("n_ranks", 8, True)
+            target = "cpu"
+            bonded_target = "cpu"
+            npme = "0"
+
+        # Overrides +openmp settings
+        if self.spec.satisfies("+cuda"):
+            self.add_experiment_variable("n_gpus", 8, True)
+            target = "gpu"
+            bonded_target = "cpu"
+            npme = "1"
+        elif self.spec.satisfies("+rocm"):
+            self.add_experiment_variable("n_gpus", 8, True)
+            target = "gpu"
+            bonded_target = "cpu"
+            npme = "1"
+
+        input_variables = {
+            "target": f"{target}",
+            "size": "1536",
+            "dlb": "no",
+            "pin": "off",
+            "maxh": "0.05",
+            "nsteps": "1000",
+            "nstlist": "200",
+            "npme": f"{npme}",
+        }
+
+        other_input_variables = {
+            "nb": f"{target}",
+            "pme": "auto",
+            "bonded": f"{bonded_target}",
+            "update": f"{target}",
+        }
+
+        for k, v in input_variables.items():
+            self.add_experiment_variable(k, v, True)
+        for k, v in other_input_variables.items():
+            self.add_experiment_variable(k, v)
+
+    def compute_spack_section(self):
+        # get package version
+        app_version = self.spec.variants["version"][0]
+
+        # get system config options
+        # TODO: Get compiler/mpi/package handles directly from system.py
+        system_specs = {}
+        system_specs["compiler"] = "default-compiler"
+        system_specs["mpi"] = "default-mpi"
+        system_specs["blas"] = "default-blas"
+        system_specs["lapack"] = "default-lapack"
+
+        # set package spack specs
+        # empty package_specs value implies external package
+        self.add_spack_spec(system_specs["mpi"])
+        # 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["lapack"])
+
+        spack_specs = "+mpi~hwloc"
+        spack_specs += "+sycl" if self.spec.satisfies("+rocm") else "~sycl"
+
+        if self.spec.satisfies("+cuda") or self.spec.satisfies("+rocm"):
+            spack_specs += f" gpu-aware-mpi={self.spec.variants['gpu-aware-mpi'][0]} "
+            spack_specs += " ~double "
+        else:
+            spack_specs += " gpu-aware-mpi=off "
+
+        self.add_spack_spec(
+            self.name,
+            [f"gromacs@{app_version} {spack_specs}", system_specs["compiler"]],
+        )
diff --git a/lib/benchpark/experiment.py b/lib/benchpark/experiment.py
@@ -144,6 +144,9 @@ def add_experiment_variable(self, name, values, use_in_expr_name=False):
         if use_in_expr_name:
             self.expr_name.append(f"{{{name}}}")
 
+    def set_environment_variable(self, name, values):
+        self.set_env_vars[name] = values
+
     def zip_experiment_variables(self, name, variable_names):
         self.zips[name] = list(variable_names)
 
@@ -165,6 +168,7 @@ def compute_applications_section(self):
 
     def compute_applications_section_wrapper(self):
         self.expr_name = []
+        self.set_env_vars = {}
         self.variables = {}
         self.zips = {}
         self.matrix = []

diff --git a/repo/gromacs/package.py b/repo/gromacs/package.py
@@ -8,10 +8,9 @@
 import llnl.util.filesystem as fs
 
 from spack.package import *
-from spack.pkg.benchpark.rocm_consistency import RocmConsistency as RocmConsistency
 
 
-class Gromacs(CMakePackage, CudaPackage, ROCmPackage, RocmConsistency):
+class Gromacs(CMakePackage, CudaPackage, ROCmPackage):
     """GROMACS is a molecular dynamics package primarily designed for simulations
     of proteins, lipids and nucleic acids. It was originally developed in
     the Biophysical Chemistry department of University of Groningen, and is now
@@ -89,6 +88,16 @@ class Gromacs(CMakePackage, CudaPackage, ROCmPackage, RocmConsistency):
     variant(
         "mpi", default=True, description="Activate MPI support (disable for Thread-MPI support)"
     )
+    # off: turn off GPU-aware MPI
+    # on: turn on, but allow groamcs to disable it if GPU-aware MPI is not supported
+    # force: turn on and force gromacs to use GPU-aware MPI. May result in error if unsupported
+    variant(
+        "gpu-aware-mpi",
+        default="on",
+        values=("on", "off", "force"),
+        when="@2021: +mpi",
+        description="Use GPU-aware MPI",
+    )
     variant("shared", default=True, description="Enables the build of shared libraries")
     variant(
         "double",
@@ -279,6 +288,9 @@ class Gromacs(CMakePackage, CudaPackage, ROCmPackage, RocmConsistency):
     depends_on("[email protected]:3", type="build", when="%fj")
     depends_on("cuda", when="+cuda")
 
+    for target in ("none", "gfx803", "gfx900", "gfx906", "gfx908", "gfx90a", "gfx942"):
+        requires(f"^[email protected]+rocm amdgpu_target={target}", when=f"gromacs@2024+rocm amdgpu_target={target}")
+
     with when("+rocm"):
         depends_on("sycl")
         depends_on("hip")
@@ -407,6 +419,17 @@ def setup_run_environment(self, env):
         if self.compiler.extra_rpaths:
             for rpath in self.compiler.extra_rpaths:
                 env.prepend_path("LD_LIBRARY_PATH", rpath)
+        if "+mpi" in self.spec:
+            if self.spec["mpi"].extra_attributes and "ldflags" in self.spec["mpi"].extra_attributes:
+                env.append_flags("LDFLAGS", self.spec["mpi"].extra_attributes["ldflags"])
+
+    def setup_build_environment(self, env):
+        if self.compiler.extra_rpaths:
+            for rpath in self.compiler.extra_rpaths:
+                env.prepend_path("LD_LIBRARY_PATH", rpath)
+        if "+mpi" in self.spec:
+            if self.spec["mpi"].extra_attributes and "ldflags" in self.spec["mpi"].extra_attributes:
+                env.append_flags("LDFLAGS", self.spec["mpi"].extra_attributes["ldflags"])
 
 class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
     @run_after("build")
@@ -497,6 +520,27 @@ def cmake_args(self):
                         "-DMPI_CXX_COMPILER=%s" % self.spec["mpi"].mpicxx,
                     ]
                 )
+            if 'on' in self.spec.variants['gpu-aware-mpi'].value:
+                options.extend(
+                    [
+                        "-DGMX_ENABLE_DIRECT_GPU_COMM=ON",
+                        "-DGMX_FORCE_GPU_AWARE_MPI=OFF",
+                    ]
+                )
+            elif 'force' in self.spec.variants['gpu-aware-mpi'].value:
+                options.extend(
+                    [
+                        "-DGMX_ENABLE_DIRECT_GPU_COMM=ON",
+                        "-DGMX_FORCE_GPU_AWARE_MPI=ON",
+                    ]
+                )
+            if "+rocm" in self.spec:
+                if self.spec["mpi"].extra_attributes and "ldflags" in self.spec["mpi"].extra_attributes:
+                    options.extend(
+                        [
+                            "-DCMAKE_EXE_LINKER_FLAGS=%s" % self.spec["mpi"].extra_attributes["ldflags"],
+                        ]
+                    )
         else:
             options.extend(
                 [
@@ -570,14 +614,14 @@ def cmake_args(self):
             options.append("-DCUDA_TOOLKIT_ROOT_DIR:STRING=" + self.spec["cuda"].prefix)
 
         target = self.spec.target
-        if "+cuda" in self.spec and target.family == "ppc64le":
+        if target.family == "ppc64le":
             options.append("-DGMX_EXTERNAL_LAPACK:BOOL=OFF")
         else:
             options.append("-DGMX_EXTERNAL_LAPACK:BOOL=ON")
             if self.spec["lapack"].libs:
                 options.append("-DGMX_LAPACK_USER={0}".format(self.spec["lapack"].libs.joined(";")))
 
-        if "+cuda" in self.spec and target.family == "ppc64le":
+        if target.family == "ppc64le":
             options.append("-DGMX_EXTERNAL_BLAS:BOOL=OFF")
         else:
             options.append("-DGMX_EXTERNAL_BLAS:BOOL=ON")

diff --git a/repo/rocm-consistency/package.py b/repo/rocm-consistency/package.py
diff --git a/systems/llnl-cluster/system.py b/systems/llnl-cluster/system.py
@@ -37,6 +37,18 @@ class LlnlCluster(System):
         description="Which compiler to use",
     )
 
+    variant(
+        "lapack",
+        default="intel-oneapi-mkl",
+        description="Which lapack to use",
+    )
+
+    variant(
+        "blas",
+        default="intel-oneapi-mkl",
+        description="Which blas to use",
+    )
+
     def initialize(self):
         super().initialize()
 
@@ -56,26 +68,22 @@ def generate_description(self, output_dir):
     def external_pkg_configs(self):
         externals = LlnlCluster.resource_location / "externals"
 
-        compiler = self.spec.variants["compiler"][0]
-
         selections = [externals / "base" / "00-packages.yaml"]
 
-        if compiler == "gcc":
+        if self.spec.satisfies("compiler=gcc"):
             selections.append(externals / "mpi" / "00-gcc-packages.yaml")
-        elif compiler == "intel":
+        elif self.spec.satisfies("compiler=intel"):
             selections.append(externals / "mpi" / "01-intel-packages.yaml")
 
         return selections
 
     def compiler_configs(self):
         compilers = LlnlCluster.resource_location / "compilers"
 
-        compiler = self.spec.variants["compiler"][0]
-
         selections = []
-        if compiler == "gcc":
+        if self.spec.satisfies("compiler=gcc"):
             selections.append(compilers / "gcc" / "00-gcc-12-compilers.yaml")
-        elif compiler == "intel":
+        elif self.spec.satisfies("compiler=intel"):
             selections.append(compilers / "intel" / "00-intel-2021-6-0-compilers.yaml")
 
         return selections
@@ -85,13 +93,17 @@ def sw_description(self):
         will fail if these variables are not defined though, so for now
         they are still generated (but with more-generic values).
         """
-        return """\
+        return f"""\
 software:
   packages:
     default-compiler:
       pkg_spec: gcc
     default-mpi:
       pkg_spec: mvapich2
+    default-lapack:
+      pkg_spec: {self.spec.variants["lapack"][0]}
+    default-blas:
+      pkg_spec: {self.spec.variants["blas"][0]}
     compiler-gcc:
       pkg_spec: gcc
     compiler-intel: