Update

[ghstack-poisoned]
pytorch · kurtamohler · Mar 4, 2025 · Mar 7, 2025 · Mar 7, 2025 · Mar 7, 2025
commit 7c0f0820bd0f4ec98b7cb10a47e5a2c11bcc3f5d
diff --git a/c10/core/impl/COW.cpp b/c10/core/impl/COW.cpp
@@ -1,15 +1,19 @@
 #include <c10/core/impl/COW.h>
 
 #include <c10/core/Allocator.h>
+#include <c10/core/DeviceGuard.h>
 #include <c10/core/StorageImpl.h>
 #include <c10/core/alignment.h>
 #include <c10/core/impl/COWDeleter.h>
-#include <c10/cuda/CUDACachingAllocator.h>
-#include <c10/cuda/CUDAFunctions.h>
 #include <c10/util/Exception.h>
 #include <c10/util/ParallelGuard.h>
 #include <c10/util/UniqueVoidPtr.h>
 
+#ifdef USE_CUDA
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAFunctions.h>
+#endif
+
 #include <memory>
 #include <optional>
 
@@ -121,9 +125,12 @@ c10::intrusive_ptr<StorageImpl> lazy_clone_storage(
           c10::impl::cow::cow_deleter);
       device_type = device.type();
 
+#ifdef USE_CUDA
       if (device_type == c10::kCUDA) {
         allocator = c10::cuda::CUDACachingAllocator::get();
-      } else {
+      } else
+#endif
+      {
         allocator = c10::GetAllocator(device.type());
       }
 
@@ -170,9 +177,11 @@ C10_API void materialize_cow_storage(StorageImpl& storage) {
        // that the data will remain while we copy it.
     new_data_ptr = storage.allocator()->clone(data_ptr.get(), storage.nbytes());
     if (!devices_match) {
+#ifdef USE_CUDA
       if (storage.device().type() == c10::kCUDA) {
         c10::cuda::device_synchronize();
       }
+#endif
     }
   }
 

diff --git a/c10/test/core/impl/cow_test.cpp b/c10/test/core/impl/cow_test.cpp
@@ -45,7 +45,8 @@ class ContextTest : public testing::Test {
 };
 
 TEST_F(ContextTest, Basic) {
-  auto& context = *new cow::COWDeleterContext(new_delete_tracker());
+  Device device(c10::kCPU);
+  auto& context = *new cow::COWDeleterContext(new_delete_tracker(), device);
   ASSERT_THAT(delete_count(), testing::Eq(0));
 
   context.increment_refcount();
@@ -76,7 +77,8 @@ TEST_F(ContextTest, Basic) {
 
 TEST_F(ContextTest, cow_deleter) {
   // This is effectively the same thing as decrement_refcount() above.
-  auto& context = *new cow::COWDeleterContext(new_delete_tracker());
+  Device device(c10::kCPU);
+  auto& context = *new cow::COWDeleterContext(new_delete_tracker(), device);
   ASSERT_THAT(delete_count(), testing::Eq(0));
 
   cow::cow_deleter(&context);
@@ -150,14 +152,15 @@ TEST(lazy_clone_storage_test, already_copy_on_write) {
       new std::byte[5],
       +[](void* bytes) { delete[] static_cast<std::byte*>(bytes); });
   void* data_ptr = data.get();
+  Device device(c10::kCPU);
   StorageImpl original_storage(
       {},
       /*size_bytes=*/5,
       at::DataPtr(
           /*data=*/data_ptr,
-          /*ctx=*/new cow::COWDeleterContext(std::move(data)),
+          /*ctx=*/new cow::COWDeleterContext(std::move(data), device),
           cow::cow_deleter,
-          Device(Device::Type::CPU)),
+          device),
       /*allocator=*/nullptr,
       /*resizable=*/false);
 
@@ -193,14 +196,15 @@ TEST(materialize_test, copy_on_write_single_reference) {
       new std::byte[4],
       +[](void* bytes) { delete[] static_cast<std::byte*>(bytes); });
   void* data_ptr = data.get();
+  Device device(c10::kCPU);
   StorageImpl storage(
       {},
       /*size_bytes=*/4,
       at::DataPtr(
           /*data=*/data_ptr,
-          /*ctx=*/new cow::COWDeleterContext(std::move(data)),
+          /*ctx=*/new cow::COWDeleterContext(std::move(data), device),
           cow::cow_deleter,
-          Device(Device::Type::CPU)),
+          device),
       /*allocator=*/nullptr,
       /*resizable=*/false);
 

diff --git a/test/test_torch.py b/test/test_torch.py
@@ -5342,8 +5342,23 @@ def test_lazy_clone_to_device(self, device):
             ('cpu', 'cuda:1'),
             ('cuda:1', 'cuda:0'),
             ('cuda:0', 'cuda:1'),
-            # TODO: Figure out why CUDA to CPU segfaults
             # ('cuda', 'cpu'),
+            # NOTE: CUDA -> CPU doesn't work at the moment. Traceback shows that
+            # apparently the CPU Allocator's `clone` method doesn't know how to
+            # deal with a CUDA data pointer. May not need to solve this problem
+            # at the moment, because the whole point is MPS-CPU, not CUDA-CPU.
+            #
+            #   Thread 1 "python" received signal SIGSEGV, Segmentation fault.
+            #   __memcpy_avx_unaligned () at ../sysdeps/x86_64/multiarch/memmove-vec-unaligned-erms.S:220
+            #   warning: 220    ../sysdeps/x86_64/multiarch/memmove-vec-unaligned-erms.S: No such file or directory
+            #   (gdb) bt
+            #   #0  __memcpy_avx_unaligned () at ../sysdeps/x86_64/multiarch/memmove-vec-unaligned-erms.S:220
+            #   #1  0x00007ffff3d22e40 in c10::Allocator::clone (this=0x7ffff3df8010 <c10::g_cpu_alloc>,
+            #       data=0x7ffc89000200, n=40)
+            #       at /home/kurtamohler/develop/pytorch-0/c10/util/UniqueVoidPtr.h:61
+            #   #2  0x00007ffff3d751cf in c10::impl::cow::materialize_cow_storage (storage=...)
+            #       at /home/kurtamohler/develop/pytorch-0/c10/util/UniqueVoidPtr.h:61
+            #
         ]
         for from_device, to_device in device_pairs:
             from_device_check = torch.empty(0, device=from_device).device