[PIR+CINN]refactor cinn pir compiler temp buffer

paddle/cinn/hlir/framework/op_lowering_impl.cc

@@ -44,6 +44,45 @@ using cinn::common::Type;
 
 using cinn::hlir::op::ExternalApiRegistry;
 
+// Collect the temporary tensors from a computational graph.
+std::vector<ir::Buffer> GetTempBuffers(
+    const std::vector<cinn::ir::Tensor>& tensor_args, Expr body) {
+  std::unordered_set<std::string> tensor_arg_names;
+  std::unordered_set<std::string> buffer_arg_names;
+  for (auto& tensor : tensor_args) {
+    tensor_arg_names.insert(tensor->name);
+    if (tensor->buffer.defined()) {
+      buffer_arg_names.insert(tensor->buffer->name);
+    }
+  }
+  std::map<std::string, ir::Buffer>
+      name_to_buffer;  // used to avoid duplication.
+
+  auto all_temp_tensors =
+      ir::ir_utils::CollectIRNodesWithoutTensor(body, [&](const Expr* x) {
+        return x->as_tensor() && x->as_tensor()->buffer.defined() &&
+               ((!buffer_arg_names.count(x->as_tensor()->buffer->name) &&
+                 !tensor_arg_names.count(x->as_tensor()->name)) ||
+                utils::Endswith(x->as_tensor()->buffer->name, "temp_buffer"));
+      });
+  for (auto& e : all_temp_tensors) {
+    auto buffer_name = e.as_tensor()->buffer->name;
+    if (!name_to_buffer.count(buffer_name)) {
+      name_to_buffer[buffer_name] = e.as_tensor()->buffer;
+    } else {
+      // TODO(phlrain): why update
+      if (e.as_tensor()->buffer->numel() <
+          name_to_buffer[buffer_name]->numel()) {
+        name_to_buffer[buffer_name] = e.as_tensor()->buffer;
+      }
+    }
+  }
+
+  std::vector<ir::Buffer> temp_buffers;
+  for (auto& i : name_to_buffer) temp_buffers.push_back(i.second);
+  return temp_buffers;
+}
+
 OpLowererImpl::OpLowererImpl(
     const absl::flat_hash_map<std::string, Type>& type_dict,
     const absl::flat_hash_map<std::string, shape_t>& shape_dict,
@@ -300,9 +339,7 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
   }
 #endif
   // 2.Prepare temp buffers
-  poly::StageMap stages;
-  auto temp_buffers =
-      lang::GetTempBuffers(*group_func_arg_tensors, stages, func_body);
+  auto temp_buffers = GetTempBuffers(*group_func_arg_tensors, func_body);
   // 3.Building LoweredFunc
   auto func = ir::_LoweredFunc_::Make(group->GetFuncName(),
                                       group_func_args,