support pir jit prim (#57561)

cyber-pioneer · web-flow · commit b13dcb85918b · 2023-09-21T14:13:23.000+08:00
diff --git a/python/paddle/jit/dy2static/newir_partial_program.py b/python/paddle/jit/dy2static/newir_partial_program.py
@@ -642,11 +642,15 @@ def _insert_aggregation_ops_for_var(target_program, var):
     @switch_to_static_graph
     def _append_backward_desc(self, main_program):
         program = main_program
-        # if self._hooker:
-        # program = self._hooker.before_append_backward(program)
+
         targets = list(
             filter(lambda x: isinstance(x, OpResult), self._outputs.tolist())
         )
+        if self._hooker:
+            program, targets = self._hooker.before_append_backward(
+                program, targets
+            )
+            self._outputs = NestSequence(targets, need_check=True)
         inputs = list(
             filter(lambda x: isinstance(x, OpResult), self._inputs.tolist())
         )
@@ -676,11 +680,15 @@ def _append_backward_desc(self, main_program):
                     forward_outputs_grads.append(opres)
                     not_stop_gradient_num += 1
 
-            # TODO: add later.
-            # if self._hooker:
-            # program, start_idx = self._hooker.after_append_backward(
-            # program, start_idx
-            # )
+            if self._hooker:
+                (
+                    program,
+                    forward_end_idx,
+                    targets,
+                ) = self._hooker.after_append_backward(
+                    program, targets, forward_end_idx
+                )
+                self._outputs = NestSequence(targets, need_check=True)
 
             # TODO: add later
             # self.prepare_gradient_aggregation(
@@ -692,6 +700,8 @@ def _append_backward_desc(self, main_program):
         )
         hash_id = paddle.utils._hash_with_id(program, self)
         extra_info = self._program_extra_info.get(hash_id, {})
+        extra_info['forward_inputs'] = inputs
+        extra_info['forward_outputs'] = targets
         extra_info['forward_end_op_idx'] = forward_end_idx
         extra_info['forward_inputs_grads'] = list(
             map(mapping_op_result, grad_info_map)
@@ -791,8 +801,10 @@ def _get_forward_backward_program_form(
         forward_inputs_grads = self.get_program_extra(whole_program)[
             'forward_inputs_grads'
         ]
-        forward_inputs = self._inputs.tolist()
-        forward_outputs = self._outputs.tolist()
+        forward_inputs = self.get_program_extra(whole_program)['forward_inputs']
+        forward_outputs = self.get_program_extra(whole_program)[
+            'forward_outputs'
+        ]
         forward_outputs_grads = self.get_program_extra(whole_program)[
             'forward_outputs_grads'
         ]
@@ -947,9 +959,11 @@ def create_out(var_id):
                 tensor_type = paddle.dtype(8)  # SELECT ROW TENSOR
 
             # TODO(xiongkun): more elegent way to do it.
+
             ir_dtype_2_tensor_dtype = {
                 10: paddle.dtype(5),
             }
+
             out = core.eager.Tensor(
                 ir_dtype_2_tensor_dtype[int(var.dtype)],
                 var.shape,
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -21,6 +21,7 @@
 import weakref
 
 import paddle.ir.core as ir_static
+from paddle import decomposition
 from paddle.base import core, framework
 from paddle.base.data_feeder import check_type
 from paddle.base.dygraph.base import (
@@ -42,6 +43,9 @@
     get_buffers,
     get_parameters,
 )
+from .newir_partial_program import (
+    PartialProgramLayerHook as PirPartialProgramLayerHook,
+)
 from .origin_info import (
     attach_origin_info,
     create_and_update_origin_info_map,
@@ -1473,6 +1477,46 @@ def __setattr__(self, key, value):
         return super().__setattr__(key, value)
 
 
+class PirPrimHooker(PirPartialProgramLayerHook):
+    def __init__(self, original_program, backend):
+        self.backend = backend
+        self.custom_vjps = set()
+        with backend_guard(self.backend):
+            if core._is_all_prim_enabled():
+                self.custom_vjps = {
+                    op.name()
+                    for op in original_program.global_block().ops
+                    if core.has_custom_vjp(op)
+                }
+
+    def before_append_backward(self, forward_program, src_vars):
+        with backend_guard(self.backend):
+            if core._is_fwd_prim_enabled():
+                dst_vars = decomposition.decompose(
+                    forward_program, src_vars, blacklist=self.custom_vjps
+                )
+            return forward_program, dst_vars
+
+    def after_append_backward(self, whole_program, src_vars, forward_end_idx):
+        with backend_guard(self.backend):
+            backward_length = (
+                len(whole_program.global_block().ops) - forward_end_idx
+            )
+            if core._is_fwd_prim_enabled() and len(self.custom_vjps) != 0:
+                # only process backward part of block
+                dst_vars = decomposition.decompose(whole_program, src_vars)
+            new_start_index = (
+                len(whole_program.global_block().ops) - backward_length
+            )
+            return whole_program, new_start_index, dst_vars
+
+    def after_infer(self, infer_program, src_vars):
+        with backend_guard(self.backend):
+            if core._is_fwd_prim_enabled():
+                dst_vars = decomposition.decompose(infer_program, src_vars)
+            return infer_program, dst_vars
+
+
 class ProgramCache:
     """
     Wrapper class for the program functions defined by dygraph function.
@@ -1530,7 +1574,10 @@ def _build_once(self, cache_key):
                 raise
 
         backend = cache_key.kwargs['backend']
-        if prim_or_cinn_is_enabled(cache_key.kwargs['build_strategy'], backend):
+        if (
+            prim_or_cinn_is_enabled(cache_key.kwargs['build_strategy'], backend)
+            and not use_pir_api()
+        ):
             for var in concrete_program.main_program.list_vars():
                 if var.type not in NO_SHAPE_VAR_TYPE and -1 in var.shape:
                     warnings.warn(
@@ -1553,9 +1600,14 @@ def _build_once(self, cache_key):
             )
         with backend_guard(backend):
             if core._is_fwd_prim_enabled():
-                partial_program.set_hooker(
-                    PrimHooker(concrete_program.main_program, backend)
-                )
+                if use_pir_api():
+                    partial_program.set_hooker(
+                        PirPrimHooker(concrete_program.main_program, backend)
+                    )
+                else:
+                    partial_program.set_hooker(
+                        PrimHooker(concrete_program.main_program, backend)
+                    )
         return concrete_program, partial_program
 
     def __getitem__(self, item):
diff --git a/test/prim/new_ir_prim/CMakeLists.txt b/test/prim/new_ir_prim/CMakeLists.txt
@@ -1,5 +1,5 @@
 set(TEST_PRIM_PURE_NEW_IR_CASES test_prim_program test_prim_simpnet
-                                test_prim_custom_vjp)
+                                test_prim_custom_vjp test_prim_jit)
 
 foreach(target ${TEST_PRIM_PURE_NEW_IR_CASES})
   py_test_modules(${target} MODULES ${target} ENVS GLOG_v=1
diff --git a/test/prim/new_ir_prim/test_prim_jit.py b/test/prim/new_ir_prim/test_prim_jit.py
@@ -0,0 +1,58 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.framework import core
+
+
+class TestDy2staticNewIR(unittest.TestCase):
+    def test_basic_network_backward(self):
+        core._set_prim_all_enabled(True)
+
+        def func(x):
+            x1 = paddle.mean(x)
+            out = paddle.nn.functional.gelu(x1, False)
+            return out
+
+        # ==== dygraph computation ====
+        static_func = paddle.jit.to_static(func)
+        x = paddle.randn((8, 16, 64))
+        x.stop_gradient = False
+        ref_out = func(x) * 2
+        ref_out.backward()
+        ref_grad = x.grad.numpy()
+        x.clear_gradient()
+
+        # ==== to static compuatation ====
+        out = static_func(x)
+        actual_out = out * 2
+        actual_out.backward()
+        actual_grad = x.grad
+        core._set_prim_all_enabled(True)
+
+        np.testing.assert_allclose(
+            ref_out, actual_out.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+        np.testing.assert_allclose(
+            ref_grad, actual_grad.numpy(), atol=1e-6, rtol=1e-6
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
