From 81d99a94f2e8b1c91c878db1d874c8c9286e170d Mon Sep 17 00:00:00 2001
From: Liam Keegan <liam@keegan.ch>
Date: Wed, 27 Sep 2023 16:42:54 +0200
Subject: [PATCH] Disable `cudaDeviceReset()` calls in pmcx

- If `MCX_DISABLE_CUDA_DEVICE_RESET` is defined then `cudaDeviceReset()` calls are not made in mcx_core.cu
- cudaDeviceReset() destroys all resources on the GPU associated with the current process
- for pmcx this can leave other python GPU libraries like pytorch in a broken state
- resolves #187
---
 pmcx/setup.py   | 1 +
 src/mcx_core.cu | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/pmcx/setup.py b/pmcx/setup.py
index 051fed40..d9628d4b 100644
--- a/pmcx/setup.py
+++ b/pmcx/setup.py
@@ -49,6 +49,7 @@ def build_extension(self, ext):
             f"-DPYTHON_EXECUTABLE={sys.executable}",
             f"-DCMAKE_BUILD_TYPE={cfg}",  # not used on MSVC, but no harm
             "-DBUILD_PYTHON=true",
+            "-DCUDA_NVCC_FLAGS=-DMCX_DISABLE_CUDA_DEVICE_RESET",
         ]
         build_args = []
         # Adding CMake arguments set as environment variable
diff --git a/src/mcx_core.cu b/src/mcx_core.cu
index 803d9370..67208fd0 100644
--- a/src/mcx_core.cu
+++ b/src/mcx_core.cu
@@ -2275,7 +2275,9 @@ __global__ void mcx_main_loop(uint media[], OutputType field[], float genergy[],
  */
 void mcx_cu_assess(cudaError_t cuerr, const char* file, const int linenum) {
     if (cuerr != cudaSuccess) {
+#ifndef MCX_DISABLE_CUDA_DEVICE_RESET
         cudaDeviceReset();
+#endif
         mcx_error(-(int)cuerr, (char*)cudaGetErrorString(cuerr), file, linenum);
     }
 }
@@ -2816,7 +2818,10 @@ void mcx_run_simulation(Config* cfg, GPUInfo* gpu) {
             free(field);
             free(Pseed);
 
+#ifndef MCX_DISABLE_CUDA_DEVICE_RESET
             CUDA_ASSERT(cudaDeviceReset());
+#endif
+
         }
         #pragma omp barrier
 
@@ -3765,7 +3770,9 @@ is more than what your have specified (%d), please use the -H option to specify
     /**
      * The below call in theory is not needed, but it ensures the device is freed for other programs, especially on Windows
      */
+#ifndef MCX_DISABLE_CUDA_DEVICE_RESET
     CUDA_ASSERT(cudaDeviceReset());
+#endif
 
     /**
      * Lastly, free all host buffers, the simulation is complete.