fix the issue with while input values and args

Author: 0x3878f

paddle2onnx/mapper/exporter.cc

@@ -252,6 +252,7 @@ int32_t ModelExporter::GetMinOpsetVersion(const PaddlePirParser& pir_parser,
       current_opset = current_opset > 11 ? current_opset : 11;
     } else if (op_name == "pd_op.while") {
       auto while_op = op->dyn_cast<paddle::dialect::WhileOp>();
+      pir_parser.GetWhileInputValuesAndArgsMappings(while_op);
       current_opset = GetCfBlockMinOpsetVersion(pir_parser, while_op.body());
       current_opset = current_opset > 11 ? current_opset : 11;
 
@@ -483,7 +484,7 @@ ONNX_NAMESPACE::GraphProto ModelExporter::ExportIfBlock(
       temp_outputs.push_back(std::move(MakeValueInfo(cond_info[0])));
       if (value.defining_op() == nullptr) {
         value =
-            pir::Value(pir_parser.while_op_input_value_map[&(*(value.impl()))]);
+            pir::Value(pir_parser.while_op_values_args_map[&(*(value.impl()))]);
       }
       if (value.defining_op()->GetParent() != &block) {
         temp_inputs.push_back(std::move(MakeValueInfo(cond_info[0])));

paddle2onnx/mapper/while.cc

@@ -25,32 +25,12 @@ void ModelExporter::ExportWhile(PaddlePirParser& pir_parser,
   std::vector<TensorInfo> outputs_info;
   auto while_op = op->dyn_cast<paddle::dialect::WhileOp>();
   auto cond_info = pir_parser.GetTensorInfo(while_op.cond());
-  // mapping args and inputs in while op using while_op_input_value_map
-  std::vector<pir::detail::ValueImpl*> while_op_input_value_address;
-  std::vector<pir::detail::ValueImpl*> while_op_input_arg_address;
-  pir_parser.while_op_input_value_map
-      .clear();  // wangmingkai02: handle nested loop situations in future.
-
-  // record input value address
   for (int index = 1; index < while_op.num_operands(); index++) {
     const pir::Value& value = while_op.operand_source(index);
     inputs_info.push_back(pir_parser.GetTensorInfo(
-        pir_parser.GetOpOutputName(value), value.type()));
-    while_op_input_value_address.push_back(
-        &(*(value).impl()));  // get value address
-  }
-  // record args value address
-  std::vector<pir::Value> args = while_op.block_args();
-  for (int i = 0; i < args.size(); i++) {
-    const pir::Value& value = args[i];
-    while_op_input_arg_address.push_back(&(*(value.impl())));
-  }
-
-  // mapping
-  for (int index = 0; index < while_op_input_value_address.size(); index++) {
-    pir_parser.while_op_input_value_map[while_op_input_arg_address[index]] =
-        while_op_input_value_address[index];
+        pir_parser.GetSubBlockOpOutputName(value), value.type()));
   }
+  pir_parser.GetWhileInputValuesAndArgsMappings(while_op);
 
   std::vector<pir::Operation*> sub_blocks_ops_copy(pir_parser.sub_blocks_ops);
   pir_parser.sub_blocks_ops.clear();
@@ -120,7 +100,7 @@ void ModelExporter::ExportWhile(PaddlePirParser& pir_parser,
     input_names.push_back(inputs_info[i].name);
   }
   for (size_t i = 0; i < op->num_results(); i++) {
-    output_names.push_back(pir_parser.GetOpOutputName(op->result(i)));
+    output_names.push_back(pir_parser.GetSubBlockOpOutputName(op->result(i)));
   }
   auto loop_node = temp_helper->MakeNode("Loop", input_names, output_names);
   AddAttribute(loop_node, "body", graph);

paddle2onnx/parser/pir_parser.cc

@@ -106,10 +106,10 @@ std::string PaddlePirParser::GetOpOutputName(const pir::Value& source) const {
 
 std::string PaddlePirParser::GetSubBlockOpOutputName(
     const pir::Value& source) const {
-  auto it = while_op_input_value_map.find(&(*(source.impl())));
+  auto it = while_op_values_args_map.find(&(*(source.impl())));
   pir::Operation* op;
   uint32_t output_idx;
-  if (it != while_op_input_value_map.end()) {
+  if (it != while_op_values_args_map.end()) {
     pir::Value value(it->second);
     op = value.defining_op();
     output_idx = value.dyn_cast<pir::OpResult>().index();
@@ -1029,4 +1029,34 @@ P2ODataType PaddlePirParser::TransPirDataType2OldIrDataType(
            "PaddlePirParser::TransPirDataType2OnnxDataType.");
   }
 }
+void PaddlePirParser::GetWhileInputValuesAndArgsMappings(
+    const paddle::dialect::WhileOp& while_op) const {
+  // mapping args and inputs in while op using while_op_values_args_map
+  std::vector<pir::detail::ValueImpl*> while_op_input_value_address;
+  std::vector<pir::detail::ValueImpl*> while_op_input_arg_address;
+  // record input value address
+  for (int index = 1; index < while_op.num_operands(); index++) {
+    const pir::Value& value = while_op.operand_source(index);
+    while_op_input_value_address.push_back(
+        &(*(value).impl()));  // get value address
+  }
+  // record args value address
+  std::vector<pir::Value> args = while_op.block_args();
+  for (int i = 0; i < args.size(); i++) {
+    const pir::Value& value = args[i];
+    while_op_input_arg_address.push_back(&(*(value.impl())));
+  }
+
+  // mapping
+  for (int index = 0; index < while_op_input_value_address.size(); index++) {
+    auto arg_addr = while_op_input_arg_address[index];
+    if (while_op_values_args_map.count(arg_addr)) continue;
+    auto value_addr = while_op_input_value_address[index];
+    while (while_op_values_args_map.count(value_addr)) {
+      value_addr = while_op_values_args_map[value_addr];
+    }
+    while_op_values_args_map[arg_addr] = value_addr;
+  }
+}
+
 }  // namespace paddle2onnx

paddle2onnx/parser/pir_parser.h

@@ -22,6 +22,7 @@
 #include "paddle/pir/include/core/value.h"
 #include "paddle2onnx/parser/tensor_utils.h"
 #include "paddle2onnx/proto/p2o_paddle.pb.h"
+#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 namespace paddle2onnx {
 class PaddlePirParser {
  public:
@@ -40,8 +41,8 @@ class PaddlePirParser {
   // recoring set of operators for all blocks
   std::set<pir::Operation*> total_blocks_ops;
   // recording args of while op body name info
-  std::unordered_map<pir::detail::ValueImpl*, pir::detail::ValueImpl*>
-      while_op_input_value_map;
+  mutable std::unordered_map<pir::detail::ValueImpl*, pir::detail::ValueImpl*>
+      while_op_values_args_map;
   int NumOfBlocks() const;
   // int NumOfOps(int block_idx) const;
   int NumOfProgramOps() const;
@@ -265,6 +266,8 @@ class PaddlePirParser {
                           std::string tensor_arr_name) const;
   std::string GetTensorArrayName(int64_t op_id, bool if_in_sub_block) const;
   std::string GenOpInputOutputName(const std::string& name) const;
+  void GetWhileInputValuesAndArgsMappings(
+      const paddle::dialect::WhileOp& while_op) const;
 
  private:
   bool IsAttrVar(const pir::Operation* op, const int64_t& attr_id) const;