Call _C_ops.sum in new ir (#57016)

0x45f · web-flow · commit 878e1821ae51 · 2023-09-11T14:14:58.000+08:00
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -163,6 +163,7 @@ def __setattr__(self, name, val):
     DataType.BOOL: core.VarDesc.VarType.BOOL,
     DataType.FLOAT16: core.VarDesc.VarType.FP16,
     DataType.UINT16: core.VarDesc.VarType.BF16,
+    DataType.BFLOAT16: core.VarDesc.VarType.BF16,
     DataType.FLOAT32: core.VarDesc.VarType.FP32,
     DataType.FLOAT64: core.VarDesc.VarType.FP64,
     DataType.INT8: core.VarDesc.VarType.INT8,
@@ -1083,6 +1084,9 @@ def convert_np_dtype_to_dtype_(np_dtype):
         core.VarDesc.VarType / core.DataType : The data type in Paddle.
 
     """
+    if in_new_ir_mode():
+        return ir.core.convert_np_dtype_to_dtype_(np_dtype)
+
     # Convert the data type string to numpy data type.
     if isinstance(np_dtype, str) and np_dtype == "bfloat16":
         dtype = np.uint16
diff --git a/python/paddle/ir/core.py b/python/paddle/ir/core.py
@@ -28,7 +28,7 @@
     np.dtype("int16"): DataType.INT16,
     np.dtype("int64"): DataType.INT64,
     np.dtype("bool_"): DataType.BOOL,
-    np.dtype("uint16"): DataType.UINT16,
+    np.dtype("uint16"): DataType.BFLOAT16,
     np.dtype("uint8"): DataType.UINT8,
     np.dtype("int8"): DataType.INT8,
     np.dtype("complex64"): DataType.COMPLEX64,
@@ -50,7 +50,9 @@ def convert_np_dtype_to_dtype_(np_dtype):
     """
     # Convert the data type string to numpy data type.
     if isinstance(np_dtype, str) and np_dtype == "bfloat16":
-        dtype = np.uint16
+        # since there is still no support for bfloat16 in NumPy,
+        # uint16 is used for casting bfloat16
+        dtype = np.dtype("uint16")
     else:
         dtype = np.dtype(np_dtype)
 
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -1517,17 +1517,12 @@ def sum(x, axis=None, dtype=None, keepdim=False, name=None):
 
     dtype_flag = False
     if dtype is not None:
-        if paddle.ir.core._use_new_ir_api():
-            dtype = paddle.ir.core.convert_np_dtype_to_dtype_(dtype)
-        else:
-            dtype_flag = True
-            dtype = convert_np_dtype_to_dtype_(dtype)
+        dtype_flag = True
+        dtype = convert_np_dtype_to_dtype_(dtype)
 
-    if in_dynamic_mode():
+    if in_dynamic_or_new_ir_mode():
         return _C_ops.sum(x, axis, dtype, keepdim)
     else:
-        if paddle.ir.core._use_new_ir_api():
-            return paddle._ir_ops.sum(x, axis, dtype, keepdim)
         reduce_all, axis = _get_reduce_axis_with_tensor(axis, x)
         attrs = {'dim': axis, 'keep_dim': keepdim, 'reduce_all': reduce_all}
 
