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