Merge pull request #17 from yaoyic/upstream-master

[BUGFIX] Synchronize CUDA context before force copying

platforms/cuda/src/CudaTorchKernels.cpp

@@ -33,11 +33,21 @@
 #include "CudaTorchKernelSources.h"
 #include "openmm/internal/ContextImpl.h"
 #include <map>
+#include <cuda_runtime_api.h>
 
 using namespace TorchPlugin;
 using namespace OpenMM;
 using namespace std;
 
+// macro for checking the result of synchronization operation on CUDA
+// copied from `openmm/platforms/cuda/src/CudaParallelKernels.cpp`
+#define CHECK_RESULT(result, prefix) \
+if (result != CUDA_SUCCESS) { \
+    std::stringstream m; \
+    m<<prefix<<": "<<cu.getErrorString(result)<<" ("<<result<<")"<<" at "<<__FILE__<<":"<<__LINE__; \
+    throw OpenMMException(m.str());\
+}
+
 CudaCalcTorchForceKernel::~CudaCalcTorchForceKernel() {
 }
 
@@ -80,11 +90,17 @@ double CudaCalcTorchForceKernel::execute(ContextImpl& context, bool includeForce
     vector<torch::jit::IValue> inputs = {posTensor};
     if (usePeriodic)
         inputs.push_back(boxTensor);
+    // synchronizing the current context before switching to PyTorch
+    CHECK_RESULT(cuCtxSynchronize(), "Error synchronizing CUDA context");
     torch::Tensor energyTensor = module.forward(inputs).toTensor();
     if (includeForces) {
         energyTensor.backward();
         // Note: "forceTensor" needs to be cloned due to a shared context (https://github.com/openmm/openmm-torch/issues/13)
         torch::Tensor forceTensor = posTensor.grad().clone();
+        // make sure that all calculations on PyTorch side is properly finished before changing CUDA context or starting the `addForcesKernel` of this plugin
+        // cudaDeviceSynchronize();  // synchronizing the whole device is not necessary and may even cause problem
+        // synchronizing the current context and check the return status
+        CHECK_RESULT(cuCtxSynchronize(), "Error synchronizing CUDA context");
         cu.setAsCurrent();
         void* data;
         if (cu.getUseDoublePrecision()) {