[PIR save/load]Fix bug in PirInterpreterEngine and open jit layer test (#72101)

changeyoung98 · web-flow · commit 3e00f80703cc · 2025-04-09T10:38:58.000+08:00
* fix bug in PirInterpreterEngine and open jit layer test

* refine
diff --git a/paddle/fluid/jit/engine/pir_interpreter_engine.cc b/paddle/fluid/jit/engine/pir_interpreter_engine.cc
@@ -35,6 +35,7 @@ PirInterpreterEngine::PirInterpreterEngine(
                     common::errors::PreconditionNotMet(
                         "There is no operator in ProgramDesc."));
   utils::ShareParamsIntoScope(info_->ParamNames(), params_dict_, &scope_);
+  prog_ = paddle::dialect::PdOpLowerToKernelPass(prog_.get(), place_);
   CreateInterpreterCore();
 }
 
@@ -60,7 +61,6 @@ std::vector<Tensor> PirInterpreterEngine::operator()(
 
 std::vector<DenseTensor> PirInterpreterEngine::operator()(
     const std::vector<DenseTensor> &inputs) {
-  prog_ = paddle::dialect::PdOpLowerToKernelPass(prog_.get(), place_);
   utils::ShareIntoScope(info_->InputArgNames(), inputs, &scope_);
 
   // the latter can be moved to python side.
diff --git a/test/legacy_test/test_jit_layer.py b/test/legacy_test/test_jit_layer.py
@@ -68,26 +68,23 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
     def test_multi_load(self):
-        with paddle.pir_utils.OldIrGuard():
-            paddle.disable_static()
-            x = paddle.full([2, 4], 2)
-            model = Net()
-            with enable_to_static_guard(False):
-                forward_out1 = model.forward(x)
-                infer_out1 = model.infer(x)
-            model_path = os.path.join(self.temp_dir.name, 'multi_program')
-            paddle.jit.save(model, model_path, combine_params=True)
-            place = paddle.CPUPlace()
-            if paddle.is_compiled_with_cuda():
-                place = paddle.CUDAPlace(0)
-            jit_layer = Layer()
-            jit_layer.load(model_path, place)
-            forward_out2 = jit_layer.forward(x)
-            infer_out2 = jit_layer.infer(x)
-            np.testing.assert_allclose(
-                forward_out1, forward_out2[0], rtol=1e-05
-            )
-            np.testing.assert_allclose(infer_out1, infer_out2[0], rtol=1e-05)
+        paddle.disable_static()
+        x = paddle.full([2, 4], 2)
+        model = Net()
+        with enable_to_static_guard(False):
+            forward_out1 = model.forward(x)
+            infer_out1 = model.infer(x)
+        model_path = os.path.join(self.temp_dir.name, 'multi_program')
+        paddle.jit.save(model, model_path, combine_params=True)
+        place = paddle.CPUPlace()
+        if paddle.is_compiled_with_cuda():
+            place = paddle.CUDAPlace(0)
+        jit_layer = Layer()
+        jit_layer.load(model_path, place)
+        forward_out2 = jit_layer.forward(x)
+        infer_out2 = jit_layer.infer(x)
+        np.testing.assert_allclose(forward_out1, forward_out2[0], rtol=1e-05)
+        np.testing.assert_allclose(infer_out1, infer_out2[0], rtol=1e-05)
 
     def test_multi_jit_load(self):
         x = paddle.full([2, 4], 2)
@@ -127,20 +124,18 @@ def tearDown(self):
         self.temp_dir.cleanup()
 
     def test_mkl_output(self):
-        with paddle.pir_utils.OldIrGuard():
-            paddle.disable_static()
-            with _dygraph_place_guard(place=paddle.CPUPlace()):
-                net = SaveLinear()
-                model_path = os.path.join(self.temp_dir.name, 'save_linear')
-                paddle.jit.save(net, model_path, combine_params=True)
-
-                layer = Layer()
-                print("load ", model_path)
-                layer.load(model_path, paddle.CPUPlace())
-                x = paddle.ones([498, 80])
-                out = layer.forward(x)
-                out = paddle.unsqueeze(out[0], 0)
-                np.testing.assert_equal(out.shape, [1, 498, 80])
+        paddle.disable_static()
+        with _dygraph_place_guard(place=paddle.CPUPlace()):
+            net = SaveLinear()
+            model_path = os.path.join(self.temp_dir.name, 'save_linear')
+            paddle.jit.save(net, model_path, combine_params=True)
+
+            layer = Layer()
+            layer.load(model_path, paddle.CPUPlace())
+            x = paddle.ones([498, 80])
+            out = layer.forward(x)
+            out = paddle.unsqueeze(out[0], 0)
+            np.testing.assert_equal(out.shape, [1, 498, 80])
 
     def test_mkl_jit_output(self):
         with _dygraph_place_guard(place=paddle.CPUPlace()):
