Fix ref leak in dtype.to_complex()/to_real() (#125154)

malfet · pytorchmergebot · commit 744f341aa4ea · 2024-04-29T23:59:27.000Z
By using `Py_NewRef` Also, wrap `THPDtype_to_real`/`THPDtype_to_complex` calls with `HANDLE_TH_ERRORS` Add regression test for the above issues, by calling to_complex for integral dtypes, that raises an exception and by preserving reference count to the same to_complex/to_real call to detect if leak is happeneing. Replace ```cpp auto dtype = (PyObject*)torch::getTHPDtype(current_dtype); Py_INCREF(dtype); return dtype; ``` with a more compact/streamlined equivalent ```cpp return Py_NewRef(torch::getTHPDtype(current_dtype)); ``` Fixes #124868 Pull Request resolved: #125154 Approved by: https://github.com/Skylion007, https://github.com/albanD
diff --git a/test/test_type_info.py b/test/test_type_info.py
@@ -6,6 +6,7 @@
 # sharding on sandcastle. This line silences flake warnings
 load_tests = load_tests
 
+import sys
 import torch
 import unittest
 
@@ -23,6 +24,8 @@ def test_invalid_input(self):
         for dtype in [torch.int64, torch.int32, torch.int16, torch.int8, torch.uint8, torch.bool]:
             with self.assertRaises(TypeError):
                 _ = torch.finfo(dtype)
+            with self.assertRaises(RuntimeError):
+                dtype.to_complex()
 
     @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
     def test_iinfo(self):
@@ -88,6 +91,26 @@ def test_finfo(self):
         self.assertEqual(xinfo.resolution, 1.0)
         self.assertEqual(xinfo.dtype, "float8_e4m3fn")
 
+    def test_to_complex(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/124868
+        # If reference count is leaked this would be a set of 10 elements
+        ref_cnt = {sys.getrefcount(torch.float32.to_complex()) for _ in range(10)}
+        self.assertLess(len(ref_cnt), 3)
+
+        self.assertEqual(torch.float64.to_complex(), torch.complex128)
+        self.assertEqual(torch.float32.to_complex(), torch.complex64)
+        self.assertEqual(torch.float16.to_complex(), torch.complex32)
+
+    def test_to_real(self):
+        # Regression test for https://github.com/pytorch/pytorch/issues/124868
+        # If reference count is leaked this would be a set of 10 elements
+        ref_cnt = {sys.getrefcount(torch.cfloat.to_real()) for _ in range(10)}
+        self.assertLess(len(ref_cnt), 3)
+
+        self.assertEqual(torch.complex128.to_real(), torch.double)
+        self.assertEqual(torch.complex64.to_real(), torch.float32)
+        self.assertEqual(torch.complex32.to_real(), torch.float16)
+
 if __name__ == '__main__':
     TestCase._default_dtype_check_enabled = True
     run_tests()
diff --git a/torch/csrc/Dtype.cpp b/torch/csrc/Dtype.cpp
@@ -7,6 +7,7 @@
 #include <torch/csrc/utils/object_ptr.h>
 #include <torch/csrc/utils/python_numbers.h>
 #include <torch/csrc/utils/python_strings.h>
+#include <torch/csrc/utils/pythoncapi_compat.h>
 #include <torch/csrc/utils/tensor_dtypes.h>
 #include <torch/csrc/utils/tensor_types.h>
 #include <cstring>
@@ -74,21 +75,25 @@ PyObject* THPDtype_reduce(PyObject* _self, PyObject* noargs) {
 }
 
 PyObject* THPDtype_to_real(PyObject* _self, PyObject* noargs) {
+  HANDLE_TH_ERRORS
   auto* self = (THPDtype*)_self;
   auto scalar_type = self->scalar_type;
   if (!at::isFloatingType(self->scalar_type)) {
     scalar_type = at::toRealValueType(self->scalar_type);
   }
-  return (PyObject*)torch::getTHPDtype(scalar_type);
+  return Py_NewRef(torch::getTHPDtype(scalar_type));
+  END_HANDLE_TH_ERRORS
 }
 
 PyObject* THPDtype_to_complex(PyObject* _self, PyObject* noargs) {
+  HANDLE_TH_ERRORS
   auto* self = (THPDtype*)_self;
   auto scalar_type = self->scalar_type;
   if (!at::isComplexType(self->scalar_type)) {
     scalar_type = at::toComplexType(self->scalar_type);
   }
-  return (PyObject*)torch::getTHPDtype(scalar_type);
+  return Py_NewRef(torch::getTHPDtype(scalar_type));
+  END_HANDLE_TH_ERRORS
 }
 
 typedef PyObject* (*getter)(PyObject*, void*);
diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp
@@ -1048,9 +1048,7 @@ PyObject* THPModule_setFlushDenormal(PyObject* _unused, PyObject* arg) {
 PyObject* THPModule_getDefaultDtype(PyObject* _unused, PyObject* arg) {
   HANDLE_TH_ERRORS
   auto scalar_type = torch::tensors::get_default_scalar_type();
-  auto dtype = (PyObject*)torch::getTHPDtype(scalar_type);
-  Py_INCREF(dtype);
-  return dtype;
+  return Py_NewRef(torch::getTHPDtype(scalar_type));
   END_HANDLE_TH_ERRORS
 }
 
diff --git a/torch/csrc/autograd/init.cpp b/torch/csrc/autograd/init.cpp
@@ -536,9 +536,7 @@ static PyObject* get_autocast_dtype(
   auto r = parser.parse(args, kwargs, parsed_args);
   auto device_type = at::Device(r.string(0)).type();
   at::ScalarType current_dtype = at::autocast::get_autocast_dtype(device_type);
-  auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
-  Py_INCREF(dtype);
-  return dtype;
+  return Py_NewRef(torch::getTHPDtype(current_dtype));
   END_HANDLE_TH_ERRORS
 }
 
@@ -735,9 +733,7 @@ static PyObject* get_autocast_gpu_dtype(PyObject* _unused, PyObject* arg) {
   TORCH_WARN_DEPRECATION(
       "torch.get_autocast_gpu_dtype() is deprecated. Please use torch.get_autocast_dtype('cuda') instead.")
   at::ScalarType current_dtype = at::autocast::get_autocast_dtype(at::kCUDA);
-  auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
-  Py_INCREF(dtype);
-  return dtype;
+  return Py_NewRef(torch::getTHPDtype(current_dtype));
   END_HANDLE_TH_ERRORS
 }
 
@@ -746,9 +742,7 @@ static PyObject* get_autocast_cpu_dtype(PyObject* _unused, PyObject* arg) {
   TORCH_WARN_DEPRECATION(
       "torch.get_autocast_cpu_dtype() is deprecated. Please use torch.get_autocast_dtype('cpu') instead.")
   at::ScalarType current_dtype = at::autocast::get_autocast_dtype(at::kCPU);
-  auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
-  Py_INCREF(dtype);
-  return dtype;
+  return Py_NewRef(torch::getTHPDtype(current_dtype));
   END_HANDLE_TH_ERRORS
 }
 
@@ -757,9 +751,7 @@ static PyObject* get_autocast_ipu_dtype(PyObject* _unused, PyObject* arg) {
   TORCH_WARN_DEPRECATION(
       "torch.get_autocast_ipu_dtype() is deprecated. Please use torch.get_autocast_dtype('ipu') instead.")
   at::ScalarType current_dtype = at::autocast::get_autocast_dtype(at::kIPU);
-  auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
-  Py_INCREF(dtype);
-  return dtype;
+  return Py_NewRef(torch::getTHPDtype(current_dtype));
   END_HANDLE_TH_ERRORS
 }
 
@@ -768,9 +760,7 @@ static PyObject* get_autocast_xla_dtype(PyObject* _unused, PyObject* arg) {
   TORCH_WARN_DEPRECATION(
       "torch.get_autocast_xla_dtype() is deprecated. Please use torch.get_autocast_dtype('xla') instead.")
   at::ScalarType current_dtype = at::autocast::get_autocast_dtype(at::kXLA);
-  auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
-  Py_INCREF(dtype);
-  return dtype;
+  return Py_NewRef(torch::getTHPDtype(current_dtype));
   END_HANDLE_TH_ERRORS
 }
 
