Avoid using PME decomposition in tests when it's not available

Previously, setting `GMX_GPU_PME_DECOMPOSITION=1` environment
variable would cause some tests to try to use GPU PME
decomposition even if it is not available.

Does not affect any production runs, only unit tests.

Fixes #4687

src/gromacs/fft/tests/fft_mpi.cpp

@@ -56,6 +56,7 @@
 #    include "gromacs/gpu_utils/devicebuffer.h"
 #endif
 #include "gromacs/utility/arrayref.h"
+#include "gromacs/utility/mpiinfo.h"
 #include "gromacs/utility/stringutil.h"
 
 #include "testutils/mpitest.h"
@@ -75,6 +76,24 @@ using GpuFftTestParams = std::tuple<std::tuple<IVec, // size of grid
 
 using GpuFftTestGridParams = std::tuple_element<0, GpuFftTestParams>::type;
 
+static GpuAwareMpiStatus getGpuAwareMpiStatusForFftBackend(const FftBackend fftBackend)
+{
+    /* This is not a 100%-reliable test. In SYCL, we check that the MPI is supported
+     * by the FFT library, not by the devices. E.g., we can have OpenCL Intel GPUs, which are
+     * not supported by GPU-aware Intel MPI (LevelZero backend is required). Or even NVIDIA or AMD
+     * GPUs in the same build. We do handle some common cases, however.
+     * This will be improved later in scope of #4638 */
+    switch (fftBackend)
+    {
+        case FftBackend::CuFFTMp:
+        case FftBackend::HeFFTe_CUDA:
+        case FftBackend::HeFFTe_Sycl_cuFFT: return checkMpiCudaAwareSupport();
+        case FftBackend::HeFFTe_Sycl_Rocfft: return checkMpiHipAwareSupport();
+        case FftBackend::HeFFTe_Sycl_OneMkl: return checkMpiZEAwareSupport();
+        default: return GpuAwareMpiStatus::NotSupported;
+    }
+}
+
 /*! \brief Check that the real grid after forward and backward
  * 3D transforms matches the input real grid. */
 static void checkRealGrid(const IVec           realGridSizeFull,
@@ -115,6 +134,11 @@ class GpuFftTest3D : public ::testing::Test, public ::testing::WithParamInterfac
     {
         const auto& deviceList = getTestHardwareEnvironment()->getTestDeviceList();
 
+        if (deviceList.empty())
+        {
+            GTEST_SKIP() << "No compatible GPUs detected";
+        }
+
         int rank;
         MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 
@@ -133,6 +157,13 @@ class GpuFftTest3D : public ::testing::Test, public ::testing::WithParamInterfac
         std::tie(gridParams, backend)                        = param;
         std::tie(realGridSizeFull, numDomainsX, numDomainsY) = gridParams;
 
+        const GpuAwareMpiStatus gpuAwareMpiStatus = getGpuAwareMpiStatusForFftBackend(backend);
+        if (gpuAwareMpiStatus != GpuAwareMpiStatus::Supported
+            && gpuAwareMpiStatus != GpuAwareMpiStatus::Forced)
+        {
+            GTEST_SKIP() << "Test requires GPU-aware MPI library";
+        }
+
         // define local grid sizes - this follows same logic as GROMACS implementation
         std::vector<int> localGridSizesX(numDomainsX);
         for (unsigned int i = 0; i < localGridSizesX.size(); ++i)

src/programs/mdrun/tests/pmetest.cpp

@@ -57,6 +57,7 @@
 #include "gromacs/utility/cstringutil.h"
 #include "gromacs/utility/enumerationhelpers.h"
 #include "gromacs/utility/message_string_collector.h"
+#include "gromacs/utility/mpiinfo.h"
 #include "gromacs/utility/stringutil.h"
 
 #include "testutils/mpitest.h"
@@ -261,10 +262,25 @@ MessageStringCollector PmeTest::getSkipMessagesIfNecessary(const CommandLine& co
     {
         messages.appendIf(getCompatibleDevices(s_hwinfo->deviceInfoList).empty(),
                           "it targets GPU execution, but no compatible devices were detected");
-        if ((getenv("GMX_GPU_PME_DECOMPOSITION")) == nullptr)
+
+        if (!commandLineTargetsPmeOnlyRanks && numRanks > 1)
         {
-            messages.appendIf(!commandLineTargetsPmeOnlyRanks && numRanks > 1,
+            const bool pmeDecompositionSupported = GMX_USE_cuFFTMp || GMX_USE_Heffte;
+            messages.appendIf(!pmeDecompositionSupported,
                               "it targets PME decomposition, but that is not supported");
+            if (pmeDecompositionSupported)
+            {
+                const bool pmeDecompositionActive = (getenv("GMX_GPU_PME_DECOMPOSITION") != nullptr);
+                messages.appendIf(!pmeDecompositionActive,
+                                  "it targets PME decomposition, but that is not enabled");
+                // The check below only handles CUDA, see #4638
+                GpuAwareMpiStatus gpuAwareMpiStatus = checkMpiCudaAwareSupport();
+                const bool        gpuAwareMpiActive = gpuAwareMpiStatus == GpuAwareMpiStatus::Forced
+                                               || gpuAwareMpiStatus == GpuAwareMpiStatus::Supported;
+                messages.appendIf(!gpuAwareMpiActive,
+                                  "it targets PME decomposition, which requires GPU-aware MPI, but "
+                                  "that is not detected");
+            }
         }
 
         std::optional<std::string_view> pmeFftOptionArgument = commandLine.argumentOf("-pmefft");