[PIR]Add Dtype patch and fix test_tensor_scalar_type_promotion_dynamic in pir mode (#64187)

YuanRisheng · web-flow · commit 93e62b9ef5fc · 2024-05-13T16:41:18.000+08:00
* fix test_tensor_scalar_type_promotion_dynamic in pir mode

* fix year

* fix py3
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -34,12 +34,13 @@
 # We need remove the duplicated code here once we fix
 # the illogical implement in the monkey-patch methods later.
 from .framework import monkey_patch_math_tensor, monkey_patch_variable
-from .pir import monkey_patch_program, monkey_patch_value
+from .pir import monkey_patch_dtype, monkey_patch_program, monkey_patch_value
 
 monkey_patch_variable()
 monkey_patch_math_tensor()
 monkey_patch_value()
 monkey_patch_program()
+monkey_patch_dtype()
 
 from .base.dataset import *  # noqa: F403
 from .framework import (
diff --git a/python/paddle/base/data_feeder.py b/python/paddle/base/data_feeder.py
@@ -19,7 +19,7 @@
 from paddle import pir
 
 from ..pir import Value
-from ..pir.core import ParameterMeta
+from ..pir.core import _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE, ParameterMeta
 from . import core
 from .framework import (
     Variable,
@@ -62,21 +62,6 @@
     'complex128': core.VarDesc.VarType.COMPLEX128,
 }
 
-_PADDLE_PIR_DTYPE_2_NUMPY_DTYPE = {
-    core.DataType.BOOL: 'bool',
-    core.DataType.FLOAT16: 'float16',
-    core.DataType.BFLOAT16: 'uint16',
-    core.DataType.FLOAT32: 'float32',
-    core.DataType.FLOAT64: 'float64',
-    core.DataType.INT8: 'int8',
-    core.DataType.INT16: 'int16',
-    core.DataType.INT32: 'int32',
-    core.DataType.INT64: 'int64',
-    core.DataType.UINT8: 'uint8',
-    core.DataType.COMPLEX64: 'complex64',
-    core.DataType.COMPLEX128: 'complex128',
-}
-
 
 def convert_float_to_uint16(data, data_format="NCHW"):
     if data.size == 0:
diff --git a/python/paddle/pir/__init__.py b/python/paddle/pir/__init__.py
@@ -41,6 +41,7 @@
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
 
 from . import core  # noqa: F401
+from .dtype_patch import monkey_patch_dtype  # noqa: F401
 from .math_op_patch import monkey_patch_value  # noqa: F401
 from .program_patch import monkey_patch_program  # noqa: F401
 
diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -74,6 +74,21 @@
     np.complex128: DataType.COMPLEX128,
 }
 
+_PADDLE_PIR_DTYPE_2_NUMPY_DTYPE = {
+    DataType.BOOL: 'bool',
+    DataType.FLOAT16: 'float16',
+    DataType.BFLOAT16: 'uint16',
+    DataType.FLOAT32: 'float32',
+    DataType.FLOAT64: 'float64',
+    DataType.INT8: 'int8',
+    DataType.INT16: 'int16',
+    DataType.INT32: 'int32',
+    DataType.INT64: 'int64',
+    DataType.UINT8: 'uint8',
+    DataType.COMPLEX64: 'complex64',
+    DataType.COMPLEX128: 'complex128',
+}
+
 
 def convert_np_dtype_to_dtype_(np_dtype):
     """
diff --git a/python/paddle/pir/dtype_patch.py b/python/paddle/pir/dtype_patch.py
@@ -0,0 +1,40 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle.base.libpaddle import DataType
+from paddle.pir.core import _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE
+
+_already_patch_dtype_repr = False
+
+
+def monkey_patch_dtype():
+    global _already_patch_dtype_repr
+    if not _already_patch_dtype_repr:
+        # NOTE(zhiqiu): pybind11 will set a default __str__ method of enum class.
+        # So, we need to overwrite it to a more readable one.
+        # See details in https://github.com/pybind/pybind11/issues/2537.
+        origin = DataType.__str__
+
+        def dtype_str(dtype):
+            if dtype in _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE:
+                numpy_dtype = _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE[dtype]
+                if numpy_dtype == 'uint16':
+                    numpy_dtype = 'bfloat16'
+                prefix = 'paddle.'
+                return prefix + numpy_dtype
+            else:
+                return origin(dtype)
+
+        DataType.__str__ = dtype_str
+        _already_patch_dtype_repr = True
diff --git a/test/deprecated/ir/pir/test_special_op_translator.py b/test/deprecated/ir/pir/test_special_op_translator.py
@@ -534,7 +534,7 @@ def test_data_op(self):
         self.assertTrue(l.global_block().ops[0].name() == "pd_op.data")
         data_op = l.global_block().ops[0]
         self.assertIn("dtype", data_op.attrs())
-        self.assertEqual(str(data_op.attrs()["dtype"]), "DataType.INT64")
+        self.assertEqual(str(data_op.attrs()["dtype"]), "paddle.int64")
 
 
 class TestCheckUnregisteredOp(unittest.TestCase):
