#include "pxr/pxr.h"
#include "pxr/base/work/threadLimits.h"
#include "pxr/base/tf/envSetting.h"
#include <tbb/atomic.h>
#include <tbb/task_scheduler_init.h>
#include <algorithm>
PXR_NAMESPACE_USING_DIRECTIVE
TF_DEFINE_ENV_SETTING(
PXR_WORK_THREAD_LIMIT, 0,
"Limits the number of threads the application may spawn. 0 (default) "
"allows for maximum concurrency as determined by the number of physical "
"cores, or the process's affinity mask, whichever is smaller. Note that "
"the environment variable (if set to a non-zero value) will override any "
"value passed to Work thread-limiting API calls.");
PXR_NAMESPACE_OPEN_SCOPE
static tbb::atomic<unsigned> _threadLimit;
static tbb::task_scheduler_init *_tbbTaskSchedInit;
unsigned
WorkGetPhysicalConcurrencyLimit()
{
return tbb::task_scheduler_init::default_num_threads();
}
static unsigned
Work_NormalizeThreadCount(const int n)
{
return n >= 0 ? n : std::max<int>(1, n + WorkGetPhysicalConcurrencyLimit());
}
static unsigned
Work_GetConcurrencyLimitSetting()
{
return Work_NormalizeThreadCount(TfGetEnvSetting(PXR_WORK_THREAD_LIMIT));
}
static unsigned
Work_OverrideConcurrencyLimit(unsigned weakValue, unsigned strongValue)
{
return strongValue ? strongValue : weakValue;
}
static void
Work_InitializeThreading()
{
const unsigned settingVal = Work_GetConcurrencyLimitSetting();
const unsigned physicalLimit = WorkGetPhysicalConcurrencyLimit();
_threadLimit = Work_OverrideConcurrencyLimit(physicalLimit, settingVal);
if (settingVal)
_tbbTaskSchedInit = new tbb::task_scheduler_init(_threadLimit);
}
static int _forceInitialization = (Work_InitializeThreading(), 0);
void
WorkSetConcurrencyLimit(unsigned n)
{
if (n) {
const unsigned settingVal = Work_GetConcurrencyLimitSetting();
_threadLimit = Work_OverrideConcurrencyLimit(n, settingVal);
}
if (_tbbTaskSchedInit) {
_tbbTaskSchedInit->terminate();
_tbbTaskSchedInit->initialize(_threadLimit);
} else {
_tbbTaskSchedInit = new tbb::task_scheduler_init(_threadLimit);
}
}
void
WorkSetMaximumConcurrencyLimit()
{
WorkSetConcurrencyLimit(WorkGetPhysicalConcurrencyLimit());
}
void
WorkSetConcurrencyLimitArgument(int n)
{
WorkSetConcurrencyLimit(Work_NormalizeThreadCount(n));
}
unsigned
WorkGetConcurrencyLimit()
{
return _threadLimit;
}
PXR_NAMESPACE_CLOSE_SCOPE