Remove CudaContextManager and add cudaSetDevice(m_Device)

RTKConsortium · Feb 9, 2024 · dba5972 · dba5972
1 parent e045737
commit dba5972
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 158 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,6 +76,7 @@ CudaImageDataManager<ImageType>::MakeCPUBufferUpToDate()
       std::cout << this << ": GPU->CPU data copy" << std::endl;
 #endif
 
+      CUDA_CHECK(cudaSetDevice(m_Device));
       errid = cudaMemcpy(m_CPUBuffer, m_GPUBuffer->GetPointer(), m_BufferSize, cudaMemcpyDeviceToHost);
       CudaCheckError(errid, __FILE__, __LINE__, ITK_LOCATION);
 
@@ -114,6 +115,7 @@ CudaImageDataManager<ImageType>::MakeGPUBufferUpToDate()
       std::cout << "CPU->GPU data copy" << std::endl;
 #endif
 
+      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/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,6 +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
+  m_Device = 0;
+  CUDA_CHECK(cudaSetDevice(m_Device));
 
   m_CPUBuffer = nullptr;
   m_GPUBuffer = GPUMemPointer::New();
@@ -53,7 +54,6 @@ CudaDataManager::CudaDataManager()
 CudaDataManager::~CudaDataManager()
 {
   m_GPUBuffer = nullptr;
-  CudaContextManager::DestroyInstance();
 }
 
 void
@@ -88,6 +88,7 @@ CudaDataManager::Free()
   {
     try
     {
+      CUDA_CHECK(cudaSetDevice(m_Device));
       m_GPUBuffer->Free();
     }
     catch (itk::ExceptionObject & e)
@@ -165,6 +166,7 @@ CudaDataManager::UpdateCPUBuffer()
       std::cout << this << "::UpdateCPUBuffer GPU->CPU data copy " << m_GPUBuffer->GetPointer() << "->" << m_CPUBuffer
                 << " : " << m_BufferSize << std::endl;
 #endif
+      CUDA_CHECK(cudaSetDevice(m_Device));
       CUDA_CHECK(cudaMemcpy(m_CPUBuffer, m_GPUBuffer->GetPointer(), m_BufferSize, cudaMemcpyDeviceToHost));
       m_IsCPUBufferDirty = false;
     }
@@ -203,6 +205,7 @@ CudaDataManager::UpdateGPUBuffer()
       std::cout << this << "::UpdateGPUBuffer CPU->GPU data copy " << m_CPUBuffer << "->" << m_GPUBuffer->GetPointer()
                 << " : " << m_BufferSize << std::endl;
 #endif
+      CUDA_CHECK(cudaSetDevice(m_Device));
       CUDA_CHECK(cudaMemcpy(m_GPUBuffer->GetPointer(), m_CPUBuffer, m_BufferSize, cudaMemcpyHostToDevice));
     }
     m_IsGPUBufferDirty = false;
@@ -247,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;