GHEX with NCCL uenv

ci/ci.yml

@@ -7,7 +7,8 @@ pipeline-configure:
   tags: [languard-k8s-lightweight]
   image: docker.io/python:latest
   script:
-    - git clone https://github.com/eth-cscs/uenv-pipeline.git
+    - git clone https://github.com/msimberg/uenv-pipeline.git
+    - git -C ./uenv-pipeline checkout v5-libfabric-buildable
     - ./uenv-pipeline/configure-pipeline -c./config.yaml -r./recipes -s$system -u$uenv -a$uarch -o./pipeline.yml
   artifacts:
     paths:

recipes/icon-dsl/25.12/a100/environments.yaml

@@ -1,26 +1,34 @@
 icon:
   compiler:
-      - toolchain: gcc
-        spec: gcc
-      - toolchain: llvm
-        spec: nvhpc
+  - toolchain: gcc
+    spec: gcc
+  - toolchain: llvm
+    spec: nvhpc
   mpi:
-      spec: [email protected]%nvhpc
-      gpu: cuda
-  unify: true
+    spec: [email protected]%nvhpc
+    gpu: cuda
   specs:
   # icon4py deps
+  - [email protected] +cuda +gdrcopy fabrics=cxi,rxm,tcp
+  - [email protected]
+  - [email protected]
+  - [email protected]
   - [email protected]
-  - [email protected] +cuda cuda_arch=80
-  - ghex@async-mpi +python +cuda cuda_arch=80
+  - [email protected] +cuda cuda_arch=80
+  - [email protected]=0.5.0 backend=nccl # nccl-2
+  - [email protected]=0.5.0 # nccl-context
+  - [email protected]
+  - [email protected]
+  - nccl-tests
+  - [email protected]
   - boost +thread
   - py-mpi4py
-  - icon4py@0.0.15
-  - icon-exclaim @0.3.0 %nvhpc dsl=substitute +grib2 +eccodes-definitions +ecrad
+  - icon4py@0.1.0
+  - [email protected] %nvhpc dsl=substitute +grib2 +eccodes-definitions +ecrad
     ~emvorado +art +dace gpu=nvidia-80 +mpi-gpu +realloc-buf ~aes ~jsbach ~ocean ~coupling
     ~rte-rrtmgp ~loop-exchange ~async-io-rma ~cuda-graphs fflags="-traceback"
   # icon deps
-  - cuda@12.4
+  - cuda@12.6:12
   - [email protected]%nvhpc +tools +fortran +aec +openmp jp2k=jasper
   - [email protected]
   - hdf5%nvhpc +szip +hl +fortran +mpi
@@ -30,14 +38,15 @@ icon:
   - [email protected]%nvhpc
   - openblas
   # for validation
-  - cdo
+  # - cdo # TODO: Enable? Currently broken.
   # everything needed for nccl on SS11
   - nccl
   - nccl-tests
   # extras
   - libfyaml
   - zlib-ng
   - cmake
+  unify: true
   variants:
   - +mpi
   - +cuda
@@ -46,6 +55,67 @@ icon:
     default:
       link: roots
       uenv:
-        add_compilers: true
+        prefix_paths:
+          LD_LIBRARY_PATH: [lib, lib64]
+# TODO: What's the cleanest way to do this? We want both dsl=substitute and
+# dsl=verify available in the uenv. dsl=verify does not have to be available
+# easily.
+icon-verify:
+  compiler:
+  - toolchain: gcc
+    spec: gcc
+  - toolchain: llvm
+    spec: nvhpc
+  mpi:
+    spec: [email protected]%nvhpc
+    gpu: cuda
+  specs:
+  # icon4py deps
+  - [email protected] +cuda +gdrcopy fabrics=cxi,rxm,tcp
+  - [email protected]
+  - [email protected]
+  - [email protected]
+  - [email protected]
+  - [email protected] +cuda cuda_arch=80
+  - [email protected]=0.5.0 backend=nccl # nccl-2
+  - [email protected]=0.5.0 # nccl-context
+  - [email protected]
+  - [email protected]
+  - nccl-tests
+  - [email protected]
+  - boost +thread
+  - py-mpi4py
+  - [email protected]
+  - [email protected] %nvhpc dsl=verify +grib2 +eccodes-definitions +ecrad
+    ~emvorado +art +dace gpu=nvidia-80 +mpi-gpu +realloc-buf ~aes ~jsbach ~ocean ~coupling
+    ~rte-rrtmgp ~loop-exchange ~async-io-rma ~cuda-graphs fflags="-traceback"
+  # icon deps
+  - [email protected]:12
+  - [email protected]%nvhpc +tools +fortran +aec +openmp jp2k=jasper
+  - [email protected]
+  - hdf5%nvhpc +szip +hl +fortran +mpi
+  - [email protected]%nvhpc
+  - [email protected]%nvhpc
+  - [email protected]%nvhpc
+  - [email protected]%nvhpc
+  - openblas
+  # for validation
+  # - cdo # TODO: Enable? Currently broken.
+  # everything needed for nccl on SS11
+  - nccl
+  - nccl-tests
+  # extras
+  - libfyaml
+  - zlib-ng
+  - cmake
+  unify: true
+  variants:
+  - +mpi
+  - +cuda
+  - cuda_arch=80
+  views:
+    default-verify:
+      link: roots
+      uenv:
         prefix_paths:
           LD_LIBRARY_PATH: [lib, lib64]

recipes/icon-dsl/25.12/a100/pre-install

@@ -11,10 +11,18 @@ packages=(
     "oomph"
     "uv"
     "icon4py"
-    # "icon"
+    "libfabric"
+    "nccl"
+    "libcxi"
+    "cxi-driver"
+    "cassini-headers"
+    "aws-ofi-nccl"
+    "gdrcopy"
+    "gettext"
+    "libfuse"
     # "icon-exclaim"
 )
-SPACK_C2SM_TAG="v0.22.2.6"
+SPACK_C2SM_TAG="sc-nccl"
 
 echo "----------------------------------------"
 echo "SPACK C2SM Setup Script"

recipes/icon-dsl/25.12/a100/repo/packages/icon-exclaim/package.py

@@ -365,9 +365,9 @@ class IconExclaim(Icon):
 
     version('develop', branch='icon-dsl', submodules=True)
     version("0.3.0",
-            commit="5c5b742a969af2bd491e26cd0a05a35838f121c4",
+            commit="a0be2c3e0448ec2dc92024e3b38ec635435ac0dd",
             submodules=True)
-  
+
     # EXCLAIM-GT4Py specific features:
     dsl_values = ('substitute', 'verify')
     variant('dsl',
@@ -377,10 +377,19 @@ class IconExclaim(Icon):
             description='Build with GT4Py dynamical core',
             multi=True)
 
-    depends_on("[email protected]", when="@0.3.0")       
     for x in dsl_values:
         depends_on('icon4py', type="build", when=f"dsl={x}")
 
+
+    # TODO: Should this be set here or in the icon4py package?
+    def setup_build_environment(self, env):
+        if self.spec.variants['dsl'].value != ('none', ):
+            # TODO: clean up
+            print(f"adding {self.spec['icon4py'].prefix.share.venv.bin} to PATH for icon4py bindings because +dsl is enabled")
+            env.prepend_path("PATH", self.spec["icon4py"].prefix.share.venv.bin)
+            env.append_path("PATH", self.spec["icon4py"].prefix.share.venv.bin)
+
+
     def configure_args(self):
         raw_args = super().configure_args()
 
@@ -404,21 +413,14 @@ def configure_args(self):
         dsl = self.spec.variants['dsl'].value
         if dsl != ('none', ):
             if 'substitute' in dsl:
-                args_flags.append('--enable-py2f=substitute')
+                args_flags.append('--enable-icon4py=substitute')
             elif 'verify' in dsl:
-                args_flags.append('--enable-py2f=verify')
+                args_flags.append('--enable-icon4py=verify')
             else:
                 raise ValueError(
                     f"Unknown DSL variant '{dsl}'. "
                     f"Valid options are: {', '.join(('none',) + dsl_values)}")
 
-            # Add icon4py paths and libs
-            icon4py_prefix = self.spec["icon4py"].prefix
-            bindings_dir = os.path.join(icon4py_prefix, "src")
-
-            ldflags.append(f"-L{bindings_dir} -Wl,-rpath,{bindings_dir}")
-            libs.append("-licon4py_bindings")
-
         # Remove duplicates
         icon_ldflags = list(dict.fromkeys(icon_ldflags))
         ldflags = list(dict.fromkeys(ldflags))
@@ -434,25 +436,3 @@ def configure_args(self):
             final_args.append("LIBS=" + " ".join(libs))
 
         return final_args
-
-    def build(self, spec, prefix):
-        # Check the variant
-        dsl = self.spec.variants['dsl'].value
-        if dsl != ('none', ):
-            file = "icon4py_bindings.f90"
-
-            bindings_dir = os.path.join(self.spec["icon4py"].prefix, "src")
-            src_file = os.path.join(bindings_dir, file)
-
-            build_py2f_dir = os.path.join(self.stage.source_path, "src",
-                                          "build_py2f")
-            os.makedirs(build_py2f_dir, exist_ok=True)
-            dest_file = os.path.join(build_py2f_dir, file)
-
-            shutil.copy2(src_file, dest_file)
-            print(
-                f"Copied {src_file} to build directory {dest_file} because +dsl is enabled"
-            )
-
-        # Proceed with the normal build
-        super().build(spec, prefix)

recipes/icon-dsl/25.12/a100/repo/packages/py-cupy/package.py

@@ -18,43 +18,19 @@ class PyCupy(PythonPackage, CudaPackage, ROCmPackage):
 
     homepage = "https://cupy.dev/"
     pypi = "cupy/cupy-8.0.0.tar.gz"
+    version("13.5.1", sha256="3dba2f30258463482d52deb420862fbbbaf2c446165a5e8d67377ac6cb5c0870")
+    version("13.4.0", sha256="d4b60e5a1d3b89be40fad0845bb9fc467a653abe8660f752416fd38d24ab7fdb")
+    version("13.3.0", sha256="9a2a17af2b99cce91dd1366939c3805e3f51f9de5046df64f29ccbad3bdf78ed")
+    version("13.2.0", sha256="e4dbd2b2ed4159a5cc0c0f98a710a014950eb2c16eeb455e956128f3b3bd0d51")
+    version("13.1.0", sha256="5caf62288481a27713384523623045380ff42e618be4245f478238ed1786f32d")
+    version("12.1.0", sha256="f6d31989cdb2d96581da12822e28b102f29e254427195c2017eac327869b7320")
+    version("12.0.0", sha256="61ddbbef73d50d606bd5087570645f3c91ec9176c2566784c1d486d6a3404545")
+    version("11.6.0", sha256="53dbb840072bb32d4bfbaa6bfa072365a30c98b1fcd1f43e48969071ad98f1a7")
+    version("11.5.0", sha256="4bc8565bded22cc89b210fd9fb48a5d5316f30701e12bb23852a60314e1f9f6e")
+    version("11.4.0", sha256="03d52b2626e02a3a2b46d714c1cd03e702c8fe33915fcca6ed8de5c539964f49")
+    version("11.3.0", sha256="d057cc2f73ecca06fae8b9c270d9e14116203abfd211a704810cc50a453b4c9e")
+    version("11.2.0", sha256="c33361f117a347a63f6996ea97446d17f1c038f1a1f533e502464235076923e2")
 
-    # version(
-    #     "13.5.1",
-    #     sha256="3dba2f30258463482d52deb420862fbbbaf2c446165a5e8d67377ac6cb5c0870",
-    # )
-    version(
-        "13.1.0",
-        sha256="5caf62288481a27713384523623045380ff42e618be4245f478238ed1786f32d",
-    )
-    version(
-        "12.1.0",
-        sha256="f6d31989cdb2d96581da12822e28b102f29e254427195c2017eac327869b7320",
-    )
-    version(
-        "12.0.0",
-        sha256="61ddbbef73d50d606bd5087570645f3c91ec9176c2566784c1d486d6a3404545",
-    )
-    version(
-        "11.6.0",
-        sha256="53dbb840072bb32d4bfbaa6bfa072365a30c98b1fcd1f43e48969071ad98f1a7",
-    )
-    version(
-        "11.5.0",
-        sha256="4bc8565bded22cc89b210fd9fb48a5d5316f30701e12bb23852a60314e1f9f6e",
-    )
-    version(
-        "11.4.0",
-        sha256="03d52b2626e02a3a2b46d714c1cd03e702c8fe33915fcca6ed8de5c539964f49",
-    )
-    version(
-        "11.3.0",
-        sha256="d057cc2f73ecca06fae8b9c270d9e14116203abfd211a704810cc50a453b4c9e",
-    )
-    version(
-        "11.2.0",
-        sha256="c33361f117a347a63f6996ea97446d17f1c038f1a1f533e502464235076923e2",
-    )
 
     variant(
         "all", default=False, description="Enable optional py-scipy, optuna, and cython"
@@ -63,32 +39,48 @@ class PyCupy(PythonPackage, CudaPackage, ROCmPackage):
     depends_on("cxx", type="build")  # generated
 
     depends_on("[email protected]:", when="@:11", type=("build", "run"))
-    depends_on("[email protected]:", when="@12:", type=("build", "run"))
+    depends_on("[email protected]:3.11", when="@12", type=("build", "run"))
+    depends_on("[email protected]:3.13", when="@13", type=("build", "run"))
+    depends_on("[email protected]:", when="@14", type=("build", "run"))
     depends_on("py-setuptools", type="build")
-    depends_on("[email protected]:2", type="build")
+    depends_on("py-setuptools@:73", when="@:13.3", type="build")
+    depends_on("[email protected]:0.29", type="build", when="@:13.3")
+    depends_on(
+        "py-cython@3:3.0.10,3.0.12:", type="build", when="@13.4:"
+    )  # 3.0.11 broken likely because of cython#6335, fixed in 3.0.12
+    depends_on("[email protected]:0.29", when="@:13.3 +all", type=("build", "run"))
+    depends_on("py-cython@3:", when="@13.4: +all", type=("build", "run"))
     depends_on("[email protected]:", type=("build", "run"))
     depends_on("[email protected]:1.25", when="@:11", type=("build", "run"))
-    depends_on("[email protected]:1.26", when="@12:", type=("build", "run"))
-    depends_on("[email protected]:1.28", when="@13:", type=("build", "run"))
-
-    depends_on("[email protected]:1.12", when="@:12+all", type=("build", "run"))
-    depends_on("[email protected]:1.13", when="@13:+all", type=("build", "run"))
-    depends_on("[email protected]:2", when="+all", type=("build", "run"))
+    depends_on("[email protected]:1.26", when="@12", type=("build", "run"))
+    depends_on("[email protected]:1", when="@13.1", type=("build", "run"))
+    depends_on("[email protected]:2.0", when="@13.2", type=("build", "run"))
+    depends_on("[email protected]:2.2", when="@13.4", type=("build", "run"))
+    depends_on("[email protected]:2.3", when="@13.5", type=("build", "run"))
+    depends_on("[email protected]:2", when="@14", type=("build", "run"))
+    depends_on("[email protected]:1.11", when="@:12+all", type=("build", "run"))
+    depends_on("[email protected]:1.16", when="@13+all", type=("build", "run"))
+    depends_on("[email protected]:1.16", when="@14+all", type=("build", "run"))
     depends_on("py-optuna@2:", when="+all", type=("build", "run"))
+    depends_on("py-optuna@3:", when="@12:+all", type=("build", "run"))
 
     # Based on https://github.com/cupy/cupy/releases
     depends_on("cuda@:11.9", when="@:11 +cuda")
     depends_on("cuda@:12.1", when="@12:12.1.0 +cuda")
-    depends_on("cuda@:12.4", when="@13:13.4.1 +cuda")
-    depends_on("cuda@:12.9", when="@13.5.1: +cuda")
+    depends_on("cuda@:12.1", when="@13.0 +cuda")
+    depends_on("cuda@:12.4", when="@13.1:13.2 +cuda")
+    depends_on("cuda@:12.6", when="@13.3 +cuda")
+    depends_on("cuda@:12.8", when="@13.4 +cuda")
+    depends_on("cuda@:12.9", when="@13.5 +cuda")
 
     for a in CudaPackage.cuda_arch_values:
+        depends_on("nccl +cuda cuda_arch={0}".format(a), when="+cuda cuda_arch={0}".format(a))
         depends_on(
-            "nccl +cuda cuda_arch={0}".format(a), when="+cuda cuda_arch={0}".format(a)
+            "nccl@2.16:2.26 +cuda cuda_arch={0}".format(a), when="@13+cuda cuda_arch={0}".format(a)
         )
 
     depends_on("cutensor", when="@:12.1.0 +cuda")
-    depends_on("[email protected].1.2", when="@13.1: +cuda")
+    depends_on("[email protected]", when="@13.1: +cuda")
 
     for _arch in ROCmPackage.amdgpu_targets:
         arch_str = "amdgpu_target={0}".format(_arch)