Update pin_memory implementation follow stock pytorch change

Signed-off-by: majing <[email protected]>
intel · Jul 23, 2024 · a6fb4b7 · a6fb4b7
1 parent 5f4970b
commit a6fb4b7
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 63 deletions.
diff --git a/src/ATen/native/xpu/PinnedMemoryAllocator.cpp b/src/ATen/native/xpu/PinnedMemoryAllocator.cpp
diff --git a/test/xpu/test_torch_xpu.py b/test/xpu/test_torch_xpu.py
@@ -8572,31 +8572,17 @@ def test_new(self) -> None:
         # TypeError would be better
         self.assertRaises(RuntimeError, lambda: x.new(z.storage()))
 
-    @unittest.skipIf(PYTORCH_CUDA_MEMCHECK, "is_pinned uses failure to detect pointer property")
     def test_pin_memory(self):
         x = torch.randn(3, 5)
         self.assertFalse(x.is_pinned())
-        if not torch.cuda.is_available() or not torch.xpu.is_available():
-            self.assertRaises(RuntimeError, lambda: x.pin_memory())
-        else:
-            if torch.xpu.is_available():
-                device = 'xpu'
-                self.assertFalse(x.is_pinned(device))
-                pinned = x.pin_memory(device)
-                self.assertTrue(pinned.is_pinned(device))
-                self.assertEqual(pinned, x)
-                self.assertNotEqual(pinned.data_ptr(), x.data_ptr())
-                # test that pin_memory on already pinned tensor has no effect
-                self.assertIs(pinned, pinned.pin_memory(device))
-                self.assertEqual(pinned.data_ptr(), pinned.pin_memory(device).data_ptr())
-            else:
-                pinned = x.pin_memory()
-                self.assertTrue(pinned.is_pinned())
-                self.assertEqual(pinned, x)
-                self.assertNotEqual(pinned.data_ptr(), x.data_ptr())
-                # test that pin_memory on already pinned tensor has no effect
-                self.assertIs(pinned, pinned.pin_memory())
-                self.assertEqual(pinned.data_ptr(), pinned.pin_memory().data_ptr())
+        if torch.cuda.is_available() or torch.xpu.is_available():
+            pinned = x.pin_memory()
+            self.assertTrue(pinned.is_pinned())
+            self.assertEqual(pinned, x)
+            self.assertNotEqual(pinned.data_ptr(), x.data_ptr())
+            # test that pin_memory on already pinned tensor has no effect
+            self.assertIs(pinned, pinned.pin_memory())
+            self.assertEqual(pinned.data_ptr(), pinned.pin_memory().data_ptr())
 
 
 

diff --git a/yaml/xpu_functions.yaml b/yaml/xpu_functions.yaml
@@ -539,8 +539,6 @@ supported:
   - sgn.out
   - sgn_
   - _cdist_forward
-  - _pin_memory
-  - is_pinned
   - is_set_to
   - bucketize.Tensor
   - bucketize.Tensor_out