Skip to content

Commit 8a6b916

Browse files
committed
feat(itkPipeline): add --threads flag
Common flag to specify the threads used by the pipeline, if supported.
1 parent 41528c3 commit 8a6b916

File tree

2 files changed

+45
-0
lines changed

2 files changed

+45
-0
lines changed

include/itkPipeline.h

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,13 @@
2525
#include "itkMacro.h"
2626
#include "itkImage.h"
2727
#include "itkVectorImage.h"
28+
#include "itkConfigure.h"
29+
#if defined(ITK_USE_PTHREADS) || defined(ITK_USE_WIN32_THREADS)
30+
#define ITK_WASM_PIPELINE_USE_THREADS 1
31+
#endif
32+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
33+
#include "itkMultiThreaderBase.h"
34+
#endif
2835

2936
#include "glaze/glaze.hpp"
3037

@@ -164,6 +171,16 @@ class WebAssemblyInterface_EXPORT Pipeline : public CLI::App
164171
return m_UseMemoryIO;
165172
}
166173

174+
int
175+
get_threads() const
176+
{
177+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
178+
return m_Threads;
179+
#else
180+
return 1;
181+
#endif
182+
}
183+
167184
int
168185
get_argc() const
169186
{
@@ -198,6 +215,9 @@ class WebAssemblyInterface_EXPORT Pipeline : public CLI::App
198215
int m_argc;
199216
char ** m_argv;
200217
std::string m_Version;
218+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
219+
int m_Threads;
220+
#endif
201221
};
202222

203223

src/itkPipeline.cxx

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ Pipeline ::Pipeline(std::string name, std::string description, int argc, char **
3131
, m_argc(argc)
3232
, m_argv(argv)
3333
, m_Version("0.1.0")
34+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
35+
, m_Threads(8)
36+
#endif
3437
{
3538
this->footer("Enjoy ITK!");
3639

@@ -39,17 +42,39 @@ Pipeline ::Pipeline(std::string name, std::string description, int argc, char **
3942
this->add_flag("--memory-io", m_UseMemoryIO, "Use itk-wasm memory IO")->group("");
4043
this->set_version_flag("--version", m_Version);
4144

45+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
46+
this->add_option("--threads", m_Threads, "Number of threads to use for processing")->group("");
47+
#endif
48+
4249
// Set m_UseMemoryIO before it is used by other memory parsers
4350
this->preparse_callback([this](size_t arg) {
4451
m_UseMemoryIO = false;
52+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
53+
m_Threads = 8; // Reset to default
54+
#endif
4555
for (int ii = 0; ii < this->m_argc; ++ii)
4656
{
4757
const std::string arg(this->m_argv[ii]);
4858
if (arg == "--memory-io")
4959
{
5060
m_UseMemoryIO = true;
5161
}
62+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
63+
if (arg == "--threads" && ii + 1 < this->m_argc)
64+
{
65+
try {
66+
m_Threads = std::stoi(this->m_argv[ii + 1]);
67+
} catch (...) {
68+
// Keep default value if parsing fails
69+
}
70+
}
71+
#endif
5272
}
73+
#ifdef ITK_WASM_PIPELINE_USE_THREADS
74+
// Set ITK's global number of threads
75+
itk::MultiThreaderBase::SetGlobalDefaultNumberOfThreads(m_Threads);
76+
std::cout << "m_Threads: " << m_Threads << std::endl;
77+
#endif
5378
});
5479

5580
#ifndef ITK_WASM_NO_FILESYSTEM_IO

0 commit comments

Comments
 (0)