fix jit prim pir eval cinn (#59356)

cyber-pioneer · web-flow · commit c9d945d96f7f · 2023-11-26T17:46:52.000+08:00
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -536,11 +536,11 @@ def _create_program(self, is_infer_mode=False):
         if is_infer_mode:
             # TODO(xiongkun) who to transfer the pruning program?
             infer_program = self.origin_runable_program.clone()
+            if self._hooker:
+                self._hooker.after_infer(infer_program)
             infer_program = PirPassContext.apply(
                 infer_program, self._build_strategy
             )
-            if self._hooker:
-                self._hooker.after_infer(infer_program)
             return infer_program
         else:
             train_program: RunableProgram = self.origin_runable_program.clone()
diff --git a/test/ir/pir/cinn/test_cinn_sub_graph.py b/test/ir/pir/cinn/test_cinn_sub_graph.py
@@ -226,5 +226,43 @@ def test_forward(self):
         # np.testing.assert_allclose(cinn_out.numpy(), dy_out.numpy(), atol=1e-8)
 
 
+class TestCinnEvalPrim(TestCinnSubGraphBase):
+    def prepare_data(self):
+        self.shape = [1, 2048, 768]
+        self.hidden_states = paddle.randn(self.shape, dtype="float32")
+        self.hidden_states.stop_gradient = False
+
+    def eval(self, use_cinn):
+        paddle.seed(2022)
+        net = CINNSoftmaxSubGraphNet()
+        if use_cinn:
+            net = apply_to_static(net, True)
+        net.eval()
+        out = net(self.hidden_states)
+
+        if use_cinn:
+            ops = [
+                op.name()
+                for op in net.forward.program_cache.last()[-1][-1]
+                .train_program.program.global_block()
+                .ops
+            ]
+            assert (
+                "pd_op.softmax" not in ops
+            ), f"after prim, pd_op.softmax should not exist, but got {ops}"
+            assert (
+                "pd_op.exp" in ops
+            ), f"after prim, pd_op.softmax should not exist, but got {ops}"
+
+        return out
+
+    def test_eval(self):
+        cinn_out = self.eval(use_cinn=True)
+        dy_out = self.eval(use_cinn=False)
+        np.testing.assert_allclose(
+            cinn_out.numpy(), dy_out.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+
 if __name__ == '__main__':
     unittest.main()
