feat: allow adjustment of compiler threads (#99)

* feat: allow adjustment of compiler threads

Compiler threads may be changed in the `Advanced->Compiler Threads`
menu. Defaults to the available cores - 1, but may be adjusted from 1 to
maximum cores. Compiler threads also have been changed to `background`
priority to prevent lock ups on start.

* fix: fix ctd when terminating compiler threads

Switch to BS:thread_pool to abstract thread management and avoid
overhead of killing and starting threads. Instead, only add jobs to
pool when the number of active and queued threads is less than the
limit. The reason it's total and not just active is to avoid the case
where all tasks are dumped into the thread pool since the thread pool is
at the hardware maximum.

During compilation, one manager job is spawned to handle task
allocation and the remaining threads do the compilation work.

src/Menu.cpp

@@ -14,6 +14,7 @@
 #include "Features/WaterBlending.h"
 
 #define SETTING_MENU_TOGGLEKEY "Toggle Key"
+#define SETTING_MENU_SKIPKEY "Skip Compilation Key"
 #define SETTING_MENU_FONTSCALE "Font Scale"
 
 void SetupImGuiStyle()
@@ -77,6 +78,9 @@ void Menu::Load(json& o_json)
 	if (o_json[SETTING_MENU_TOGGLEKEY].is_number_unsigned()) {
 		toggleKey = o_json[SETTING_MENU_TOGGLEKEY];
 	}
+	if (o_json[SETTING_MENU_SKIPKEY].is_number_unsigned()) {
+		skipCompilationKey = o_json[SETTING_MENU_SKIPKEY];
+	}
 	if (o_json[SETTING_MENU_FONTSCALE].is_number_float()) {
 		fontScale = o_json[SETTING_MENU_FONTSCALE];
 	}
@@ -86,6 +90,7 @@ void Menu::Save(json& o_json)
 {
 	json menu;
 	menu[SETTING_MENU_TOGGLEKEY] = toggleKey;
+	menu[SETTING_MENU_SKIPKEY] = skipCompilationKey;
 	menu[SETTING_MENU_FONTSCALE] = fontScale;
 
 	o_json["Menu"] = menu;
@@ -211,11 +216,18 @@ RE::BSEventNotifyControl Menu::ProcessEvent(RE::InputEvent* const* a_event, RE::
 			switch (button->device.get()) {
 			case RE::INPUT_DEVICE::kKeyboard:
 				if (!button->IsPressed()) {
+					logger::trace("Detected key code {} ({})", KeyIdToString(key), key);
 					if (settingToggleKey) {
 						toggleKey = key;
 						settingToggleKey = false;
+					} else if (settingSkipCompilationKey) {
+						skipCompilationKey = key;
+						settingSkipCompilationKey = false;
 					} else if (key == toggleKey) {
 						IsEnabled = !IsEnabled;
+					} else if (key == skipCompilationKey) {
+						auto& shaderCache = SIE::ShaderCache::Instance();
+						shaderCache.backgroundCompilation = true;
 					}
 				}
 
@@ -412,10 +424,25 @@ void Menu::DrawSettings()
 
 				ImGui::AlignTextToFramePadding();
 				ImGui::SameLine();
-				if (ImGui::Button("Change")) {
+				if (ImGui::Button("Change##toggle")) {
 					settingToggleKey = true;
 				}
 			}
+			if (settingSkipCompilationKey) {
+				ImGui::Text("Press any key to set as Skip Compilation Key...");
+			} else {
+				ImGui::AlignTextToFramePadding();
+				ImGui::Text("Skip Compilation Key:");
+				ImGui::SameLine();
+				ImGui::AlignTextToFramePadding();
+				ImGui::TextColored(ImVec4(1, 1, 0, 1), "%s", KeyIdToString(skipCompilationKey));
+
+				ImGui::AlignTextToFramePadding();
+				ImGui::SameLine();
+				if (ImGui::Button("Change##skip")) {
+					settingSkipCompilationKey = true;
+				}
+			}
 
 			if (ImGui::SliderFloat("Font Scale", &fontScale, -2.f, 2.f, "%.2f")) {
 				float trueScale = exp2(fontScale);
@@ -479,9 +506,19 @@ void Menu::DrawSettings()
 				ImGui::BeginTooltip();
 				ImGui::PushTextWrapPos(ImGui::GetFontSize() * 35.0f);
 				ImGui::Text(
-					"Number of threads to compile shaders with. "
-					"The more threads the faster compilation will finish but may make the system unresponsive. "
-					"This should only be changed between restarts. ");
+					"Number of threads to use to compile shaders. "
+					"The more threads the faster compilation will finish but may make the system unresponsive. ");
+				ImGui::PopTextWrapPos();
+				ImGui::EndTooltip();
+			}
+			ImGui::SliderInt("Background Compiler Threads", &shaderCache.backgroundCompilationThreadCount, 1, static_cast<int32_t>(std::thread::hardware_concurrency()));
+			if (ImGui::IsItemHovered()) {
+				ImGui::BeginTooltip();
+				ImGui::PushTextWrapPos(ImGui::GetFontSize() * 35.0f);
+				ImGui::Text(
+					"Number of threads to use to compile shaders while playing game. "
+					"This is activated if the startup compilation is skipped. "
+					"The more threads the faster compilation will finish but may make the system unresponsive. ");
 				ImGui::PopTextWrapPos();
 				ImGui::EndTooltip();
 			}
@@ -575,7 +612,9 @@ void Menu::DrawOverlay()
 
 	auto failed = shaderCache.GetFailedTasks();
 	auto hide = shaderCache.IsHideErrors();
-	auto progressTitle = fmt::format("Compiling Shaders: {}", shaderCache.GetShaderStatsString(!state->IsDeveloperMode()).c_str());
+	auto progressTitle = fmt::format("{}Compiling Shaders: {}",
+		shaderCache.backgroundCompilation ? "Background " : "",
+		shaderCache.GetShaderStatsString(!state->IsDeveloperMode()).c_str());
 	auto percent = (float)compiledShaders / (float)totalShaders;
 	auto progressOverlay = fmt::format("{}/{} ({:2.1f}%)", compiledShaders, totalShaders, 100 * percent);
 	if (shaderCache.IsCompiling()) {
@@ -587,6 +626,13 @@ void Menu::DrawOverlay()
 		}
 		ImGui::TextUnformatted(progressTitle.c_str());
 		ImGui::ProgressBar(percent, ImVec2(0.0f, 0.0f), progressOverlay.c_str());
+		if (!shaderCache.backgroundCompilation && shaderCache.menuLoaded) {
+			auto skipShadersText = fmt::format(
+				"Press {} to proceed without completing shader compilation. "
+				"WARNING: Uncompiled shaders will have visual errors or cause stuttering when loading.",
+				KeyIdToString(skipCompilationKey));
+			ImGui::TextUnformatted(skipShadersText.c_str());
+		}
 
 		ImGui::End();
 	} else if (failed && !hide) {

src/Menu.h

@@ -27,8 +27,9 @@ class Menu : public RE::BSTEventSink<RE::InputEvent*>
 
 private:
 	uint32_t toggleKey = VK_END;
+	uint32_t skipCompilationKey = VK_ESCAPE;
 	bool settingToggleKey = false;
-
+	bool settingSkipCompilationKey = false;
 	float fontScale = 0.f;  // exponential
 
 	Menu() {}

src/ShaderCache.cpp

@@ -4,6 +4,7 @@
 
 #include <d3d11.h>
 #include <d3dcompiler.h>
+#include <fmt/std.h>
 #include <wrl/client.h>
 
 #include "Features/ExtendedMaterials.h"
@@ -1386,6 +1387,27 @@ namespace SIE
 		Clear();
 	}
 
+	void ShaderCache::AdjustThreadCount()
+	{
+		auto size = compilationThreads.size();
+		if (size == compilationThreadCount)
+			return;
+		if (size && std::this_thread::get_id() != compilationThreads.front().get_id())
+			// only allow first thread to adjust threads
+			return;
+		logger::debug("Adjusting active threads {}/{}", (int)size, (int)compilationThreadCount);
+		if (size && size > compilationThreadCount) {
+			auto& thread = compilationThreads.back();
+			logger::debug("Stopping thread {}: active {}/{}", thread.get_id(), (int)size - 1, (int)compilationThreadCount);
+			thread.request_stop();
+			compilationThreads.pop_back();
+		} else if (size < compilationThreadCount) {
+			compilationThreads.push_back(std::jthread(&ShaderCache::ProcessCompilationSet, this, ssource.get_token()));
+			auto& thread = compilationThreads.back();
+			logger::debug("Starting new thread {}: active {}/{}", thread.get_id(), (int)size + 1, (int)compilationThreadCount);
+		}
+	}
+
 	void ShaderCache::Clear()
 	{
 		for (auto& shaders : vertexShaders) {
@@ -1401,6 +1423,7 @@ namespace SIE
 			shaders.clear();
 		}
 
+		ssource.request_stop();
 		compilationSet.Clear();
 		std::unique_lock lock{ mapMutex };
 		shaderMap.clear();
@@ -1546,10 +1569,8 @@ namespace SIE
 
 	ShaderCache::ShaderCache()
 	{
-		logger::debug("ShaderCache initialized with {} compiler threads", compilationThreadCount);
-		for (size_t threadIndex = 0; threadIndex < compilationThreadCount; ++threadIndex) {
-			compilationThreads.push_back(std::jthread(&ShaderCache::ProcessCompilationSet, this));
-		}
+		logger::debug("ShaderCache initialized with {} compiler threads", (int)compilationThreadCount);
+		AdjustThreadCount();
 	}
 
 	RE::BSGraphics::VertexShader* ShaderCache::MakeAndAddVertexShader(const RE::BSShader& shader,
@@ -1642,13 +1663,16 @@ namespace SIE
 		hideError = !hideError;
 	}
 
-	void ShaderCache::ProcessCompilationSet()
+	void ShaderCache::ProcessCompilationSet(std::stop_token stoken)
 	{
-		SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_NORMAL);
-		while (true) {
-			const auto& task = compilationSet.WaitTake();
-			task.Perform();
-			compilationSet.Complete(task);
+		SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_BELOW_NORMAL);
+		while (!stoken.stop_requested()) {
+			const auto& task = compilationSet.WaitTake(stoken);
+			if (!task.has_value())
+				break;  // exit because thread told to end
+			task.value().Perform();
+			compilationSet.Complete(task.value());
+			AdjustThreadCount();
 		}
 	}
 
@@ -1684,15 +1708,23 @@ namespace SIE
 		return GetId() == other.GetId();
 	}
 
-	ShaderCompilationTask CompilationSet::WaitTake()
+	std::optional<ShaderCompilationTask> CompilationSet::WaitTake(std::stop_token stoken)
 	{
 		std::unique_lock lock(compilationMutex);
-		conditionVariable.wait(lock, [this]() { return !availableTasks.empty(); });
+		if (!conditionVariable.wait(
+				lock, stoken,
+				[this, &shaderCache]() { return !availableTasks.empty() &&
+			                                    // check against all tasks in queue to trickle the work. It cannot be the active tasks count because the thread pool itself is maximum.
+			                                    (int)shaderCache.compilationPool.get_tasks_total() <=
+			                                        (!shaderCache.backgroundCompilation ? shaderCache.compilationThreadCount : shaderCache.backgroundCompilationThreadCount); })) {
+			/*Woke up because of a stop request. */
+			return std::nullopt;
+		}
 		if (!ShaderCache::Instance().IsCompiling()) {  // we just got woken up because there's a task, start clock
 			lastCalculation = lastReset = high_resolution_clock::now();
 		}
 		auto node = availableTasks.extract(availableTasks.begin());
-		auto task = node.value();
+		auto& task = node.value();
 		tasksInProgress.insert(std::move(node));
 		return task;
 	}

src/ShaderCache.h

@@ -60,7 +60,7 @@ namespace SIE
 	class CompilationSet
 	{
 	public:
-		ShaderCompilationTask WaitTake();
+		std::optional<ShaderCompilationTask> WaitTake(std::stop_token stoken);
 		void Add(const ShaderCompilationTask& task);
 		void Complete(const ShaderCompilationTask& task);
 		void Clear();
@@ -77,7 +77,7 @@ namespace SIE
 		std::unordered_set<ShaderCompilationTask> availableTasks;
 		std::unordered_set<ShaderCompilationTask> tasksInProgress;
 		std::unordered_set<ShaderCompilationTask> processedTasks;  // completed or failed
-		std::condition_variable conditionVariable;
+		std::condition_variable_any conditionVariable;
 		std::chrono::steady_clock::time_point lastReset = high_resolution_clock::now();
 		std::chrono::steady_clock::time_point lastCalculation = high_resolution_clock::now();
 		double totalMs = (double)duration_cast<std::chrono::milliseconds>(lastReset - lastReset).count();
@@ -124,7 +124,11 @@ namespace SIE
 		void DeleteDiskCache();
 		void ValidateDiskCache();
 		void WriteDiskCacheInfo();
-
+		/// <summary>
+		/// Adjust the compiler threads based on the compileThreadCount.
+		/// </summary>
+		/// This will terminate or generate threads as required to match compileThreadCount.
+		void AdjustThreadCount();
 		void Clear();
 
 		bool AddCompletedShader(ShaderClass shaderClass, const RE::BSShader& shader, uint32_t descriptor, ID3DBlob* a_blob);
@@ -152,10 +156,14 @@ namespace SIE
 		bool IsHideErrors();
 
 		int32_t compilationThreadCount = std::max(static_cast<int32_t>(std::thread::hardware_concurrency()) - 1, 1);
+		int32_t backgroundCompilationThreadCount = std::max(static_cast<int32_t>(std::thread::hardware_concurrency()) / 2, 1);
+		BS::thread_pool compilationPool{};
+		bool backgroundCompilation = false;
+		bool menuLoaded = false;
 
 	private:
 		ShaderCache();
-		void ProcessCompilationSet();
+		void ProcessCompilationSet(std::stop_token stoken);
 
 		~ShaderCache();
 
@@ -173,6 +181,7 @@ namespace SIE
 		bool hideError = false;
 
 		eastl::vector<std::jthread> compilationThreads;
+		std::stop_source ssource;
 		std::mutex vertexShadersMutex;
 		std::mutex pixelShadersMutex;
 		CompilationSet compilationSet;

src/State.cpp

@@ -92,6 +92,8 @@ void State::Load()
 			SetDefines(advanced["Shader Defines"]);
 		if (advanced["Compiler Threads"].is_number_integer())
 			shaderCache.compilationThreadCount = std::clamp(advanced["Compiler Threads"].get<int32_t>(), 1, static_cast<int32_t>(std::thread::hardware_concurrency()));
+		if (advanced["Background Compiler Threads"].is_number_integer())
+			shaderCache.backgroundCompilationThreadCount = std::clamp(advanced["Compiler Threads"].get<int32_t>(), 1, static_cast<int32_t>(std::thread::hardware_concurrency()));
 	}
 
 	if (settings["General"].is_object()) {
@@ -145,6 +147,7 @@ void State::Save()
 	advanced["Log Level"] = logLevel;
 	advanced["Shader Defines"] = shaderDefinesString;
 	advanced["Compiler Threads"] = shaderCache.compilationThreadCount;
+	advanced["Background Compiler Threads"] = shaderCache.backgroundCompilationThreadCount;
 	settings["Advanced"] = advanced;
 
 	json general;

src/XSEPlugin.cpp

@@ -115,8 +115,8 @@ void MessageHandler(SKSE::MessagingInterface::Message* message)
 				RE::BSInputDeviceManager::GetSingleton()->AddEventSink(Menu::GetSingleton());
 
 				auto& shaderCache = SIE::ShaderCache::Instance();
-
-				while (shaderCache.IsCompiling()) {
+				shaderCache.menuLoaded = true;
+				while (shaderCache.IsCompiling() && !shaderCache.backgroundCompilation) {
 					std::this_thread::sleep_for(100ms);
 				}