[Dy2St] Constant promotion dtype align to other branch (#69097)

SigureMo · web-flow · commit 4d0da395533d · 2024-11-02T10:10:55.000+08:00
diff --git a/python/paddle/jit/dy2static/convert_operators.py b/python/paddle/jit/dy2static/convert_operators.py
@@ -41,16 +41,19 @@
 __all__ = []
 
 
-def to_static_variable(x):
+def to_static_variable(x, dtype=None):
     '''
     Translate a Python Tensor to PaddlePaddle static graph Tensor
     '''
     if isinstance(x, bool):
-        return paddle.full(shape=[], dtype='bool', fill_value=x)
+        dtype = 'bool' if dtype is None else dtype
+        return paddle.full(shape=[], dtype=dtype, fill_value=x)
     if isinstance(x, float):
-        return paddle.full(shape=[], dtype='float64', fill_value=x)
+        dtype = 'float64' if dtype is None else dtype
+        return paddle.full(shape=[], dtype=dtype, fill_value=x)
     if isinstance(x, int):
-        return paddle.full(shape=[], dtype='int64', fill_value=x)
+        dtype = 'int64' if dtype is None else dtype
+        return paddle.full(shape=[], dtype=dtype, fill_value=x)
     if not use_pir_api() and (isinstance(x, UndefinedVar) or x is None):
         """
         for early return case, we need a variable to represent None, current we use data_layer_not_check.
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
@@ -1302,6 +1302,14 @@ def create_container_by_items_and_indices(*items_indices_pairs):
     return container
 
 
+def run_with_block(fn, block):
+    def new_fn(*args, **kwargs):
+        with block():
+            return fn(*args, **kwargs)
+
+    return new_fn
+
+
 class OutputSelector:
     def __init__(
         self, if_op, flattened_true_output, flattened_false_output, names
@@ -1331,6 +1339,13 @@ def unified_output(self):
                 ],
                 name,
             )
+            (true_out, false_out) = OutputSelector.precision_promotion(
+                [
+                    (true_out, self.if_op.true_block),
+                    (false_out, self.if_op.false_block),
+                ],
+                name,
+            )
             unified_true_output.append(true_out)
             unified_false_output.append(false_out)
         return unified_true_output, unified_false_output
@@ -1407,58 +1422,11 @@ def all_has_same_type(outs):
                 return True
             return all(type(out) is type(outs[0]) for out in outs[1:])
 
-        def all_has_same_dtype(outs):
-            if len(outs) <= 1:
-                return True
-            return all(out.dtype == outs[0].dtype for out in outs[1:])
-
-        def constant_to_variable_with_block(constant, block_context_manager):
-            with block_context_manager():
-                return to_static_variable(constant)
-
-        def promote_precision(out_with_blocks):
-            def get_expected_precision(out_with_blocks):
-                if len(out_with_blocks) <= 1:
-                    return core.DataType.FLOAT32
-                # now only support FLOAT16 to FLOAT32
-                if any(
-                    out.dtype == core.DataType.FLOAT16
-                    for out, _ in out_with_blocks
-                ) and any(
-                    out.dtype == core.DataType.FLOAT32
-                    for out, _ in out_with_blocks
-                ):
-                    return core.DataType.FLOAT32
-                else:
-                    return out_with_blocks[0][0].dtype
-
-            new_outs = []
-            expected_dtype = get_expected_precision(out_with_blocks)
-            for out, block in out_with_blocks:
-                if expected_dtype != out.dtype:
-                    with block():
-                        out = paddle.cast(
-                            out, _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE[expected_dtype]
-                        )
-                new_outs.append(out)
-            return new_outs
-
-        if all(isinstance(out, paddle.pir.Value) for out in outs):
-            amp_attrs = core._get_amp_attrs()
-            amp_level = amp_attrs._amp_level
-            apply_amp_level_list = [
-                core.AmpLevel.O1,
-                core.AmpLevel.O2,
-            ]
-            if (amp_level in apply_amp_level_list) and (
-                not all_has_same_dtype(outs)
-            ):
-                warnings.warn(
-                    f"Return results from different branches in cond has different dtype: true value dtype is '{outs[0].dtype}' and false value dtype is '{outs[1].dtype}', "
-                    "so we will promote the lower precision to the higher one."
-                )
-                return promote_precision(out_with_blocks)
-            return outs
+        def get_first_value_dtype(outs):
+            for out in outs:
+                if isinstance(out, paddle.pir.Value):
+                    return out.dtype
+            return None
 
         if all(arg is None for arg in outs):
             return outs
@@ -1475,7 +1443,9 @@ def get_expected_precision(out_with_blocks):
                     "so we will promote the constant to variable."
                 )
                 return [
-                    constant_to_variable_with_block(out, block)
+                    # TODO(SigureMo): Should we use the same dtype for all the constants?
+                    # e.g. in true branch var is 3, else branch var is 2, then the dtype should be float64.
+                    run_with_block(to_static_variable, block)(out, dtype=None)
                     for out, block in out_with_blocks
                 ]
 
@@ -1489,7 +1459,9 @@ def get_expected_precision(out_with_blocks):
                 + f"'{type(outs[0])}'"
             )
             return [
-                constant_to_variable_with_block(out, block)
+                run_with_block(to_static_variable, block)(
+                    out, dtype=get_first_value_dtype(outs)
+                )
                 for out, block in out_with_blocks
             ]
 
@@ -1505,6 +1477,61 @@ def get_expected_precision(out_with_blocks):
             f"returned `{name}` by false_fn is `{outs[0]}` and true_var of true_fn is `{outs[1]}`"
         )
 
+    @staticmethod
+    def precision_promotion(out_with_blocks, name):
+        # Only support promotion from fp16 to fp32 in AMP mode
+        outs, _ = zip(*out_with_blocks)
+
+        amp_attrs = core._get_amp_attrs()
+        amp_level = amp_attrs._amp_level
+        apply_amp_level_list = [
+            core.AmpLevel.O1,
+            core.AmpLevel.O2,
+        ]
+        if amp_level not in apply_amp_level_list:
+            return outs
+
+        def all_has_same_dtype(outs):
+            if len(outs) <= 1:
+                return True
+            return all(out.dtype == outs[0].dtype for out in outs[1:])
+
+        def promote_precision(out_with_blocks):
+            def get_expected_precision(out_with_blocks):
+                if len(outs) <= 1:
+                    return outs[0].dtype
+                # now only support fp16 to fp32
+                if any(
+                    out.dtype == paddle.float16 for out, _ in out_with_blocks
+                ) and any(
+                    out.dtype == paddle.float32 for out, _ in out_with_blocks
+                ):
+                    return paddle.float32
+                else:
+                    return out_with_blocks[0][0].dtype
+
+            new_outs = []
+            expected_dtype = get_expected_precision(out_with_blocks)
+            for out, block in out_with_blocks:
+                if expected_dtype != out.dtype:
+                    out = run_with_block(paddle.cast, block)(
+                        out, _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE[expected_dtype]
+                    )
+                new_outs.append(out)
+            return new_outs
+
+        if all(
+            isinstance(out, paddle.pir.Value) for out in outs
+        ) and not all_has_same_dtype(outs):
+            warnings.warn(
+                f"Return results from different branches in cond has different dtype: true value dtype is '{outs[0].dtype}' and false value dtype is '{outs[1].dtype}', "
+                "so we will promote the lower precision to the higher one."
+            )
+            outs = promote_precision(out_with_blocks)
+            return outs
+
+        return outs
+
 
 def cond(pred, true_fn=None, false_fn=None, name=None, return_names=None):
     """
diff --git a/test/dygraph_to_static/test_ifelse.py b/test/dygraph_to_static/test_ifelse.py
@@ -613,5 +613,30 @@ def test_use_undefined_var(self):
             static_fn(falsy)
 
 
+def dynamic_shape_with_constant_promotion(x):
+    x_shape0 = x.shape[0]
+    if x_shape0 < 10:
+        x_shape0 = x.shape[-1]
+    return x_shape0
+
+
+class TestDynamicShapeWithConstantPromotion(Dy2StTestBase):
+    @test_ast_only
+    @test_pir_only
+    def test_dynamic_shape_with_constant_promotion(self):
+        x = paddle.randn([5, 3])
+        static_fn = paddle.jit.to_static(
+            dynamic_shape_with_constant_promotion,
+            input_spec=[
+                paddle.static.InputSpec(
+                    shape=[None, 3],
+                    dtype='float32',
+                )
+            ],
+        )
+        out = static_fn(x)
+        self.assertEqual(out, 3)
+
+
 if __name__ == '__main__':
     unittest.main()
