Make DeviceStream into a class

[alexxy/gromacs.git] / src / gromacs / ewald / pme_only.cpp

diff --git a/src/gromacs/ewald/pme_only.cpp b/src/gromacs/ewald/pme_only.cpp
index 845b1a33ecf45b99b78d846f0107533647e4f485..fe51deb5fc4dc322b51db151006f6ad4e1eeadc9 100644 (file)
--- a/src/gromacs/ewald/pme_only.cpp
+++ b/src/gromacs/ewald/pme_only.cpp
@@ -629,16 +629,16 @@ int gmx_pmeonly(struct gmx_pme_t*         pme,
     const bool useGpuForPme = (runMode == PmeRunMode::GPU) || (runMode == PmeRunMode::Mixed);
     if (useGpuForPme)
     {
-        const void* commandStream = pme_gpu_get_device_stream(pme);
+        const DeviceStream& deviceStream = *pme_gpu_get_device_stream(pme);
 
         changePinningPolicy(&pme_pp->chargeA, pme_get_pinning_policy());
         changePinningPolicy(&pme_pp->x, pme_get_pinning_policy());
         if (c_enableGpuPmePpComms)
         {
             pme_pp->pmeCoordinateReceiverGpu = std::make_unique<gmx::PmeCoordinateReceiverGpu>(
-                    commandStream, pme_pp->mpi_comm_mysim, pme_pp->ppRanks);
+                    deviceStream, pme_pp->mpi_comm_mysim, pme_pp->ppRanks);
             pme_pp->pmeForceSenderGpu = std::make_unique<gmx::PmeForceSenderGpu>(
-                    commandStream, pme_pp->mpi_comm_mysim, pme_pp->ppRanks);
+                    deviceStream, pme_pp->mpi_comm_mysim, pme_pp->ppRanks);
         }
         GMX_RELEASE_ASSERT(
                 deviceContext != nullptr,
@@ -646,7 +646,7 @@ int gmx_pmeonly(struct gmx_pme_t*         pme,
         // TODO: Special PME-only constructor is used here. There is no mechanism to prevent from using the other constructor here.
         //       This should be made safer.
         stateGpu = std::make_unique<gmx::StatePropagatorDataGpu>(
-                commandStream, *deviceContext, GpuApiCallBehavior::Async,
+                &deviceStream, *deviceContext, GpuApiCallBehavior::Async,
                 pme_gpu_get_padding_size(pme), wcycle);
     }