Fix UntypedStorage.resize_ to keep same CUDA device index (#113386)

kurtamohler · pytorchmergebot · commit 8bdce9bb74e7 · 2023-11-10T01:57:25.000Z
Fixes #113300 Pull Request resolved: #113386 Approved by: https://github.com/albanD
diff --git a/aten/src/ATen/native/cuda/Resize.cpp b/aten/src/ATen/native/cuda/Resize.cpp
@@ -4,6 +4,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/PeerToPeerAccess.h>
 #include <ATen/native/ResizeCommon.h>
+#include <c10/cuda/CUDAGuard.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
 #include <ATen/NativeFunctions.h>
@@ -18,18 +19,18 @@ void resize_bytes_cuda(StorageImpl* storage, size_t size_bytes) {
   auto allocator = storage->allocator();
   TORCH_CHECK(allocator != nullptr, "Trying to resize storage without an allocator");
 
-  auto device = at::cuda::current_device();
+  c10::Device device = storage->device();
+
   if (size_bytes == 0) {
-    storage->set_data_ptr_noswap(at::DataPtr(nullptr, at::Device(at::DeviceType::CUDA, device)));
+    storage->set_data_ptr_noswap(at::DataPtr(nullptr, device));
     storage->set_nbytes(0);
     return;
   }
 
+  c10::cuda::CUDAGuard guard(device.index());
   at::DataPtr data = allocator->allocate(size_bytes);
   if (storage->data_ptr()) {
-    // Enable p2p access when the memcpy is across devices
     at::globalContext().lazyInitCUDA();
-    at::cuda::get_p2p_access(device, storage->device().index());
 
     C10_CUDA_CHECK(
         cudaMemcpyAsync(
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -6783,6 +6783,33 @@ def test_from_buffer(self):
         self.assertEqual(bytes.tolist(), [1, 2, 3, 4])
         self.assertTrue(isinstance(bytes, torch.ByteStorage))
 
+    # Check that after `UntypedStorage.resize_` is called, the storage remains
+    # on the same CUDA device index it was initialized on, regardless of the
+    # default index
+    @unittest.skipIf(torch.cuda.device_count() < 2, "Requires 2 GPUs")
+    def test_untyped_storage_resize_cuda_device(self):
+        test_cases = [
+            # start_size, storage_resize
+            ((2, 3), 0),
+            ((2, 3), 100),
+            ((2, 3), 10),
+            (0, 10),
+            (0, 0),
+        ]
+        default_indices = [0, 1]
+        devices = [
+            torch.device('cuda:0'),
+            torch.device('cuda:1')]
+
+        for default_idx, device, (start_size, storage_resize) in product(default_indices, devices, test_cases):
+            with torch.cuda.device(default_idx):
+                a = torch.zeros(start_size, device=device)
+                a.untyped_storage().resize_(storage_resize)
+
+            self.assertEqual(a.device, device)
+            self.assertEqual(a.untyped_storage().device, device)
+            self.assertEqual(a.untyped_storage().nbytes(), storage_resize)
+
     def test_storage_error(self):
         quantized_storages = [
             torch.QInt32Storage,
