From 6ff7050d2d3b7718612657bb7fe7c04b30fb40e2 Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Tue, 3 Oct 2023 19:38:22 -0400
Subject: [PATCH] enforce that the number of threads is properly reset for /omp
 styles

---
 src/OPENMP/thr_omp.cpp         | 7 +++++++
 tools/lammps-gui/lammpsgui.cpp | 8 ++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/OPENMP/thr_omp.cpp b/src/OPENMP/thr_omp.cpp
index c27a99028d..88344a900d 100644
--- a/src/OPENMP/thr_omp.cpp
+++ b/src/OPENMP/thr_omp.cpp
@@ -34,6 +34,10 @@
 
 #include <cstring>
 
+#if defined(_OPENMP)
+#include <omp.h>
+#endif
+
 using namespace LAMMPS_NS;
 using MathConst::THIRD;
 
@@ -44,6 +48,9 @@ ThrOMP::ThrOMP(LAMMPS *ptr, int style) : lmp(ptr), fix(nullptr), thr_style(style
   // register fix omp with this class
   fix = static_cast<FixOMP *>(lmp->modify->get_fix_by_id("package_omp"));
   if (!fix) lmp->error->all(FLERR, "The 'package omp' command is required for /omp styles");
+#if defined(_OPENMP)
+  omp_set_num_threads(lmp->comm->nthreads);
+#endif
 }
 
 // clang-format off
diff --git a/tools/lammps-gui/lammpsgui.cpp b/tools/lammps-gui/lammpsgui.cpp
index 6995023593..b6b590e855 100644
--- a/tools/lammps-gui/lammpsgui.cpp
+++ b/tools/lammps-gui/lammpsgui.cpp
@@ -1329,9 +1329,9 @@ void LammpsGui::preferences()
         // must delete LAMMPS instance after preferences have changed that require
         // using different command line flags when creating the LAMMPS instance like
         // suffixes or package commands
+        int newthreads = settings.value("nthreads", 1).toInt();
         if ((oldaccel != settings.value("accelerator", AcceleratorTab::None).toInt()) ||
-            (oldthreads != settings.value("nthreads", 1).toInt()) ||
-            (oldecho != settings.value("echo", false).toBool()) ||
+            (oldthreads != newthreads) || (oldecho != settings.value("echo", false).toBool()) ||
             (oldcite != settings.value("cite", false).toBool())) {
             if (lammps.is_running()) {
                 stop_run();
@@ -1340,6 +1340,10 @@ void LammpsGui::preferences()
             }
             lammps.close();
             lammpsstatus->hide();
+#if defined(_OPENMP)
+            qputenv("OMP_NUM_THREADS", std::to_string(newthreads).c_str());
+            omp_set_num_threads(newthreads);
+#endif
         }
         if (imagewindow) imagewindow->createImage();
         settings.beginGroup("reformat");