Clarify CPU-specific thread warning

Author: caic99

deepmd/env.py

@@ -68,15 +68,21 @@ def set_env_if_empty(key: str, value: str, verbose: bool = True) -> None:
             )
 
 
-def set_default_nthreads() -> None:
+def set_default_nthreads(use_cpu: bool = False) -> None:
     """Set internal number of threads to default=automatic selection.
 
+    Parameters
+    ----------
+    use_cpu : bool, optional
+        If ``True``, suppress warnings about thread configuration,
+        by default ``False``.
+
     Notes
     -----
     `DP_INTRA_OP_PARALLELISM_THREADS` and `DP_INTER_OP_PARALLELISM_THREADS`
     control configuration of multithreading.
     """
-    if (
+    if not use_cpu and (
         "OMP_NUM_THREADS" not in os.environ
         # for backward compatibility
         or (
@@ -89,10 +95,10 @@ def set_default_nthreads() -> None:
         )
     ):
         log.warning(
-            "To get the best performance, it is recommended to adjust "
-            "the number of threads by setting the environment variables "
-            "OMP_NUM_THREADS, DP_INTRA_OP_PARALLELISM_THREADS, and "
-            "DP_INTER_OP_PARALLELISM_THREADS. See "
+            "To get the best CPU performance, adjust the number of threads by "
+            "setting the environment variables OMP_NUM_THREADS, "
+            "DP_INTRA_OP_PARALLELISM_THREADS, and DP_INTER_OP_PARALLELISM_THREADS. "
+            "These variables are only effective when running on CPU. See "
             "https://deepmd.rtfd.io/parallelism/ for more information."
         )
     if "TF_INTRA_OP_PARALLELISM_THREADS" not in os.environ:

deepmd/pd/utils/env.py

@@ -113,7 +113,7 @@ def to_bool(flag: int | bool | str) -> bool:
 DEFAULT_PRECISION = "float64"
 
 # throw warnings if threads not set
-set_default_nthreads()
+set_default_nthreads(use_cpu=DEVICE == "cpu")
 inter_nthreads, intra_nthreads = get_default_nthreads()
 # if inter_nthreads > 0:  # the behavior of 0 is not documented
 #     os.environ['OMP_NUM_THREADS'] = str(inter_nthreads)

deepmd/pt/utils/env.py

@@ -79,7 +79,7 @@
 DEFAULT_PRECISION = "float64"
 
 # throw warnings if threads not set
-set_default_nthreads()
+set_default_nthreads(use_cpu=DEVICE.type == "cpu")
 inter_nthreads, intra_nthreads = get_default_nthreads()
 if inter_nthreads > 0:  # the behavior of 0 is not documented
     torch.set_num_interop_threads(inter_nthreads)

deepmd/tf/env.py

@@ -263,7 +263,7 @@ def get_tf_session_config() -> Any:
     Any
         session configure object
     """
-    set_tf_default_nthreads()
+    set_tf_default_nthreads(use_cpu=os.environ.get("DEVICE") == "cpu")
     intra, inter = get_tf_default_nthreads()
     if int(os.environ.get("DP_JIT", 0)):
         set_env_if_empty("TF_XLA_FLAGS", "--tf_xla_auto_jit=2")