Merge branch 'fix/nvidia-device-mpi' into dev

AFD-Illinois · Jan 29, 2024 · ae325e0 · ae325e0
2 parents 8c67810 + f19d63a
commit ae325e0
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 13 deletions.
diff --git a/external/kokkos-kernels/KokkosBatched_Util.hpp b/external/kokkos-kernels/KokkosBatched_Util.hpp
@@ -772,8 +772,8 @@ KOKKOS_INLINE_FUNCTION auto subview_wrapper(ViewType v, IdxType1 i1,
 }
 template <class ViewType, class IdxType1>
 KOKKOS_INLINE_FUNCTION auto subview_wrapper(ViewType v, IdxType1 i1,
-                                            Kokkos::ALL_t i2,
-                                            Kokkos::ALL_t i3,
+                                            Kokkos::Impl::ALL_t i2,
+                                            Kokkos::Impl::ALL_t i3,
                                             const BatchLayout::Left &layout_tag,
                                             const Trans::Transpose) {
   auto sv_nt = subview_wrapper(v, i1, i3, i2, layout_tag);
@@ -805,7 +805,7 @@ KOKKOS_INLINE_FUNCTION auto subview_wrapper(
 }
 template <class ViewType, class IdxType1>
 KOKKOS_INLINE_FUNCTION auto subview_wrapper(
-    ViewType v, IdxType1 i1, Kokkos::ALL_t i2, Kokkos::ALL_t i3,
+    ViewType v, IdxType1 i1, Kokkos::Impl::ALL_t i2, Kokkos::Impl::ALL_t i3,
     const BatchLayout::Right &layout_tag, const Trans::Transpose &) {
   auto sv_nt = subview_wrapper(v, i1, i3, i2, layout_tag);
 

diff --git a/external/parthenon b/external/parthenon
diff --git a/machines/chicoma.sh b/machines/chicoma.sh
@@ -11,7 +11,7 @@ if [[ "$HOST" == "ch-fe"* || "$HOST" == "nid00"* ]]; then
 
   # Cray environments get confused easy
   # Make things as simple as possible
-  # TODO ONLY NVHPC WORKS
+  # Only default "nvhpc" module is supported, see HPC docs
   module purge
   export CRAY_CPU_TARGET="x86-64"
   if [[ "$ARGS" == *"cuda"* ]]; then
@@ -30,7 +30,7 @@ if [[ "$HOST" == "ch-fe"* || "$HOST" == "nid00"* ]]; then
       C_NATIVE=nvc
       CXX_NATIVE=nvc++
     else
-      module load PrgEnv-nvhpc
+      module load PrgEnv-nvhpc/8.3.3
     fi
     module load craype-accel-nvidia80
     # GPU runtime opts
@@ -40,18 +40,32 @@ if [[ "$HOST" == "ch-fe"* || "$HOST" == "nid00"* ]]; then
     unset OMP_PLACES
 
     # Sometimes device-side buffers don't work
-    # if [conditions]
-    EXTRA_FLAGS="-DPARTHENON_ENABLE_HOST_COMM_BUFFERS=ON $EXTRA_FLAGS"
-    #else
-    #export MPICH_GPU_SUPPORT_ENABLED=1
+    if [[ "$ARGS" == *"hostside"* ]]; then
+      EXTRA_FLAGS="-DPARTHENON_ENABLE_HOST_COMM_BUFFERS=ON $EXTRA_FLAGS"
+    else
+      export MPICH_GPU_SUPPORT_ENABLED=1
+      export MPICH_GPU_MANAGED_MEMORY_SUPPORT_ENABLED=1
+      # Use the Cray wrappers else runtime errors (?)
+      C_NATIVE=cc
+      CXX_NATIVE=CC
+    fi
 
   else
+    # CPU CASE
     module load PrgEnv-aocc
     MPI_EXTRA_ARGS="--cpus-per-task=2"
   fi
-  module load cmake #cray-hdf5-parallel
+
+  # If using system HDF5, disable compression since it's old
+  if [[ "$ARGS" != *"hdf5"* ]]; then
+    module load cray-hdf5-parallel
+    EXTRA_FLAGS="-DPARTHENON_DISABLE_HDF5_COMPRESSION=ON $EXTRA_FLAGS"
+  fi
+
+  module load cmake
 
   # Runtime opts
   MPI_EXE="srun"
-  MPI_NUM_PROCS=""
+  MPI_NUM_PROCS=${MPI_NUM_PROCS:-4}
 fi
+
diff --git a/make.sh b/make.sh
@@ -309,8 +309,13 @@ if [[ "$ARGS" == *"clean"* ]]; then
       git apply ../patches/variant-hip.patch
     fi
     cd -
-  fi
 
+    # HIP also prefers new Kokkos.
+    # TODO work something out if on HIP machines w/o internet
+    cd external/parthenon
+    git submodule update --remote external/Kokkos
+    cd -
+  fi
 
   rm -rf build
 fi