add decomp guard

Author: cyber-pioneer

python/paddle/base/core.py

@@ -518,7 +518,17 @@ 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():
-    if os.getenv("FLAGS_prim_skip_dynamic") == "1":
+    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")
+    if flag and flag.lower() in ("1", "true"):
         return True
     else:
         return False

python/paddle/decomposition/decomp.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 import typing
 import warnings
 
@@ -24,11 +25,25 @@
     has_decomp,
 )
 from paddle.base.libpaddle.pir import Block, Operation, Program
+from paddle.base.wrapped_decorator import signature_safe_contextmanager
 from paddle.framework import core
 
 from . import register
 
 
+# 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,)
@@ -176,16 +191,6 @@ def decompose(
     src_vars,
     blacklist=frozenset(),
     whitelist=frozenset(),
-):
-    blacklist = core.prim_config["forward_blacklist"] | blacklist
-    return core.decomp_tmp(program, src_vars, blacklist, whitelist)
-
-
-def decompose_(
-    program,
-    src_vars,
-    blacklist=frozenset(),
-    whitelist=frozenset(),
 ):
     """
     Search nonbasic ops which have be registered composite rules and replace them with primitive ops.
@@ -208,6 +213,9 @@ def decompose_(
     Returns:
         dst_vars (list): A list contains all vars which replace origin ones in src_vars.
     """
+    if core._enable_sink_decomp():
+        blacklist = core.prim_config["forward_blacklist"] | blacklist
+        return core.decomp_tmp(program, src_vars, blacklist, whitelist)
     if not core._is_fwd_prim_enabled():
         return src_vars
     if not isinstance(program, Program):

test/prim/pir_prim/test_prim_dynamic.py

@@ -65,12 +65,18 @@ def base_net(self, flag=None):
         return outs
 
     def test_prim_all_dynamic(self):
-        # os.environ["FLAGS_prim_skip_dynamic"] = "1"
         res_ref = self.base_net()
         res = self.base_net("all")
         for ref, actual in zip(res_ref, res):
             np.testing.assert_allclose(ref, actual, rtol=1e-6)
 
+    def test_aprim_all_dynamic_sink(self):
+        with decomp.sink_decomp_guard():
+            res_ref = self.base_net()
+            res = self.base_net("all")
+            for ref, actual in zip(res_ref, res):
+                np.testing.assert_allclose(ref, actual, rtol=1e-6)
+
 
 if __name__ == "__main__":
     unittest.main()