From 78027ea63c4617b17c842cbb9d2bd1cca2f2d92f Mon Sep 17 00:00:00 2001
From: Daniel Wolf <dwolf@dannad.de>
Date: Thu, 11 Aug 2016 10:29:01 +0200
Subject: [PATCH] Thread count can be limited via command-line argument

---
 src/audio/voiceActivityDetection.cpp | 4 ++--
 src/audio/voiceActivityDetection.h   | 2 +-
 src/main.cpp                         | 6 ++++++
 src/phoneExtraction.cpp              | 6 +++---
 src/phoneExtraction.h                | 1 +
 5 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/audio/voiceActivityDetection.cpp b/src/audio/voiceActivityDetection.cpp
index b795ed6..c03983b 100644
--- a/src/audio/voiceActivityDetection.cpp
+++ b/src/audio/voiceActivityDetection.cpp
@@ -66,7 +66,7 @@ BoundedTimeline<void> webRtcDetectVoiceActivity(const AudioClip& audioClip, Prog
 	return activity;
 }
 
-BoundedTimeline<void> detectVoiceActivity(const AudioClip& inputAudioClip, ProgressSink& progressSink) {
+BoundedTimeline<void> detectVoiceActivity(const AudioClip& inputAudioClip, int maxThreadCount, ProgressSink& progressSink) {
 	// Prepare audio for VAD
 	const unique_ptr<AudioClip> audioClip = inputAudioClip.clone() | resample(16000) | removeDCOffset();
 
@@ -74,7 +74,7 @@ BoundedTimeline<void> detectVoiceActivity(const AudioClip& inputAudioClip, Progr
 	std::mutex activityMutex;
 
 	// Split audio into segments and perform parallel VAD
-	int segmentCount = getProcessorCoreCount();
+	const int segmentCount = maxThreadCount;
 	centiseconds audioLength = audioClip->getTruncatedRange().getLength();
 	vector<TimeRange> audioSegments;
 	for (int i = 0; i < segmentCount; ++i) {
diff --git a/src/audio/voiceActivityDetection.h b/src/audio/voiceActivityDetection.h
index 381df23..df0ffd9 100644
--- a/src/audio/voiceActivityDetection.h
+++ b/src/audio/voiceActivityDetection.h
@@ -3,4 +3,4 @@
 #include <BoundedTimeline.h>
 #include <ProgressBar.h>
 
-BoundedTimeline<void> detectVoiceActivity(const AudioClip& audioClip, ProgressSink& progressSink);
+BoundedTimeline<void> detectVoiceActivity(const AudioClip& audioClip, int maxThreadCount, ProgressSink& progressSink);
diff --git a/src/main.cpp b/src/main.cpp
index d7afc41..0f7045a 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -16,6 +16,7 @@
 #include "stringTools.h"
 #include <boost/range/adaptor/transformed.hpp>
 #include <boost/filesystem/fstream.hpp>
+#include "parallel.h"
 
 using std::exception;
 using std::string;
@@ -115,6 +116,7 @@ int main(int argc, char *argv[]) {
 	tclap::ValuesConstraint<logging::Level> logLevelConstraint(logLevels);
 	tclap::ValueArg<logging::Level> logLevel("", "logLevel", "The minimum log level to log", false, logging::Level::Debug, &logLevelConstraint, cmd);
 	tclap::ValueArg<string> logFileName("", "logFile", "The log file path.", false, string(), "string", cmd);
+	tclap::ValueArg<int> maxThreadCount("", "threads", "The maximum number of worker threads to use.", false, getProcessorCoreCount(), "number", cmd);
 	tclap::ValueArg<string> dialogFile("d", "dialogFile", "A file containing the text of the dialog.", false, string(), "string", cmd);
 	auto exportFormats = vector<ExportFormat>(ExportFormatConverter::get().getValues());
 	tclap::ValuesConstraint<ExportFormat> exportFormatConstraint(exportFormats);
@@ -132,6 +134,9 @@ int main(int argc, char *argv[]) {
 
 		// Parse command line
 		cmd.parse(argc, argv);
+		if (maxThreadCount.getValue() < 1) {
+			throw std::runtime_error("Thread count must be 1 or higher.");
+		}
 
 		// Set up log file
 		if (logFileName.isSet()) {
@@ -151,6 +156,7 @@ int main(int argc, char *argv[]) {
 			phones = detectPhones(
 				*createAudioClip(inputFileName.getValue()),
 				dialogFile.isSet() ? readTextFile(path(dialogFile.getValue())) : boost::optional<u32string>(),
+				maxThreadCount.getValue(),
 				progressBar);
 		}
 		std::cerr << "Done" << std::endl;
diff --git a/src/phoneExtraction.cpp b/src/phoneExtraction.cpp
index bf48b7d..a40dadd 100644
--- a/src/phoneExtraction.cpp
+++ b/src/phoneExtraction.cpp
@@ -359,6 +359,7 @@ Timeline<void> getUnknownSounds(const Timeline<void>& utterances, const Timeline
 BoundedTimeline<Phone> detectPhones(
 	const AudioClip& inputAudioClip,
 	optional<u32string> dialog,
+	int maxThreadCount,
 	ProgressSink& progressSink)
 {
 	ProgressMerger totalProgressMerger(progressSink);
@@ -371,7 +372,7 @@ BoundedTimeline<Phone> detectPhones(
 	// Split audio into utterances
 	BoundedTimeline<void> utterances;
 	try {
-		utterances = detectVoiceActivity(*audioClip, voiceActivationProgressSink);
+		utterances = detectVoiceActivity(*audioClip, maxThreadCount, voiceActivationProgressSink);
 	}
 	catch (...) {
 		std::throw_with_nested(runtime_error("Error detecting segments of speech."));
@@ -437,8 +438,7 @@ BoundedTimeline<Phone> detectPhones(
 	try {
 		// Determine how many parallel threads to use
 		int threadCount = std::min({
-			// Don't use more threads than there are CPU cores
-			getProcessorCoreCount(),
+			maxThreadCount,
 			// Don't use more threads than there are utterances to be processed
 			static_cast<int>(utterances.size()),
 			// Don't waste time creating additional threads (and decoders!) if the recording is short
diff --git a/src/phoneExtraction.h b/src/phoneExtraction.h
index f4b9f91..1f5d14d 100644
--- a/src/phoneExtraction.h
+++ b/src/phoneExtraction.h
@@ -8,4 +8,5 @@
 BoundedTimeline<Phone> detectPhones(
 	const AudioClip& audioClip,
 	boost::optional<std::u32string> dialog,
+	int maxThreadCount,
 	ProgressSink& progressSink);