final touches

Aidyn-A · Aidyn-A · commit 88d7b8bbbe07 · 2023-04-07T11:47:08.000-07:00
diff --git a/c10/cuda/CUDAFunctions.cpp b/c10/cuda/CUDAFunctions.cpp
@@ -204,7 +204,7 @@ cudaError_t GetDeviceCount(int* dev_count) {
 // x = torch.empty(1, device=“cuda:1”) # no CUDA context on cuda:0 after this
 // call y = torch.empty(1, device=“cuda”) # CUDA context is created on cuda:0
 // ```
-#if CUDA_VERSION >= 11000
+#if CUDA_VERSION >= 12000
 thread_local int targetDeviceIndex = -1;
 
 cudaError_t GetDevice(int* device) {
@@ -223,9 +223,7 @@ cudaError_t SetDevice(int device) {
   if (device == cur_device) {
     return cudaSuccess;
   }
-  cudaError_t err = cudaSetDevice(device);
-  C10_CUDA_CHECK(cudaFree(0));
-  return err;
+  return cudaSetDevice(device);
 }
 
 cudaError_t MaybeSetDevice(int device) {
@@ -236,6 +234,8 @@ cudaError_t MaybeSetDevice(int device) {
   return cudaSuccess;
 }
 
+// This function always initializes the CUDA context
+// on to_device
 int ExchangeDevice(int to_device) {
   int cur_device = targetDeviceIndex;
   targetDeviceIndex = -1;
@@ -246,10 +246,11 @@ int ExchangeDevice(int to_device) {
     }
   }
   C10_CUDA_CHECK(cudaSetDevice(to_device));
-  C10_CUDA_CHECK(cudaFree(0));
   return cur_device;
 }
 
+// This function does not initialize the CUDA context
+// on to_device if it does not already exist
 int MaybeExchangeDevice(int to_device) {
   int cur_device = -1;
   C10_CUDA_CHECK(cudaGetDevice(&cur_device));
@@ -258,7 +259,6 @@ int MaybeExchangeDevice(int to_device) {
   }
   if (hasPrimaryContext(to_device)) {
     C10_CUDA_CHECK(cudaSetDevice(to_device));
-    C10_CUDA_CHECK(cudaFree(0));
   } else {
     targetDeviceIndex = to_device;
   }

Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ cudaError_t GetDeviceCount(int* dev_count) {`
`204`	`204`	`// x = torch.empty(1, device=“cuda:1”) # no CUDA context on cuda:0 after this`
`205`	`205`	`// call y = torch.empty(1, device=“cuda”) # CUDA context is created on cuda:0`
`206`	`206`	// ```
`207`		`-#if CUDA_VERSION >= 11000`
	`207`	`+#if CUDA_VERSION >= 12000`
`208`	`208`	`thread_local int targetDeviceIndex = -1;`
`209`	`209`
`210`	`210`	`cudaError_t GetDevice(int* device) {`
`@@ -223,9 +223,7 @@ cudaError_t SetDevice(int device) {`
`223`	`223`	`if (device == cur_device) {`
`224`	`224`	`return cudaSuccess;`
`225`	`225`	`}`
`226`		`- cudaError_t err = cudaSetDevice(device);`
`227`		`- C10_CUDA_CHECK(cudaFree(0));`
`228`		`- return err;`
	`226`	`+ return cudaSetDevice(device);`
`229`	`227`	`}`
`230`	`228`
`231`	`229`	`cudaError_t MaybeSetDevice(int device) {`
`@@ -236,6 +234,8 @@ cudaError_t MaybeSetDevice(int device) {`
`236`	`234`	`return cudaSuccess;`
`237`	`235`	`}`
`238`	`236`
	`237`	`+// This function always initializes the CUDA context`
	`238`	`+// on to_device`
`239`	`239`	`int ExchangeDevice(int to_device) {`
`240`	`240`	`int cur_device = targetDeviceIndex;`
`241`	`241`	`targetDeviceIndex = -1;`
`@@ -246,10 +246,11 @@ int ExchangeDevice(int to_device) {`
`246`	`246`	`}`
`247`	`247`	`}`
`248`	`248`	`C10_CUDA_CHECK(cudaSetDevice(to_device));`
`249`		`- C10_CUDA_CHECK(cudaFree(0));`
`250`	`249`	`return cur_device;`
`251`	`250`	`}`
`252`	`251`
	`252`	`+// This function does not initialize the CUDA context`
	`253`	`+// on to_device if it does not already exist`
`253`	`254`	`int MaybeExchangeDevice(int to_device) {`
`254`	`255`	`int cur_device = -1;`
`255`	`256`	`C10_CUDA_CHECK(cudaGetDevice(&cur_device));`
`@@ -258,7 +259,6 @@ int MaybeExchangeDevice(int to_device) {`
`258`	`259`	`}`
`259`	`260`	`if (hasPrimaryContext(to_device)) {`
`260`	`261`	`C10_CUDA_CHECK(cudaSetDevice(to_device));`
`261`		`- C10_CUDA_CHECK(cudaFree(0));`
`262`	`262`	`} else {`
`263`	`263`	`targetDeviceIndex = to_device;`
`264`	`264`	`}`