Use CUDA primary device context instead of creating our own context

AnyDSL · Jul 11, 2024 · 7042fbf · 7042fbf
1 parent 213f959
commit 7042fbf
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/cuda_platform.cpp b/src/cuda_platform.cpp
@@ -110,8 +110,8 @@ CudaPlatform::CudaPlatform(Runtime* runtime)
         devices_[i].compute_capability = (CUjit_target)(major * 10 + minor);
         debug("  (%) %, Compute capability: %.%", i, name, major, minor);
 
-        err = cuCtxCreate(&devices_[i].ctx, CU_CTX_MAP_HOST, devices_[i].dev);
-        CHECK_CUDA(err, "cuCtxCreate()");
+        err = cuDevicePrimaryCtxRetain(&devices_[i].ctx, devices_[i].dev);
+        CHECK_CUDA(err, "cuDevicePrimaryCtxRetain()");
     }
 }
 
@@ -139,7 +139,7 @@ void CudaPlatform::erase_profiles(bool erase_all) {
 CudaPlatform::~CudaPlatform() {
     erase_profiles(true);
     for (size_t i = 0; i < devices_.size(); i++)
-        cuCtxDestroy(devices_[i].ctx);
+        cuDevicePrimaryCtxRelease(devices_[i].dev);
 }
 
 void* CudaPlatform::alloc(DeviceId dev, int64_t size) {