[CINN] refine recompute for data，parameter (#69954)

wanghuancoder · web-flow · commit 18e8791f6af4 · 2024-12-05T18:38:39.000+08:00
* fix some cinn bug
diff --git a/python/paddle/decomposition/recompute.py b/python/paddle/decomposition/recompute.py
@@ -736,12 +736,17 @@ def getIdx(program, op):
             define_op = recompute_value.get_defining_op()
             if define_op in marked_recompute_ops or define_op is None:
                 return
+            if define_op.name() in [
+                "builtin.parameter",
+                "pd_op.data",
+            ]:
+                if recompute_value not in needed_saved_values:
+                    needed_saved_values.add(recompute_value)
+                return
             op_inputs = define_op.operands_source()
             if len(op_inputs) == 0 and define_op.name() not in [
                 "pd_op.full",
                 "pd_op.full_int_array",
-                "builtin.parameter",
-                "pd_op.data",
             ]:
 
                 def getIdx(program, op):
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -774,18 +774,17 @@ def get_kwargs_forward_matched_value(kw_name, kw_value):
                         elif kw_name in forward_name_value_map:
                             return forward_name_value_map[kw_name]
                         else:
-                            return None
+                            raise Exception(f"kw_args: {kw_name} not found")
 
                     for [kw_name, kw_value] in (
                         backward_program.global_block().kwargs().items()
                     ):
                         forward_matched_value = (
                             get_kwargs_forward_matched_value(kw_name, kw_value)
                         )
-                        if forward_matched_value is not None:
-                            share_symbol_shape_from_forward_to_backward(
-                                forward_matched_value, kw_value
-                            )
+                        share_symbol_shape_from_forward_to_backward(
+                            forward_matched_value, kw_value
+                        )
 
                 if cse_is_enabled():
                     paddle.base.libpaddle.pir.apply_cse_pass(forward_program)
