Upcxx 2022.3.0 update (#30393)

* upcxx,gasnet: Add 2022.3.0 version hashes
* gasnet: Add ofi and ucx experimental conduits
* gasnet: Add CUDA and ROCm/HIP variants
* upcxx: Add Cray EX/Shasta detection and auto-config
* upcxx: Add ROCm/HIP support
* Use implicit dependencies from CudaPackage, ROCmPackage

2 files changed, 80 insertions, 21 deletions

diff --git a/var/spack/repos/builtin/packages/gasnet/package.py b/var/spack/repos/builtin/packages/gasnet/package.py
index b6a043e4c3..70679d55d2 100644
--- a/var/spack/repos/builtin/packages/gasnet/package.py
+++ b/var/spack/repos/builtin/packages/gasnet/package.py
@@ -8,7 +8,7 @@ import os
 from spack import *
 
 
-class Gasnet(Package):
+class Gasnet(Package, CudaPackage, ROCmPackage):
     """GASNet is a language-independent, networking middleware layer that
        provides network-independent, high-performance communication primitives
        including Remote Memory Access (RMA) and Active Messages (AM). It has been
@@ -36,6 +36,7 @@ class Gasnet(Package):
     version('main',    branch='stable')
     version('master',  branch='master')
 
+    version('2022.3.0',  sha256='91b59aa84c0680c807e00d3d1d8fa7c33c1aed50b86d1616f93e499620a9ba09')
     version('2021.9.0',  sha256='1b6ff6cdad5ecf76b92032ef9507e8a0876c9fc3ee0ab008de847c1fad0359ee')
     version('2021.3.0',  sha256='8a40fb3fa8bacc3922cd4d45217816fcb60100357ab97fb622a245567ea31747')
     version('2020.10.0', sha256='ed17baf7fce90499b539857ee37b3eea961aa475cffbde77e4c607a34ece06a0')
@@ -46,21 +47,32 @@ class Gasnet(Package):
 
     # The optional network backends:
     variant('conduits',
-            values=any_combination_of('smp', 'mpi', 'ibv', 'udp').with_default('smp'),
+            values=any_combination_of('smp', 'mpi', 'ibv',
+                                      'udp', 'ofi', 'ucx').with_default('smp'),
             description="The hardware-dependent network backends to enable.\n" +
-                        "(smp) = SMP conduit for single-node operation ;\n" +
-                        "(ibv) = Native InfiniBand verbs conduit ;\n" +
-                        "(udp) = Portable UDP conduit, for Ethernet networks ;\n" +
-                        "(mpi) = Low-performance/portable MPI conduit ;\n" +
-                        "For detailed recommendations, consult https://gasnet.lbl.gov")
+            "(smp) = SMP conduit for single-node operation ;\n" +
+            "(ibv) = Native InfiniBand verbs conduit ;\n" +
+            "(udp) = Portable UDP conduit, for Ethernet networks ;\n" +
+            "(mpi) = Low-performance/portable MPI conduit ;\n" +
+            "(ofi) = EXPERIMENTAL Portable OFI conduit over libfabric ;\n" +
+            "(ucx) = EXPERIMENTAL UCX conduit for Mellanox IB/RoCE ConnectX-5+ ;\n" +
+            "For detailed recommendations, consult https://gasnet.lbl.gov")
 
     variant('debug', default=False, description="Enable library debugging mode")
 
+    variant('cuda', default=False,
+            description='Enables support for the CUDA memory kind in some conduits')
+
+    variant('rocm', default=False,
+            description='Enables support for the ROCm/HIP memory kind in some conduits')
+
     depends_on('mpi', when='conduits=mpi')
 
     depends_on('autoconf@2.69', type='build', when='@master:')
     depends_on('automake@1.16:', type='build', when='@master:')
 
+    conflicts('hip@:4.4.0', when='+rocm')
+
     def install(self, spec, prefix):
         if spec.satisfies('@master:'):
             bootstrapsh = Executable("./Bootstrap")
@@ -84,6 +96,12 @@ class Gasnet(Package):
             if '+debug' in spec:
                 options.append("--enable-debug")
 
+            if '+cuda' in spec:
+                options.append("--enable-kind-cuda-uva")
+
+            if '+rocm' in spec:
+                options.append("--enable-kind-hip")
+
             if 'conduits=mpi' in spec:
                 options.append("--enable-mpi-compat")
             else:
@@ -93,6 +111,8 @@ class Gasnet(Package):
             for c in spec.variants['conduits'].value:
                 options.append("--enable-" + c)
 
+            options.append("--enable-rpath")
+
             configure(*options)
             make()
             make('install')
@@ -125,6 +145,8 @@ class Gasnet(Package):
             'smp': ['env', 'GASNET_PSHM_NODES=' + ranks],
             'mpi': [join_path(self.prefix.bin, 'gasnetrun_mpi'), '-n', ranks],
             'ibv': [join_path(self.prefix.bin, 'gasnetrun_ibv'), '-n', ranks],
+            'ofi': [join_path(self.prefix.bin, 'gasnetrun_ofi'), '-n', ranks],
+            'ucx': [join_path(self.prefix.bin, 'gasnetrun_ucx'), '-n', ranks],
             'udp': [join_path(self.prefix.bin, 'amudprun'), '-spawn', 'L', '-np', ranks]
         }
 
diff --git a/var/spack/repos/builtin/packages/upcxx/package.py b/var/spack/repos/builtin/packages/upcxx/package.py
index 48830ccb5a..32c8be8e5e 100644
--- a/var/spack/repos/builtin/packages/upcxx/package.py
+++ b/var/spack/repos/builtin/packages/upcxx/package.py
@@ -13,6 +13,11 @@ def is_CrayXC():
            (os.environ.get('CRAYPE_NETWORK_TARGET') == "aries")
 
 
+def is_CrayEX():
+    return (spack.platforms.host().name == 'cray') and \
+           (os.environ.get('CRAYPE_NETWORK_TARGET') in ['ofi', 'ucx'])
+
+
 def cross_detect():
     if is_CrayXC():
         if which('srun'):
@@ -22,10 +27,10 @@ def cross_detect():
     return 'none'
 
 
-class Upcxx(Package):
+class Upcxx(Package, CudaPackage, ROCmPackage):
     """UPC++ is a C++ library that supports Partitioned Global Address Space
     (PGAS) programming, and is designed to interoperate smoothly and
-    efficiently with MPI, OpenMP, CUDA and AMTs. It leverages GASNet-EX to
+    efficiently with MPI, OpenMP, CUDA, ROCm/HIP and AMTs. It leverages GASNet-EX to
     deliver low-overhead, fine-grained communication, including Remote Memory
     Access (RMA) and Remote Procedure Call (RPC)."""
 
@@ -39,6 +44,7 @@ class Upcxx(Package):
     version('develop', branch='develop')
     version('master',  branch='master')
 
+    version('2022.3.0', sha256='72bccfc9dfab5c2351ee964232b3754957ecfdbe6b4de640e1b1387d45019496')
     version('2021.9.0', sha256='9299e17602bcc8c05542cdc339897a9c2dba5b5c3838d6ef2df7a02250f42177')
     version('2021.3.0', sha256='3433714cd4162ffd8aad9a727c12dbf1c207b7d6664879fc41259a4b351595b7')
     version('2020.11.0', sha256='f6f212760a485a9f346ca11bb4751e7095bbe748b8e5b2389ff9238e9e321317',
@@ -53,7 +59,10 @@ class Upcxx(Package):
             description='Enables MPI-based spawners and mpi-conduit')
 
     variant('cuda', default=False,
-            description='Builds a CUDA-enabled version of UPC++')
+            description='Enables UPC++ support for the CUDA memory kind')
+
+    variant('rocm', default=False,
+            description='Enables UPC++ support for the ROCm/HIP memory kind')
 
     variant('cross', default=cross_detect(),
             description="UPC++ cross-compile target (autodetect by default)")
@@ -70,29 +79,34 @@ class Upcxx(Package):
     depends_on('gasnet conduits=none', when='+gasnet')
 
     depends_on('mpi', when='+mpi')
-    depends_on('cuda', when='+cuda')
     depends_on('python@2.7.5:', type=("build", "run"))
 
+    conflicts('hip@:4.4.0', when='+rocm')
+
     # All flags should be passed to the build-env in autoconf-like vars
     flag_handler = env_flags
 
-    def setup_run_environment(self, env):
+    def set_variables(self, env):
         env.set('UPCXX_INSTALL', self.prefix)
         env.set('UPCXX', self.prefix.bin.upcxx)
         if is_CrayXC():
             env.set('UPCXX_NETWORK', 'aries')
+        elif is_CrayEX():
+            env.set('UPCXX_NETWORK', 'ofi')
 
-    def setup_dependent_package(self, module, dep_spec):
-        dep_spec.upcxx = self.prefix.bin.upcxx
+    def setup_run_environment(self, env):
+        self.set_variables(env)
 
     def setup_dependent_build_environment(self, env, dependent_spec):
-        env.set('UPCXX_INSTALL', self.prefix)
-        env.set('UPCXX', self.prefix.bin.upcxx)
-        if is_CrayXC():
-            env.set('UPCXX_NETWORK', 'aries')
+        self.set_variables(env)
+
+    def setup_dependent_package(self, module, dep_spec):
+        dep_spec.upcxx = self.prefix.bin.upcxx
 
     def install(self, spec, prefix):
         env = os.environ
+        if (env.get('GASNET_CONFIGURE_ARGS') is None):
+            env['GASNET_CONFIGURE_ARGS'] = ''
         # UPC++ follows autoconf naming convention for LDLIBS, which is 'LIBS'
         if (env.get('LDLIBS')):
             env['LIBS'] = env['LDLIBS']
@@ -112,13 +126,12 @@ class Upcxx(Package):
                 env[var] = ":".join(
                     filter(lambda x: "libsci" not in x.lower(),
                            env[var].split(":")))
+        if is_CrayXC() or is_CrayEX():
             # Undo spack compiler wrappers:
             # the C/C++ compilers must work post-install
             real_cc = join_path(env['CRAYPE_DIR'], 'bin', 'cc')
             real_cxx = join_path(env['CRAYPE_DIR'], 'bin', 'CC')
             # workaround a bug in the UPC++ installer: (issue #346)
-            if (env.get('GASNET_CONFIGURE_ARGS') is None):
-                env['GASNET_CONFIGURE_ARGS'] = ''
             env['GASNET_CONFIGURE_ARGS'] += \
                 " --with-cc=" + real_cc + " --with-cxx=" + real_cxx
             if '+mpi' in spec:
@@ -132,6 +145,23 @@ class Upcxx(Package):
         options.append('--with-cc=' + real_cc)
         options.append('--with-cxx=' + real_cxx)
 
+        if is_CrayEX():
+            # Probe to find the right libfabric provider (SlingShot 10 vs 11)
+            fi_info = which('fi_info')('-l', output=str)
+            if fi_info.find('cxi') >= 0:
+                provider = 'cxi'
+            else:
+                provider = 'verbs;ofi_rxm'
+
+            # Append the recommended options for Cray Shasta
+            options.append('--with-pmi-version=cray')
+            options.append('--with-pmi-runcmd=\'srun -n %N -- %C\'')
+            options.append('--disable-ibv')
+            options.append('--enable-ofi')
+            options.append('--with-ofi-provider=' + provider)
+            env['GASNET_CONFIGURE_ARGS'] = \
+                '--with-ofi-spawner=pmi ' + env['GASNET_CONFIGURE_ARGS']
+
         if '+gasnet' in spec:
             options.append('--with-gasnet=' + spec['gasnet'].prefix.src)
 
@@ -144,9 +174,16 @@ class Upcxx(Package):
             options.append('--without-mpicc')
 
         if '+cuda' in spec:
-            options.append('--with-cuda')
+            options.append('--enable-cuda')
             options.append('--with-nvcc=' + spec['cuda'].prefix.bin.nvcc)
 
+        if '+rocm' in spec:
+            options.append('--enable-hip')
+            options.append('--with-ld-flags=' +
+                           self.compiler.cc_rpath_arg + spec['hip'].prefix.lib)
+
+        env['GASNET_CONFIGURE_ARGS'] = '--enable-rpath ' + env['GASNET_CONFIGURE_ARGS']
+
         configure(*options)
 
         make()