BUG: Remove CudaContextManager class and use cudaSetDevice

Prevent the destruction of the Cuda context during use (e.g. cupy package) CHANGES: * Remove CudaContextManager class * Replace cuCtxSetCurrent() by cudaSetDevice() to create the primary context introduced by Cuda 7 * Add m_Device in CudaDataManager constructor * Update CMakeLists.txt
RTKConsortium · Feb 12, 2024 · 3c1b5bc · 3c1b5bc
1 parent 9d3fe9b
commit 3c1b5bc
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 199 deletions.
diff --git a/include/itkCudaContextManager.h b/include/itkCudaContextManager.h
diff --git a/include/itkCudaDataManager.h b/include/itkCudaDataManager.h
@@ -232,6 +232,8 @@ class CudaCommon_EXPORT CudaDataManager : public Object
   void
   PrintSelf(std::ostream & os, Indent indent) const override;
 
+  int m_Device;
+
 private:
   CudaDataManager(const Self &) = delete; // purposely not implemented
   void
@@ -240,8 +242,6 @@ class CudaCommon_EXPORT CudaDataManager : public Object
 protected:
   size_t m_BufferSize; // # of bytes
 
-  CudaContextManager * m_ContextManager;
-
   /** buffer type */
   int m_MemFlags;
 

diff --git a/include/itkCudaImageDataManager.hxx b/include/itkCudaImageDataManager.hxx
@@ -76,9 +76,7 @@ CudaImageDataManager<ImageType>::MakeCPUBufferUpToDate()
       std::cout << this << ": GPU->CPU data copy" << std::endl;
 #endif
 
-      CUDA_CHECK(cuCtxSetCurrent(
-        *(this->m_ContextManager->GetCurrentContext()))); // This is necessary when running multithread to bind the host
-                                                          // CPU thread to the right context
+      CUDA_CHECK(cudaSetDevice(m_Device));
       errid = cudaMemcpy(m_CPUBuffer, m_GPUBuffer->GetPointer(), m_BufferSize, cudaMemcpyDeviceToHost);
       CudaCheckError(errid, __FILE__, __LINE__, ITK_LOCATION);
 
@@ -117,9 +115,7 @@ CudaImageDataManager<ImageType>::MakeGPUBufferUpToDate()
       std::cout << "CPU->GPU data copy" << std::endl;
 #endif
 
-      CUDA_CHECK(cuCtxSetCurrent(
-        *(this->m_ContextManager->GetCurrentContext()))); // This is necessary when running multithread to bind the host
-                                                          // CPU thread to the right context
+      CUDA_CHECK(cudaSetDevice(m_Device));
       errid = cudaMemcpy(m_GPUBuffer->GetPointer(), m_CPUBuffer, m_BufferSize, cudaMemcpyHostToDevice);
       CudaCheckError(errid, __FILE__, __LINE__, ITK_LOCATION);
 

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -1,5 +1,4 @@
 set(CudaCommon_SRCS
-  itkCudaContextManager.cxx
   itkCudaDataManager.cxx
   itkCudaUtil.cxx
   itkCudaMemoryProbe.cxx

diff --git a/src/itkCudaContextManager.cxx b/src/itkCudaContextManager.cxx
diff --git a/src/itkCudaDataManager.cxx b/src/itkCudaDataManager.cxx
@@ -24,7 +24,6 @@ namespace itk
 // constructor
 CudaDataManager::CudaDataManager()
 {
-  m_ContextManager = CudaContextManager::GetInstance();
 
   // Creating the context in the constructor allows avoiding a memory leak.
   // However, the cuda data manager is created even if there is no use of CUDA
@@ -33,9 +32,8 @@ CudaDataManager::CudaDataManager()
   // is no CUDA device available, we just do not set the context (SR). This fixes
   // the problem reported here:
   // https://www.creatis.insa-lyon.fr/pipermail/rtk-users/2015-July/000570.html
-  CUcontext * ctx = m_ContextManager->GetCurrentContext();
-  if (ctx)
-    CUDA_CHECK(cuCtxSetCurrent(*ctx));
+  m_Device = 0;
+  CUDA_CHECK(cudaSetDevice(m_Device));
 
   m_CPUBuffer = nullptr;
   m_GPUBuffer = GPUMemPointer::New();
@@ -56,7 +54,6 @@ CudaDataManager::CudaDataManager()
 CudaDataManager::~CudaDataManager()
 {
   m_GPUBuffer = nullptr;
-  CudaContextManager::DestroyInstance();
 }
 
 void
@@ -91,9 +88,7 @@ CudaDataManager::Free()
   {
     try
     {
-      CUDA_CHECK(cuCtxSetCurrent(
-        *(this->m_ContextManager->GetCurrentContext()))); // This is necessary when running multithread to bind the host
-                                                          // CPU thread to the right context
+      CUDA_CHECK(cudaSetDevice(m_Device));
       m_GPUBuffer->Free();
     }
     catch (itk::ExceptionObject & e)
@@ -171,9 +166,7 @@ CudaDataManager::UpdateCPUBuffer()
       std::cout << this << "::UpdateCPUBuffer GPU->CPU data copy " << m_GPUBuffer->GetPointer() << "->" << m_CPUBuffer
                 << " : " << m_BufferSize << std::endl;
 #endif
-      CUDA_CHECK(cuCtxSetCurrent(
-        *(this->m_ContextManager->GetCurrentContext()))); // This is necessary when running multithread to bind the host
-                                                          // CPU thread to the right context
+      CUDA_CHECK(cudaSetDevice(m_Device));
       CUDA_CHECK(cudaMemcpy(m_CPUBuffer, m_GPUBuffer->GetPointer(), m_BufferSize, cudaMemcpyDeviceToHost));
       m_IsCPUBufferDirty = false;
     }
@@ -212,9 +205,7 @@ CudaDataManager::UpdateGPUBuffer()
       std::cout << this << "::UpdateGPUBuffer CPU->GPU data copy " << m_CPUBuffer << "->" << m_GPUBuffer->GetPointer()
                 << " : " << m_BufferSize << std::endl;
 #endif
-      CUDA_CHECK(cuCtxSetCurrent(
-        *(this->m_ContextManager->GetCurrentContext()))); // This is necessary when running multithread to bind the host
-                                                          // CPU thread to the right context
+      CUDA_CHECK(cudaSetDevice(m_Device));
       CUDA_CHECK(cudaMemcpy(m_GPUBuffer->GetPointer(), m_CPUBuffer, m_BufferSize, cudaMemcpyHostToDevice));
     }
     m_IsGPUBufferDirty = false;
@@ -259,7 +250,6 @@ CudaDataManager::Graft(const CudaDataManager * data)
   if (data)
   {
     m_BufferSize = data->m_BufferSize;
-    m_ContextManager = data->m_ContextManager;
     m_GPUBuffer = data->m_GPUBuffer;
     m_CPUBuffer = data->m_CPUBuffer;
     m_IsCPUBufferDirty = data->m_IsCPUBufferDirty;