Relocate GPU initialization for PME

We need to have set up the GPU we intend to use before we would want
to e.g. compile kernels for it. We will shortly want to do this
for OpenCL support for PME.

Change-Id: Iaf22ecba8fb3ad7583b57958742abb018d1e8a0d

diff --git a/src/gromacs/ewald/pme.cu b/src/gromacs/ewald/pme.cu
index 26a199dfef53c624aee016c51ff30bcabf70a8c3..9d1d6c7cfe65cb2a2795b29b0ece90b326865d83 100644 (file)
--- a/src/gromacs/ewald/pme.cu
+++ b/src/gromacs/ewald/pme.cu
@@ -458,6 +458,9 @@ void pme_gpu_init_internal(PmeGpu *pmeGpu)
     // TODO: Consider turning on by default when we can detect nr of streams.
     pmeGpu->archSpecific->useTiming = (getenv("GMX_ENABLE_GPU_TIMING") != nullptr);
 
+    // Prepare to use the device that this PME task was assigned earlier.
+    CU_RET_ERR(cudaSetDevice(pmeGpu->deviceInfo->id), "Switching to PME CUDA device");
+
     pmeGpu->maxGridWidthX = pmeGpu->deviceInfo->prop.maxGridSize[0];
 
     /* Creating a PME CUDA stream */

diff --git a/src/gromacs/ewald/tests/pmetestcommon.cpp b/src/gromacs/ewald/tests/pmetestcommon.cpp
index 781bf53a61a18460a193a7bb1f552f676e4038e9..b30a8c77a511a2bbe3660c4befdd596b902c0ad0 100644 (file)
--- a/src/gromacs/ewald/tests/pmetestcommon.cpp
+++ b/src/gromacs/ewald/tests/pmetestcommon.cpp
@@ -109,10 +109,6 @@ static PmeSafePointer pmeInitInternal(const t_inputrec         *inputRec,
                                       )
 {
     const MDLogger dummyLogger;
-    if (gpuInfo)
-    {
-        init_gpu(dummyLogger, gpuInfo);
-    }
     const auto     runMode       = (mode == CodePath::CPU) ? PmeRunMode::CPU : PmeRunMode::GPU;
     t_commrec      dummyCommrec  = {0};
     NumPmeDomains  numPmeDomains = { 1, 1 };

diff --git a/src/gromacs/ewald/tests/testhardwarecontexts.cpp b/src/gromacs/ewald/tests/testhardwarecontexts.cpp
index a4f0c7cfce2294c1579f646aa90ca06d042bc70f..29e4e31ea843eef53fbe634004779784281c73ce 100644 (file)
--- a/src/gromacs/ewald/tests/testhardwarecontexts.cpp
+++ b/src/gromacs/ewald/tests/testhardwarecontexts.cpp
@@ -114,14 +114,18 @@ void PmeTestEnvironment::SetUp()
         // PME can only run on the CPU, so don't make any more test contexts.
         return;
     }
+    const MDLogger dummyLogger;
     // Constructing contexts for all compatible GPUs - will be empty on non-GPU builds
     for (int gpuIndex : getCompatibleGpus(hardwareInfo_->gpu_info))
     {
+        gmx_device_info_t *deviceInfo = getDeviceInfo(hardwareInfo_->gpu_info, gpuIndex);
+        init_gpu(dummyLogger, deviceInfo);
+
         char        stmp[200] = {};
         get_gpu_device_info_string(stmp, hardwareInfo_->gpu_info, gpuIndex);
         std::string description = "(GPU " + std::string(stmp) + ") ";
         // TODO should this be CodePath::GPU?
-        hardwareContexts_.emplace_back(TestHardwareContext(CodePath::CUDA, description.c_str(), getDeviceInfo(hardwareInfo_->gpu_info, gpuIndex)));
+        hardwareContexts_.emplace_back(TestHardwareContext(CodePath::CUDA, description.c_str(), deviceInfo));
     }
 }