[PIR] separate translation context for body block of while op (#58625)

kangguangli · web-flow · commit 79c6258790fb · 2023-11-07T11:27:55.000+08:00
* fix

* fix

* trigger CI

* fix
diff --git a/paddle/fluid/ir_adaptor/translator/program_translator.cc b/paddle/fluid/ir_adaptor/translator/program_translator.cc
@@ -524,41 +524,35 @@ void ProgramTranslator::TranslateWhileOperation(
       translation_ctx->at(loop_vars_reverse[0].first).value};
   std::vector<pir::Type> op_outputs_type;
   auto body_block = new pir::Block();
-  std::vector<TCValue> param_status;
+  auto* body_block_context = translation_ctx->CreateInnerContext();
   for (size_t idx = loop_vars_reverse.size() - 1u; idx > 0; --idx) {
     auto& name = loop_vars_reverse[idx].first;
     auto& tc_value = translation_ctx->at(name);
     auto val_type = tc_value.value.type();
     op_inputs.push_back(tc_value.value);
     op_outputs_type.push_back(val_type);
-    param_status.emplace_back(tc_value);
-    translation_ctx->PushValue(name, body_block->AddArgument(val_type));
+    body_block_context->PushValue(name, body_block->AddArgument(val_type));
   }
   pir::Operation* while_op =
       pir::Operation::Create(op_inputs, {}, op_outputs_type, op_info, 1);
   dst_block->push_back(while_op);
   while_op->region(0).push_back(body_block);
-  TranslateBlock(sub_block, 0, index + 1, translation_ctx, body_block);
+  TranslateBlock(sub_block, 0, index + 1, body_block_context, body_block);
 
   auto yeild_info = ctx_->GetRegisteredOpInfo(pir::YieldOp::name());
   std::vector<pir::Value> yeild_inputs{
-      translation_ctx->at(loop_vars_reverse[0].second).value};
+      body_block_context->at(loop_vars_reverse[0].second).value};
   for (size_t idx = loop_vars_reverse.size() - 1u; idx > 0; --idx) {
     auto& name = loop_vars_reverse[idx].second;
-    yeild_inputs.push_back(translation_ctx->at(name).value);
+    yeild_inputs.push_back(body_block_context->at(name).value);
   }
   body_block->push_back(
       pir::Operation::Create(yeild_inputs, {}, {}, yeild_info));
-
-  index = 0;
-  for (size_t idx = loop_vars_reverse.size() - 1u; idx > 0; --idx) {
-    auto& name = loop_vars_reverse[idx].first;
-    translation_ctx->PushValue(name, param_status[index++]);
-  }
   auto name_iter = loop_vars_reverse.rbegin();
   for (size_t idx = 0; idx < while_op->num_results(); ++idx) {
     translation_ctx->PushValue(name_iter++->first, while_op->result(idx));
   }
+
   while_op->Verify();
   VLOG(8) << "=============>end to translate while op:" << op;
 }
