[Fix] Fix astype copy when target dtype is same as source Tensor (#69333)

HydrogenSulfate · web-flow · commit 8b6b703c892b · 2024-11-15T11:06:12.000+08:00
* fix astype copy when target dtype is same as source Tensor

* fix annotations and refine unitest

* update

* update related code

* update unitest

* fix test_to_static_deprecated.py
diff --git a/python/paddle/base/dygraph/math_op_patch.py b/python/paddle/base/dygraph/math_op_patch.py
@@ -74,7 +74,8 @@ def monkey_patch_math_tensor():
     def astype(self: Tensor, dtype: DTypeLike) -> Tensor:
         """
 
-        Cast a Tensor to a specified data type.
+        Cast a Tensor to a specified data type if it differs from the current dtype;
+        otherwise, return the original Tensor.
 
         Args:
             dtype: The target data type.
@@ -97,6 +98,10 @@ def astype(self: Tensor, dtype: DTypeLike) -> Tensor:
         """
         if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
             dtype = convert_np_dtype_to_dtype_(dtype)
+
+        if self.dtype == dtype:
+            return self
+
         return _C_ops.cast(self, dtype)
 
     def _scalar_elementwise_op_(
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -1352,7 +1352,9 @@ def grad_var_name(var_name):
     return var_name + GRAD_VAR_SUFFIX
 
 
-def convert_np_dtype_to_proto_type(np_dtype: np.dtype | str):
+def convert_np_dtype_to_proto_type(
+    np_dtype: np.dtype | str,
+) -> core.VarDesc.VarType:
     """
     Convert the data type in numpy to the data type in Paddle.
 
@@ -1409,7 +1411,9 @@ def convert_np_dtype_to_proto_type(np_dtype: np.dtype | str):
         raise ValueError(f"Not supported numpy dtype {dtype}")
 
 
-def convert_np_dtype_to_dtype_(np_dtype):
+def convert_np_dtype_to_dtype_(
+    np_dtype: np.dtype | str,
+) -> core.VarDesc.VarType | core.DataType:
     """
     Convert the data type in numpy to the data type in Paddle.
 
diff --git a/python/paddle/base/layers/math_op_patch.py b/python/paddle/base/layers/math_op_patch.py
@@ -320,7 +320,8 @@ def astype(self, dtype):
         **Notes**:
             **The variable must be a** :ref:`api_paddle_Tensor`
 
-        Cast a variable to a specified data type.
+        Cast a variable to a specified data type if it differs from the current dtype;
+        otherwise, return the original variable.
 
         Args:
 
@@ -366,6 +367,9 @@ def astype(self, dtype):
                 original var's dtype is: paddle.float32, numpy dtype is float32
                 new var's dtype is: paddle.int64, numpy dtype is int64
         """
+        if self.dtype == dtype:
+            return self
+
         block = current_block(self)
         out = create_new_tmp_var(block, dtype)
         block.append_op(
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
 import numpy as np
 
@@ -99,7 +100,7 @@
 }
 
 
-def convert_np_dtype_to_dtype_(np_dtype):
+def convert_np_dtype_to_dtype_(np_dtype) -> DataType:
     """
     Convert the data type in numpy to the data type in Paddle.
 
@@ -108,7 +109,7 @@ def convert_np_dtype_to_dtype_(np_dtype):
             string.
 
     Returns:
-        core.DataType : The data type in Paddle.
+        DataType : The data type in Paddle.
 
     """
     # Convert the data type string to numpy data type.
@@ -123,7 +124,7 @@ def convert_np_dtype_to_dtype_(np_dtype):
     else:
         dtype = np.dtype(np_dtype)
 
-    if dtype in np_type_to_paddle_type.keys():
+    if dtype in np_type_to_paddle_type:
         return np_type_to_paddle_type[dtype]
     else:
         raise ValueError(f"Not supported numpy dtype {dtype}")
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -285,7 +285,8 @@ def astype(self, dtype):
         """
         **Notes**:
 
-        Cast a Value to a specified data type.
+        Convert a value to a specified data type if it differs from the current dtype;
+        otherwise, return the original value.
 
         Args:
 
@@ -316,6 +317,10 @@ def astype(self, dtype):
 
         if not isinstance(dtype, DataType):
             dtype = paddle.pir.core.convert_np_dtype_to_dtype_(dtype)
+
+        if self.dtype == dtype:
+            return self
+
         return _C_ops.cast(self, dtype)
 
     def _scalar_add_(var, value):
diff --git a/test/amp/test_amp_api.py b/test/amp/test_amp_api.py
@@ -472,7 +472,6 @@ def test_op_called_as_expected(self):
 
     def test_pir_op_called_as_expected(self):
         expected_fp16_calls = {
-            "pd_op.cast_": 1,
             "pd_op.layer_norm": 1,
             "pd_op.scale": 1,
             "pd_op.scale_": 2,
diff --git a/test/deprecated/auto_parallel/test_to_static_deprecated.py b/test/deprecated/auto_parallel/test_to_static_deprecated.py
@@ -112,7 +112,7 @@ def test_apply_optimizer(self):
         program_helper.to('train')
 
         forward_ops = program_helper.main_program.block(0).ops
-        self.assertEqual(len(forward_ops), 18)
+        self.assertEqual(len(forward_ops), 17)
 
         # step 2: apply optimizer to generate whole program
         optimize_ops, _ = program_helper.apply_optimizer(optimizer)
@@ -122,7 +122,7 @@ def test_apply_optimizer(self):
             for op in program_helper.main_program.block(0).ops
             if op.type == 'sgd'
         ]
-        self.assertEqual(len(all_ops), 38)
+        self.assertEqual(len(all_ops), 37)
         self.assertEqual(len(optimize_ops), len(sgd_ops))
 
         program_helper.reset()
diff --git a/test/legacy_test/test_adamw_op.py b/test/legacy_test/test_adamw_op.py
@@ -463,11 +463,15 @@ def _test_adamw_op_dygraph_place_amp_with_maingrad(
         beta2_pow_acc = paddle.ones([1]).astype(paddle.float32)
         beta2_pow_acc[0] = _beta2**10
 
-        ref_param = param.astype(paddle.float32)
-        ref_beta1_pow_acc = beta1_pow_acc.astype(paddle.float32)
-        ref_beta2_pow_acc = beta2_pow_acc.astype(paddle.float32)
-        ref_moment_1 = moment1.astype(paddle.float32)
-        ref_moment_2 = moment2.astype(paddle.float32)
+        ref_param = param.astype(paddle.float32).clone().detach()
+        ref_beta1_pow_acc = (
+            beta1_pow_acc.astype(paddle.float32).clone().detach()
+        )
+        ref_beta2_pow_acc = (
+            beta2_pow_acc.astype(paddle.float32).clone().detach()
+        )
+        ref_moment_1 = moment1.astype(paddle.float32).clone().detach()
+        ref_moment_2 = moment2.astype(paddle.float32).clone().detach()
 
         # reference code
         _, _, _, _, _, _ = paddle._C_ops.adamw_(
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -285,6 +285,10 @@ def test_astype(self):
         if not paddle.framework.use_pir_api():
             a.desc.set_need_check_feed(False)
         b = a.astype('float64')
+
+        c = a.astype(a.dtype)
+        self.assertTrue(c.is_same(a))
+
         place = base.CPUPlace()
         exe = base.Executor(place)
         a_np = np.random.uniform(-1, 1, size=[10, 1]).astype('float64')
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -501,6 +501,12 @@ def test_T(self):
                     (output_x,) = exe.run(main_program, fetch_list=[x_T])
                     self.assertEqual(output_x.shape, tuple(out_shape))
 
+    def test_astype_zero_copy_if_same_dtype(self):
+        with paddle.pir_utils.IrGuard():
+            x = paddle.static.data(name='x', shape=[3, 2, 1])
+            y = x.astype(x.dtype)
+            self.assertTrue(x.is_same(y))
+
     def test_mT(self):
         with paddle.pir_utils.IrGuard():
             shape = [1]
diff --git a/test/legacy_test/test_math_op_patch_var_base.py b/test/legacy_test/test_math_op_patch_var_base.py
@@ -596,9 +596,14 @@ def test_astype(self):
             res1 = a.astype(np.float16)
             res2 = a.astype('float16')
             res3 = a.astype(paddle.float16)
+            res4 = a.astype(a.dtype)
 
             self.assertEqual(res1.dtype, res2.dtype)
             self.assertEqual(res1.dtype, res3.dtype)
+            self.assertEqual(res4.dtype, a.dtype)
+            self.assertEqual(
+                a.data_ptr(), res4.data_ptr()
+            )  # zero-copy if same dtype
 
             np.testing.assert_array_equal(res1.numpy(), res2.numpy())
             np.testing.assert_array_equal(res1.numpy(), res3.numpy())
diff --git a/test/xpu/test_adamw_op_xpu.py b/test/xpu/test_adamw_op_xpu.py
@@ -722,11 +722,15 @@ def _test_adamw_op_dygraph_place_amp_with_maingrad(
         beta2_pow_acc = paddle.ones([1]).astype(paddle.float32)
         beta2_pow_acc[0] = _beta2**10
 
-        ref_param = param.astype(paddle.float32)
-        ref_beta1_pow_acc = beta1_pow_acc.astype(paddle.float32)
-        ref_beta2_pow_acc = beta2_pow_acc.astype(paddle.float32)
-        ref_moment_1 = moment1.astype(paddle.float32)
-        ref_moment_2 = moment2.astype(paddle.float32)
+        ref_param = param.astype(paddle.float32).clone().detach()
+        ref_beta1_pow_acc = (
+            beta1_pow_acc.astype(paddle.float32).clone().detach()
+        )
+        ref_beta2_pow_acc = (
+            beta2_pow_acc.astype(paddle.float32).clone().detach()
+        )
+        ref_moment_1 = moment1.astype(paddle.float32).clone().detach()
+        ref_moment_2 = moment2.astype(paddle.float32).clone().detach()
 
         # reference code
         _, _, _, _, _, _ = paddle._C_ops.adamw_(
