[CINN] add no_need_buffer for recompute algorithm (#70255)

chen2016013 · web-flow · commit 0b05d345d836 · 2024-12-19T15:28:16.000+08:00
* add no_need_buffer for recompute algorithm

* rerun CI

* fix test_tensor_attr_consistency

* fix bug

* bug fix
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1203,7 +1203,27 @@ void BindOperation(py::module *m) {
             self.set_attribute(
                 "chunk_id",
                 Int32Attribute::get(pir::IrContext::Instance(), chunk_id));
-          });
+          })
+      .def("is_no_need_buffer",
+           [](Operation &self, const Value &operand_source) -> bool {
+             paddle::dialect::OpYamlInfoInterface op_info_interface =
+                 self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
+             std::unique_ptr<paddle::dialect::OpYamlInfoParser> info_parser(
+                 nullptr);
+             if (op_info_interface) {
+               info_parser =
+                   std::make_unique<paddle::dialect::OpYamlInfoParser>(
+                       op_info_interface.GetOpInfo(),
+                       paddle::dialect::IsLegacyOp(self.name()));
+               auto &no_need_buffer_ids = info_parser->NoNeedBufferIds();
+               for (auto no_need_buffer_id : no_need_buffer_ids) {
+                 if (operand_source == self.operand_source(no_need_buffer_id)) {
+                   return true;
+                 }
+               }
+             }
+             return false;
+           });
   py::class_<Operation::BlockContainer> block_container(
       *m, "Operation_BlockContainer", R"DOC(
     The Operation_BlockContainer only use to walk all blocks in the operation.
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -239,9 +239,18 @@ def update(self, other: set):
     def pop(self):
         return self._set.pop()._value
 
+    def remove(self, val):
+        self._set.remove(ValueWrapper(val))
+
+    def discard(self, val):
+        self._set.discard(ValueWrapper(val))
+
     def __and__(self, other: ValueSet):
         return ValueSet(self._set & other._set)
 
+    def __sub__(self, other: ValueSet):
+        return ValueSet(self._set - other._set)
+
     def __or__(self, other: ValueSet):
         return ValueSet(self._set | other._set)
 
diff --git a/python/paddle/decomposition/recompute.py b/python/paddle/decomposition/recompute.py
@@ -157,6 +157,8 @@
 # Restricts the amount of computation recompute can do.
 MAX_DIST_FROM_BW = 3
 
+MINIMUM_WEIGHT = 0.1
+
 
 def DebugPrint(*args):
     flag = os.getenv("FLAGS_print_auto_recompute_debug")
@@ -458,7 +460,24 @@ def _is_materialized(value_node, placeholder_value_nodes):
         users = find_value_node_users(value_node)
         return not all(_is_fusible(value_node, user) for user in users)
 
-    def _get_node_weight(value_node, placeholder_value_nodes):
+    def _get_no_need_buffer_values_from_program(program):
+        need_buffer_values = backward_utils.ValueSet()
+        all_values = backward_utils.ValueSet()
+        for op in program.global_block().ops:
+            for op_operand_source in op.operands_source():
+                all_values.add(op_operand_source)
+                if op.is_no_need_buffer(op_operand_source):
+                    continue
+                need_buffer_values.add(op_operand_source)
+        no_need_buffer_values = all_values - need_buffer_values
+        return no_need_buffer_values
+
+    def _get_node_weight(
+        value_node, no_need_buffer_values, placeholder_value_nodes
+    ):
+        if value_node in no_need_buffer_values:
+            return MINIMUM_WEIGHT
+
         mem_sz = cal_value_node_size(value_node)
 
         if (
@@ -512,6 +531,7 @@ def _ban_recomputation(value_node):
 
     judge_fusion_loop = JudgeFusionLoop(program, unrecomputable_ops)
     forward_ops = set(program.global_block().ops[: fwd_op_end_idx + 1])
+    no_need_buffer_values = _get_no_need_buffer_values_from_program(program)
 
     for value_node in (
         required_fw_value_nodes
@@ -570,7 +590,9 @@ def _ban_recomputation(value_node):
             value_id_dict[value_node.id] = value_node
 
         weight = _get_node_weight(
-            value_node, placeholder_value_nodes=inputs | outputs
+            value_node,
+            no_need_buffer_values,
+            placeholder_value_nodes=inputs | outputs,
         )
 
         # Creates the weights on the "node" edge
diff --git a/test/dygraph_to_static/test_tensor_attr_consistency.py b/test/dygraph_to_static/test_tensor_attr_consistency.py
@@ -120,6 +120,7 @@
         'trainable',
         'is_sparse_coo_tensor_type',
         'is_sparse_csr_tensor_type',
+        'is_no_need_buffer',
     ]
 )
 

Original file line number	Diff line number	Diff line change
`@@ -120,6 +120,7 @@`
`120`	`120`	`'trainable',`
`121`	`121`	`'is_sparse_coo_tensor_type',`
`122`	`122`	`'is_sparse_csr_tensor_type',`
	`123`	`+ 'is_no_need_buffer',`
`123`	`124`	`]`
`124`	`125`	`)`
`125`	`126`