Follow the maximum number of threads limit in gemma.cpp

ufownl · Jun 15, 2024 · 88bcfb3 · 88bcfb3
1 parent e988a30
commit 88bcfb3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/scheduler.cpp b/src/scheduler.cpp
@@ -1,5 +1,6 @@
 #include "scheduler.hpp"
 #include <util/app.h>
+#include <algorithm>
 #include <thread>
 
 namespace {
@@ -21,7 +22,12 @@ int destroy(lua_State* L) {
 namespace cgemma {
 
 scheduler::scheduler()
-  : pool_(std::thread::hardware_concurrency()) {
+  : pool_(std::min(static_cast<size_t>(std::thread::hardware_concurrency()), gcpp::kMaxThreads)) {
+  // nop
+}
+
+scheduler::scheduler(size_t num_threads)
+  : pool_(std::min(num_threads, gcpp::kMaxThreads)) {
   // nop
 }
 

diff --git a/src/scheduler.hpp b/src/scheduler.hpp
@@ -9,7 +9,7 @@ namespace cgemma {
 class scheduler {
 public:
   scheduler();
-  explicit scheduler(size_t num_threads) : pool_(num_threads) {}
+  explicit scheduler(size_t num_threads);
 
   hwy::ThreadPool& pool() { return pool_; }