[Prim][PIR]Prim + to_static + dynamic shape (#61315)

cyber-pioneer · web-flow · commit 0428df0893d3 · 2024-01-30T19:10:46.000+08:00
* prim and jit and dynamic shape

* remove unused code

* fix test case

* polish test case
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -246,7 +246,7 @@ def _check_vjp_dynamic_shape(op, inputs):
 def dynamic_shape_prim_vjp_guard(op, inputs):
     skip_prim = (
         core._is_bwd_prim_enabled()
-        and core._enable_prim_dynamic_shape()
+        and core._enable_prim_skip_dynamic_shape()
         and _check_vjp_dynamic_shape(op, inputs)
     )
     try:
diff --git a/python/paddle/base/core.py b/python/paddle/base/core.py
@@ -516,17 +516,16 @@ def _get_batch_norm_none_var(op):
 
 
 # This api is used for development for dynamic shape in prim, and will be removed in future.
-def _enable_prim_dynamic_shape():
+def _enable_prim_skip_dynamic_shape():
     flag = os.getenv("FLAGS_prim_skip_dynamic")
     if flag and flag.lower() in ("1", "true"):
         return True
     else:
         return False
 
 
-# This api is used for development for sinking decomp in c++, and will be removed in future.
-def _enable_sink_decomp():
-    flag = os.getenv("FLAGS_sink_decomp", "true")
+def _enable_prim_dynamic_shape():
+    flag = os.getenv("FLAGS_prim_enable_dynamic")
     if flag and flag.lower() in ("1", "true"):
         return True
     else:
diff --git a/python/paddle/decomposition/decomp.py b/python/paddle/decomposition/decomp.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-import os
 import typing
 import warnings
 
@@ -25,28 +24,14 @@
     decomp_ops_contain_unused_output,
     has_decomp,
 )
-from paddle.base.libpaddle.pir import Block, Operation, Program
-from paddle.base.wrapped_decorator import signature_safe_contextmanager
+from paddle.base.libpaddle.pir import Block, Operation
 from paddle.framework import core
 
 from . import register
 
 logger = logging.getLogger(__name__)
 
 
-# For sinking decomp in c++. In future, sinking decomp will be implemented in c++ by default and then this api will be removed.
-@signature_safe_contextmanager
-def sink_decomp_guard():
-    sink_decomp = core._enable_sink_decomp()
-    try:
-        if not sink_decomp:
-            os.environ['FLAGS_sink_decomp'] = 'true'
-        yield
-    finally:
-        if not sink_decomp:
-            os.environ['FLAGS_sink_decomp'] = 'false'
-
-
 def _build_tensor_tuple(xs):
     if isinstance(xs, pir.Value):
         return (xs,)
@@ -216,161 +201,8 @@ def decompose(
     Returns:
         dst_vars (list): A list contains all vars which replace origin ones in src_vars.
     """
-    # flag default status: True
-    if core._enable_sink_decomp():
-        blacklist = core.prim_config["forward_blacklist"] | blacklist
-        return core.sinking_decomp(program, src_vars, blacklist, whitelist)
-    if not core._is_fwd_prim_enabled():
-        return src_vars
-    if not isinstance(program, Program):
-        raise TypeError(f"Expect type Program, but got type {type(program)}.")
-    block = program.global_block()
-
-    if not isinstance(blacklist, (set, frozenset)):
-        raise TypeError(
-            f'Expected type of blacklist is set|frozenset, but got {type(blacklist)}.'
-        )
-    if not isinstance(whitelist, (set, frozenset)):
-        raise TypeError(
-            f'Expected type of whitelist is set|frozenset, but got {type(whitelist)}.'
-        )
-
     blacklist = core.prim_config["forward_blacklist"] | blacklist
-
-    logger.debug(
-        f'Lowering source program into primitive program. Source program is: \n \
-{program}\n \
-The finally set that will be decomposed is: (ops & ops have decomposite rule & whitelist) - blacklist.\n \
-The blacklist which will be ignored in lowering is: {blacklist},\n \
-and the whitelist which will be processed in lowering is: {whitelist}.'
-    )
-
-    if len(blacklist) > 0 and len(whitelist) > 0:
-        op_filter = (
-            lambda x: x.name() in whitelist and x.name() not in blacklist
-        )
-    elif len(blacklist) > 0 and len(whitelist) == 0:
-        op_filter = lambda x: x.name() not in blacklist
-    elif len(blacklist) == 0 and len(whitelist) > 0:
-        op_filter = lambda x: x.name() in whitelist
-    else:
-        op_filter = lambda x: True
-    dst_vars = [None] * len(src_vars)
-    dst_vars_dct = ValueDict()
-    for idx, item in enumerate(src_vars):
-        if not isinstance(item, pir.Value):
-            raise TypeError(
-                f"Each var in dst_vars should map corresponding var in src_vars, but got type {type(item)} in {src_vars}."
-            )
-        dst_vars_dct[item] = idx
-    with pir.core.program_guard(program):
-        _decompose_subgraph(
-            block,
-            dst_vars_dct,
-            dst_vars,
-            op_filter,
-        )
-    for idx, item in enumerate(dst_vars):
-        if not isinstance(item, pir.Value):
-            if item is None:
-                dst_vars[idx] = src_vars[idx]
-            else:
-                raise TypeError(
-                    f"Each var in dst_vars should map corresponding var in src_vars, but got type {type(item)} in {dst_vars}."
-                )
-    logger.debug(
-        "Decompose composite forward ops finish: {}".format(
-            core.prim_config["composite_ops_record"]
-        )
-    )
-    return dst_vars
-
-
-def _decompose_subgraph(block, orig_vars, dst_vars, op_filter):
-    """
-    The operators in block which satisfy the filter condition will be decomposed into primitives.
-
-    Args:
-        block (Block|Sequence[Block]): The blocks of program to be processed.
-        op_filter (function): The filter to specify which ops to be processed.
-        orig_vars (dict): Origin variables of original block.
-        dst_vars (list): Corresponding replaced variables of Origin variables.
-    """
-
-    if isinstance(block, Block):
-        ops_list = block.ops
-        temp_op = None
-        for idx, op in enumerate(ops_list):
-            op_name = op.name()
-            decom_rule = register.get_decomp_rule(op_name)
-            has_sink_decomp_rule = has_decomp(op)
-            lower = (decom_rule or has_sink_decomp_rule) and op_filter(op)
-
-            if (
-                lower
-                and core._enable_prim_dynamic_shape()
-                and _check_prim_dynamic(op)
-            ):
-                lower = False
-
-            if op.name() == "builtin.combine":
-                temp_op = op
-
-            if lower:
-                core.prim_config["composite_ops_record"].add(op_name)
-                if (
-                    temp_op is not None
-                    and ops_list[idx - 1].name() == "builtin.combine"
-                ):
-                    pir.set_insertion_point(temp_op)
-                else:
-                    pir.set_insertion_point(op)
-                input_args = _prepare_python_api_arguments(op)
-                orig_outs = op.results()
-                if has_sink_decomp_rule:
-                    decomp_outs = call_decomp(op)
-                    new_outs = _analyse_decomp_results(
-                        orig_outs, decomp_outs, op
-                    )
-                else:
-                    new_outs = _build_tensor_tuple(decom_rule(*input_args))
-
-                # Todo: To cover such case: some outputs are no longer needed after decomposition.
-                _check_op_results(
-                    op_name, orig_outs, new_outs, orig_vars, dst_vars
-                )
-                if op.name() in decomp_ops_contain_unused_output.keys():
-                    for idx in range(len(orig_outs)):
-                        if (
-                            idx
-                            not in decomp_ops_contain_unused_output[op.name()]
-                        ):
-                            orig_outs[idx].replace_all_uses_with(new_outs[idx])
-                else:
-                    if op.name() in decomp_ops_contain_unused_output.keys():
-                        orig_outs[0].replace_all_uses_with(new_outs[0])
-                    else:
-                        op.replace_all_uses_with(new_outs)
-                block.remove_op(op)
-
-                if temp_op is not None:
-                    remove_op = True
-                    for item in temp_op.results():
-                        if item.has_one_use():
-                            remove_op = False
-                            break
-                    if remove_op:
-                        block.remove_op(temp_op)
-                    temp_op = None
-        return
-
-    elif isinstance(block, typing.Sequence):
-        for item in block:
-            _decompose_subgraph(item, orig_vars, dst_vars, op_filter)
-        return
-    raise TypeError(
-        f"Expect type Block or Sequence of Block, but got type {type(block)}"
-    )
+    return core.sinking_decomp(program, src_vars, blacklist, whitelist)
 
 
 def _check_combine_inputs(input1, input2):
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -350,7 +350,11 @@ def __init__(self, function, input_spec=None, **kwargs):
             self._dygraph_function = function
             self._class_instance = None
         # TODO(chenzhuo): Remove this after lowering prim into C++
-        if input_spec is not None and prim_is_enabled():
+        if (
+            input_spec is not None
+            and prim_is_enabled()
+            and not core._enable_prim_dynamic_shape()
+        ):
             from paddle.static import InputSpec
 
             for spec in flatten(input_spec):
diff --git a/test/prim/pir_prim/CMakeLists.txt b/test/prim/pir_prim/CMakeLists.txt
@@ -6,11 +6,18 @@ set(TEST_PRIM_PURE_PIR_CASES
     test_pir_prim_flags
     test_sink_decomp
     test_prim_skip_dynamic
-    test_prim_dynamic)
+    test_prim_dynamic
+    test_prim_jit_dynamic)
 
 foreach(target ${TEST_PRIM_PURE_PIR_CASES})
-  py_test_modules(${target} MODULES ${target} ENVS GLOG_v=1
-                  FLAGS_enable_pir_api=true)
+  py_test_modules(
+    ${target}
+    MODULES
+    ${target}
+    ENVS
+    GLOG_v=1
+    FLAGS_enable_pir_api=true
+    FLAGS_prim_enable_dynamic=true)
 endforeach()
 
 file(
diff --git a/test/prim/pir_prim/test_prim_jit_dynamic.py b/test/prim/pir_prim/test_prim_jit_dynamic.py
@@ -0,0 +1,85 @@
+# 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
+from paddle.static import InputSpec
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    build_strategy = paddle.static.BuildStrategy()
+    build_strategy.build_cinn_pass = use_cinn
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        build_strategy=build_strategy,
+        full_graph=True,
+    )
+
+
+def rms_norm(hidden_states, weight):
+    variance = hidden_states.pow(2).mean((0, 1), keepdim=True)
+    hidden_states = paddle.rsqrt(variance + 1e-5) * hidden_states
+    return hidden_states * weight
+
+
+class TestPrimMode(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [1, 300, 4096]
+        self.shape_y = [4096]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.y = np.random.random(self.shape_y).astype("float32")
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+        x = paddle.to_tensor(self.x)
+        y = paddle.to_tensor(self.y)
+        fn = apply_to_static(
+            rms_norm,
+            use_cinn=False,
+            input_spec=[
+                InputSpec(shape=[None, None, 4096], dtype='float32'),
+                InputSpec(shape=[4096], dtype='float32'),
+            ],
+        )
+        res = fn(x, y)
+        ops = [
+            op.name()
+            for op in fn.program_cache.last()[-1][-1]
+            .infer_program.program.global_block()
+            .ops
+        ]
+
+        if flag == "prim":
+            assert "pd_op.mean" not in ops
+            core._set_prim_all_enabled(False)
+        else:
+            assert "pd_op.mean" in ops
+        return res
+
+    def test_prim_all_dynamic(self):
+        res_ref = self.base_net()
+        res = self.base_net("prim")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,7 @@ def _check_vjp_dynamic_shape(op, inputs):`
`246`	`246`	`def dynamic_shape_prim_vjp_guard(op, inputs):`
`247`	`247`	`skip_prim = (`
`248`	`248`	`core._is_bwd_prim_enabled()`
`249`		`- and core._enable_prim_dynamic_shape()`
	`249`	`+ and core._enable_prim_skip_dynamic_shape()`
`250`	`250`	`and _check_vjp_dynamic_shape(op, inputs)`
`251`	`251`	`)`
`252`	`252`	`try:`