From ca5585e9a26b53468e850bc451a9bf5a151ca1e1 Mon Sep 17 00:00:00 2001
From: YuanRisheng <yuanrisheng@baidu.com>
Date: Fri, 25 Aug 2023 14:06:51 +0800
Subject: [PATCH] [BugFix]Fix test_build_model error (#56633)

* fix test bugs

* delete code
---
 paddle/fluid/ir/dialect/paddle_dialect/ir/api_builder.cc | 2 --
 paddle/fluid/pybind/ir.cc                                | 4 ++++
 python/paddle/ir/__init__.py                             | 1 +
 python/paddle/new_ir_utils.py                            | 2 +-
 4 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/ir/dialect/paddle_dialect/ir/api_builder.cc b/paddle/fluid/ir/dialect/paddle_dialect/ir/api_builder.cc
index c72ee0b6b2f..0ded4ee1a5d 100644
--- a/paddle/fluid/ir/dialect/paddle_dialect/ir/api_builder.cc
+++ b/paddle/fluid/ir/dialect/paddle_dialect/ir/api_builder.cc
@@ -15,14 +15,12 @@
 #include "paddle/fluid/ir/dialect/paddle_dialect/ir/api_builder.h"
 #include "paddle/ir/core/enforce.h"
 #include "paddle/ir/core/ir_context.h"
-// #include "paddle/fluid/ir/dialect/paddle_dialect/ir/pd_dialect.h"
 
 namespace paddle {
 namespace dialect {
 
 APIBuilder::APIBuilder() : builder_(nullptr) {
   ctx_ = ir::IrContext::Instance();
-  // ctx_->GetOrRegisterDialect<paddle::dialect::PaddleDialect>();
 }
 
 void APIBuilder::SetProgram(ir::Program* program) {
diff --git a/paddle/fluid/pybind/ir.cc b/paddle/fluid/pybind/ir.cc
index 1f42052ca7a..6c6957c3e00 100644
--- a/paddle/fluid/pybind/ir.cc
+++ b/paddle/fluid/pybind/ir.cc
@@ -410,6 +410,10 @@ void BindUtils(pybind11::module *m) {
          []() { APIBuilder::Instance().ResetInsertionPointToStart(); });
   m->def("reset_insertion_point_to_end",
          []() { APIBuilder::Instance().ResetInsertionPointToEnd(); });
+  m->def("register_paddle_dialect", []() {
+    ir::IrContext::Instance()
+        ->GetOrRegisterDialect<paddle::dialect::PaddleDialect>();
+  });
   m->def(
       "translate_to_new_ir",
       [](const ::paddle::framework::ProgramDesc &legacy_program) {
diff --git a/python/paddle/ir/__init__.py b/python/paddle/ir/__init__.py
index be8ddeba229..4fee1c1a064 100644
--- a/python/paddle/ir/__init__.py
+++ b/python/paddle/ir/__init__.py
@@ -28,6 +28,7 @@ from paddle.fluid.libpaddle.ir import (
     reset_insertion_point_to_start,
     reset_insertion_point_to_end,
     check_unregistered_ops,
+    register_paddle_dialect,
     PassManager,
 )  # noqa: F401
 
diff --git a/python/paddle/new_ir_utils.py b/python/paddle/new_ir_utils.py
index bb016158d16..443ac48ae82 100644
--- a/python/paddle/new_ir_utils.py
+++ b/python/paddle/new_ir_utils.py
@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import paddle
 
 
 def _switch_to_new_ir():
     if paddle.ir.core._use_new_ir_api():
         paddle.framework.set_flags({"FLAGS_enable_new_ir_in_executor": True})
+        paddle.ir.register_paddle_dialect()
         paddle.static.Program = paddle.ir.Program
         paddle.fluid.Program = paddle.ir.Program
         paddle.fluid.program_guard = paddle.ir.core.program_guard
-- 
GitLab