[Prim][CINN] Restore insertion point after decomp (#71960)

SigureMo · web-flow · commit b8e1c0f3d2c0 · 2025-03-28T10:57:00.000+08:00
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -1015,29 +1015,34 @@ void BindVjp(pybind11::module *m) {
 }
 
 void BindDecompRule(pybind11::module *m) {
-  m->def("sinking_decomp",
-         [](pir::Program *program,
-            std::vector<pir::Value> &src_vars,
-            std::set<std::string> &blacklist,
-            std::set<std::string> &whitelist,
-            int start_index,
-            int end_index) {
-           VLOG(4) << "[Prim] Bind Decomp sinking_decomp begin.";
-           py::list res;
-           DecompProgram decomp_object(
-               program, src_vars, blacklist, whitelist, start_index, end_index);
-           decomp_object.decomp_program();
-           std::vector<pir::Value> tar_vars = decomp_object.get_dst_vars();
-           for (size_t i = 0; i < tar_vars.size(); ++i) {
-             if (!tar_vars[i]) {
-               res.append(nullptr);
-             } else {
-               res.append(tar_vars[i]);
-             }
-           }
-           VLOG(4) << "[Prim] Bind Decomp sinking_decomp end.";
-           return res;
-         });
+  m->def(
+      "sinking_decomp",
+      [](pir::Program *program,
+         std::vector<pir::Value> &src_vars,
+         std::set<std::string> &blacklist,
+         std::set<std::string> &whitelist,
+         int start_index,
+         int end_index) {
+        VLOG(4) << "[Prim] Bind Decomp sinking_decomp begin.";
+        py::list res;
+        auto original_insertion_point =
+            paddle::dialect::ApiBuilder::Instance().GetCurrentInsertionPoint();
+        DecompProgram decomp_object(
+            program, src_vars, blacklist, whitelist, start_index, end_index);
+        decomp_object.decomp_program();
+        std::vector<pir::Value> tar_vars = decomp_object.get_dst_vars();
+        for (size_t i = 0; i < tar_vars.size(); ++i) {
+          if (!tar_vars[i]) {
+            res.append(nullptr);
+          } else {
+            res.append(tar_vars[i]);
+          }
+        }
+        paddle::dialect::ApiBuilder::Instance().SetInsertionPoint(
+            original_insertion_point);
+        VLOG(4) << "[Prim] Bind Decomp sinking_decomp end.";
+        return res;
+      });
 
   m->def("call_decomp_rule", [](pir::Operation &fwd_op) {
     py::list res;
diff --git a/test/dygraph_to_static/test_slice.py b/test/dygraph_to_static/test_slice.py
@@ -205,7 +205,9 @@ def test_set_value_with_save(self):
 class TestSliceSupplementSpecialCase(Dy2StTestBase):
     # unittest for slice index which abs(step)>0. eg: x[::2]
     def test_static_slice_step(self):
-        with static_guard():
+        with static_guard(), paddle.static.program_guard(
+            paddle.static.Program()
+        ):
             array = np.arange(4**3).reshape((4, 4, 4)).astype('int64')
 
             x = paddle.static.data(name='x', shape=[4, 4, 4], dtype='int64')
@@ -253,18 +255,10 @@ def test_compare_paddle_strided_slice_with_numpy(self):
         stride1 = -2
         sl = paddle.strided_slice(
             pt,
-            axes=[
-                0,
-            ],
-            starts=[
-                s1,
-            ],
-            ends=[
-                e1,
-            ],
-            strides=[
-                stride1,
-            ],
+            axes=[0],
+            starts=[s1],
+            ends=[e1],
+            strides=[stride1],
         )
 
         self.assertTrue(array[s1:e1:stride1], sl)
