From ac8711e28dc3763290408954a4a27f447a0d4c89 Mon Sep 17 00:00:00 2001
From: Dom Heinzeller <dom.heinzeller@icloud.com>
Date: Wed, 31 Jan 2024 20:20:30 -0700
Subject: [PATCH] Complete test_prebuiild/test_mpi_test

---
 test_prebuild/test_mpi_test/README.md         |  8 ++---
 test_prebuild/test_mpi_test/main.F90          | 33 ++++++++++++++++++-
 .../test_mpi_test/mpi_test_scheme.F90         | 15 +++++++--
 3 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/test_prebuild/test_mpi_test/README.md b/test_prebuild/test_mpi_test/README.md
index ab1f7f0f..a00efbbb 100644
--- a/test_prebuild/test_mpi_test/README.md
+++ b/test_prebuild/test_mpi_test/README.md
@@ -20,16 +20,16 @@ mkdir build_mpi_f90
 cd build_mpi_f90
 cmake -DMPI_F90=ON .. 2>&1 | tee log.cmake
 make 2>&1 | tee log.make
-./test_mpi_test.x
+mpiexec -np 3 ./test_mpi_test.x
 cd ..
 
-rm -fr build_mpi_f90
-mkdir build_mpi_f90
+rm -fr build_mpi_f08
+mkdir build_mpi_f08
 ../../scripts/ccpp_prebuild.py --config=ccpp_prebuild_config.py --builddir=build_mpi_f08
 cd build_mpi_f08
 cmake -DMPI_F08=ON .. 2>&1 | tee log.cmake
 make 2>&1 | tee log.make
-./test_mpi_test.x
+mpiexec -np 3 ./test_mpi_test.x
 cd ..
 
 ```
diff --git a/test_prebuild/test_mpi_test/main.F90 b/test_prebuild/test_mpi_test/main.F90
index eead3c32..975ad56b 100644
--- a/test_prebuild/test_mpi_test/main.F90
+++ b/test_prebuild/test_mpi_test/main.F90
@@ -2,6 +2,12 @@ program test_mpi
 
    use, intrinsic :: iso_fortran_env, only: error_unit
 
+#if defined MPI_F08
+   use mpi_f08
+#elif defined MPI_F90
+   use mpi
+#endif
+
    use ccpp_types, only: ccpp_t
    use data, only: ncols
    use data, only: ccpp_data, ccpp_mpi_comm
@@ -15,7 +21,24 @@ program test_mpi
    implicit none
 
    character(len=*), parameter :: ccpp_suite = 'mpi_test_suite'
-   integer :: ib, ierr
+   integer :: ib, ierr, mpi_size, mpi_rank
+
+   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+   ! MPI init step                                  !
+   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+#if defined MPI_F08 || defined MPI_F90
+   call mpi_init(ierr)
+   ccpp_mpi_comm = mpi_comm_world
+   call mpi_comm_size(mpi_comm_world, mpi_size, ierr)
+   call mpi_comm_rank(mpi_comm_world, mpi_rank, ierr)
+#else
+   ccpp_mpi_comm = 0
+   mpi_size = 1
+   mpi_rank = 0
+#endif
+
+   write(error_unit,'(a,i0,a,i0)') "test_mpi: process: ", mpi_rank+1, " of ", mpi_size
 
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! CCPP init step                                 !
@@ -80,6 +103,14 @@ program test_mpi
       stop 1
    end if
 
+   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+   ! MPI finalize step                              !
+   !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+#if defined MPI_F08 || defined MPI_F90
+   call mpi_finalize(ierr)
+#endif
+
 contains
 
 end program test_mpi
\ No newline at end of file
diff --git a/test_prebuild/test_mpi_test/mpi_test_scheme.F90 b/test_prebuild/test_mpi_test/mpi_test_scheme.F90
index 92c1db47..ad6a2d4d 100644
--- a/test_prebuild/test_mpi_test/mpi_test_scheme.F90
+++ b/test_prebuild/test_mpi_test/mpi_test_scheme.F90
@@ -29,18 +29,29 @@ module mpi_test_scheme
 !! \htmlinclude mpi_test_scheme_init.html
 !!
    subroutine mpi_test_scheme_init(ncols, ccpp_mpi_comm, errmsg, errflg)
+      ! Interface variables
       integer, intent(in)               :: ncols
       ccpp_mpi_comm_type, intent(in)    :: ccpp_mpi_comm
       character(len=*), intent(out)     :: errmsg
       integer,          intent(out)     :: errflg
+      ! Local variables
+      integer :: mpi_rank, mpi_test, ierr
 
       ! Initialize CCPP error handling variables
       errmsg = ''
       errflg = 0
+
+      ! The MPI_F08 and MPI_F90 blocks are identical except the debug output
 #if defined MPI_F08
-      write(error_unit,'(a,i3)') 'In mpi_test_scheme_init, MPI_F08 branch'
+      call mpi_comm_rank(ccpp_mpi_comm, mpi_rank, ierr)
+      mpi_test = mpi_rank
+      call mpi_bcast(mpi_test, 1, MPI_INTEGER, 0, ccpp_mpi_comm, ierr)
+      write(error_unit,'(a,i3)') 'In mpi_test_scheme_init, MPI_F08 branch: MPI root has rank', mpi_test
 #elif defined MPI_F90
-      write(error_unit,'(a,i3)') 'In mpi_test_scheme_init, MPI_F90 branch'
+      call mpi_comm_rank(ccpp_mpi_comm, mpi_rank, ierr)
+      mpi_test = mpi_rank
+      call mpi_bcast(mpi_test, 1, MPI_INTEGER, 0, ccpp_mpi_comm, ierr)
+      write(error_unit,'(a,i3)') 'In mpi_test_scheme_init, MPI_F90 branch: MPI root has rank', mpi_test
 #else
       write(error_unit,'(a,i3)') 'In mpi_test_scheme_init, NO MPI branch'
 #endif