support enabling prim via ++prim=1 (#843)

Author: HydrogenSulfate

ppsci/utils/callbacks.py

@@ -97,3 +97,12 @@ def on_job_start(self, config: DictConfig, **kwargs: Any) -> None:
             else None,
             full_cfg.log_level,
         )
+
+        # enable prim if specified
+        if "prim" in full_cfg and bool(full_cfg.prim):
+            # Mostly for dy2st running, will be removed in the future
+            from paddle.framework import core
+
+            core.set_prim_eager_enabled(True)
+            core._set_prim_all_enabled(True)
+            logger.message("Prim mode is enabled.")

ppsci/utils/config.py

@@ -219,6 +219,7 @@ class SolverConfig(BaseModel):
         use_amp: bool = False
         amp_level: Literal["O0", "O1", "O2", "OD"] = "O1"
         to_static: bool = False
+        prim: bool = False
         log_level: Literal["debug", "info", "warning", "error"] = "info"
 
         # Training related config