# Copyright 2013-2023 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
from spack.package import *
class Ams(CMakePackage, CudaPackage):
"""AMS Autonomous Multiscale Framework."""
homepage = "https://github.com/LLNL/AMS"
git = "git@github.com:LLNL/AMS.git"
maintainers("koparasy", "lpottier")
version("develop", branch="develop", submodules=False)
version("07.25.23-alpha", tag="07.25.23-alpha", submodules=False)
variant(
"faiss",
default=False,
description="Build with FAISS index as uncertainty quantification module",
)
variant(
"caliper", default=False, description="Build with caliper for gather performance counters"
)
variant("torch", default=False, description="Use torch for surrogate models")
variant("mpi", default=False, description="Enable MPI support")
variant("examples", default=False, description="Enable examples")
variant("redis", default=False, description="Enable redis database")
variant("hdf5", default=False, description="Enable HDF5 data storage")
variant("rabbitmq", default=False, description="Enable RabbitMQ as data broker")
variant(
"verbose",
default=False,
description="Enable AMSLib verbose output (controlled by environment variable)",
)
depends_on("umpire")
depends_on("mpi", when="+mpi")
depends_on("caliper+cuda", when="+caliper +cuda")
depends_on("faiss+cuda", when="+faiss +cuda")
depends_on("mfem+cuda", when="+examples +cuda")
depends_on("py-torch+cuda", when="+torch +cuda")
depends_on("py-torch~cuda", when="+torch ~cuda")
depends_on("caliper ~cuda", when="+caliper ~cuda")
depends_on("faiss ~cuda", when="+faiss ~cuda")
depends_on("mfem ~cuda", when="+examples ~cuda")
depends_on("redis-plus-plus", when="+redis")
depends_on("hdf5", when="+hdf5")
depends_on("amqp-cpp +tcp", when="+rabbitmq")
with when("+cuda"):
cuda_archs = CudaPackage.cuda_arch_values
with when("+examples"):
depends_on("mfem+cuda")
for sm_ in cuda_archs:
depends_on(
"mfem +cuda cuda_arch={0}".format(sm_), when="cuda_arch={0}".format(sm_)
)
with when("+torch"):
depends_on("py-torch+cuda")
for sm_ in cuda_archs:
depends_on(
"py-torch +cuda cuda_arch={0}".format(sm_), when="cuda_arch={0}".format(sm_)
)
with when("+caliper"):
depends_on("caliper+cuda", when="+caliper")
for sm_ in cuda_archs:
depends_on(
"caliper +cuda cuda_arch={0}".format(sm_), when="cuda_arch={0}".format(sm_)
)
depends_on("umpire+cuda")
for sm_ in cuda_archs:
depends_on("umpire +cuda cuda_arch={0}".format(sm_), when="cuda_arch={0}".format(sm_))
with when("+faiss"):
depends_on("faiss+cuda", when="+faiss")
for sm_ in cuda_archs:
depends_on(
"umpire +cuda cuda_arch={0}".format(sm_), when="cuda_arch={0}".format(sm_)
)
def cmake_args(self):
spec = self.spec
args = []
args.append("-DUMPIRE_DIR={0}".format(spec["umpire"].prefix))
args.append("-DWITH_MPI={0}".format("On" if "+mpi" in spec else "Off"))
args.append(
"-DWITH_DB={0}".format(
"On" if ("+redis" in spec or "hdf5" in spec or "+rabbitmq" in spec) else "Off"
)
)
if "+verbose" in spec:
args.append("-DWITH_AMS_DEBUG=On")
if "+hdf5" in spec:
args.append("-DWITH_HDF5=On")
args.append("-DHDF5_Dir={0}".format(spec["hdf5"].prefix))
if "+cuda" in spec:
args.append("-DWITH_CUDA=On")
cuda_arch = spec.variants["cuda_arch"].value[0]
args.append("-DAMS_CUDA_ARCH={0}".format(cuda_arch))
if "+caliper" in spec:
args.append("-DWITH_CALIPER=On")
args.append("-DCALIPER_DIR={0}/share/cmake/caliper".format(spec["caliper"].prefix))
else:
args.append("-DWITH_CALIPER=Off")
if "+faiss" in spec:
args.append("-DWITH_FAISS=On")
args.append("-DFAISS_DIR={0}".format(spec["faiss"].prefix))
else:
args.append("-DWITH_FAISS=Off")
if "+torch" in spec:
args.append("-DWITH_TORCH=On")
args.append(
"-DTorch_DIR={0}/lib/python{1}/site-packages"
"/torch/share/cmake/Torch".format(
spec["py-torch"].prefix, spec["python"].version.up_to(2)
)
)
if "+redis" in spec:
args.append("-DWITH_REDIS=On")
args.append("-DREDIS_PLUS_PLUS_DIR={0}".format(spec["redis-plus-plus"].prefix))
if "+rabbitmq" in spec:
args.append("-DWITH_RMQ=On")
args.append("-Damqpcpp_DIR={0}/cmake".format(spec["amqp-cpp"].prefix))
if "+examples" in spec:
args.append("-DWITH_EXAMPLES=On")
args.append("-DMFEM_DIR={0}".format(spec["mfem"].prefix))
return args