Skip to content

Commit

Permalink
Feature/gpu aware mpi libs (#512)
Browse files Browse the repository at this point in the history
* Add cray-mpich libs method for gtl

* cray-mpich libs method

* Add libs method for spectrum-mpi spack package

* license

* Remove redundant MPI ldflags

* cmake mpi link flag

* Fix saxy gtl

---------

Co-authored-by: Riyaz Haque <[email protected]>
Co-authored-by: pearce8 <[email protected]>
  • Loading branch information
3 people authored Jan 13, 2025
1 parent ef81107 commit 2e25190
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 21 deletions.
11 changes: 11 additions & 0 deletions repo/cray-mpich/package.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,17 @@ class CrayMpich(BuiltinCM):

variant("gtl", default=False, description="enable GPU-aware mode")

@property
def libs(self):
libs = super(CrayMpich, self).libs

if self.spec.satisfies("+gtl"):
gtl_lib_prefix = self.spec.extra_attributes["gtl_lib_path"]
libraries = ["libmpi_gtl_hsa"]
libs += find_libraries(libraries, root=gtl_lib_prefix, recursive=True)

return libs

def setup_run_environment(self, env):

super(CrayMpich, self).setup_run_environment(env)
Expand Down
14 changes: 1 addition & 13 deletions repo/gromacs/package.py
Original file line number Diff line number Diff line change
Expand Up @@ -419,17 +419,11 @@ 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")
Expand Down Expand Up @@ -534,13 +528,7 @@ def cmake_args(self):
"-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"],
]
)
options.extend([f"-DMPI_CXX_LINK_FLAGS='{self.spec['mpi'].libs.ld_flags}'"])
else:
options.extend(
[
Expand Down
9 changes: 1 addition & 8 deletions repo/saxpy/package.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,14 +42,6 @@ def stage_source(self):
def setup_build_environment(self, env):
if "+cuda" in self.spec:
env.set("CUDAHOSTCXX", self.spec["mpi"].mpicxx)
if self.spec["mpi"].extra_attributes and "gtl_lib_path" in self.spec["mpi"].extra_attributes:
if self.spec.satisfies("+rocm"): # no gtl for cpu-only runs
env.prepend_path("LD_LIBRARY_PATH", self.spec['mpi'].extra_attributes["gtl_lib_path"])

def setup_run_environment(self, env):
if self.spec["mpi"].extra_attributes and "gtl_lib_path" in self.spec["mpi"].extra_attributes:
if self.spec.satisfies("+rocm"): # no gtl for cpu-only runs
env.prepend_path("LD_LIBRARY_PATH", self.spec['mpi'].extra_attributes["gtl_lib_path"])

def cmake_args(self):
spec = self.spec
Expand Down Expand Up @@ -80,5 +72,6 @@ def cmake_args(self):
args.append("-DROCM_ARCH={0}".format(rocm_arch))

args.append(self.define_from_variant("USE_CALIPER", "caliper"))
args.append("-DMPI_CXX_LINK_FLAGS={0}".format(spec["mpi"].libs.ld_flags))

return args
25 changes: 25 additions & 0 deletions repo/spectrum-mpi/package.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
# 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.spectrum_mpi import SpectrumMpi as BuiltinSM


class SpectrumMpi(BuiltinSM):
@property
def libs(self):
libnames = [
'mpi_ibm',
'mpi_ibm_mpifh',
'mpiprofilesupport',
'mpi_ibm_usempi',
]
libs = list('lib' + x for x in libnames)
return find_libraries(
libs,
self.spec.prefix,
shared=True,
recursive=True
)

0 comments on commit 2e25190

Please sign in to comment.