experiments
/
STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Adds a MaxParallelism() method returning a suggested maximum parallelism for a specific numa node. 

PiperOrigin-RevId: 254191552
This commit is contained in:
A. Unique TensorFlower 2019-06-20 07:25:00 -07:00 committed by TensorFlower Gardener
parent 4b59a16a2b
commit 92511b345c
3 changed files with 16 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/core/common_runtime/local_device.cc
View File

@ -95,13 +95,7 @@ struct LocalDevice::EigenThreadPoolInfo {
intra_op_parallelism_threads = env_num_threads; intra_op_parallelism_threads = env_num_threads;
// If no session setting or environment, compute a reasonable default. // If no session setting or environment, compute a reasonable default.
if (intra_op_parallelism_threads == 0) { if (intra_op_parallelism_threads == 0) {
intra_op_parallelism_threads = port::NumSchedulableCPUs(); intra_op_parallelism_threads = port::MaxParallelism(numa_node);
if (numa_node != port::kNUMANoAffinity) {
// Assume that CPUs are equally distributed over available NUMA nodes.
// This may not be true, but there isn't currently a better way of
// determining the number of CPUs specific to the requested node.
intra_op_parallelism_threads /= port::NUMANumNodes();
}
} }
} }
ThreadOptions thread_opts; ThreadOptions thread_opts;

5
tensorflow/core/platform/cpu_info.h
View File

@ -46,6 +46,11 @@ int NumSchedulableCPUs();
// called during initialization, i.e., before before main() has started. // called during initialization, i.e., before before main() has started.
int MaxParallelism(); int MaxParallelism();
// Returns an estimate for the maximum parallelism for this process on the
// provided numa node, or any numa node if `numa_node` is kNUMANoAffinity.
// See MaxParallelism() for more information.
int MaxParallelism(int numa_node);
// Returns the total number of CPUs on the system. This number should // Returns the total number of CPUs on the system. This number should
// not change even if the underlying cluster management software may // not change even if the underlying cluster management software may
// change the number of schedulable CPUs. Unlike `NumSchedulableCPUs`, if the // change the number of schedulable CPUs. Unlike `NumSchedulableCPUs`, if the

10
tensorflow/core/platform/posix/port.cc
View File

@ -82,6 +82,16 @@ int NumSchedulableCPUs() {
int MaxParallelism() { return NumSchedulableCPUs(); } int MaxParallelism() { return NumSchedulableCPUs(); }
int MaxParallelism(int numa_node) {
if (numa_node != port::kNUMANoAffinity) {
// Assume that CPUs are equally distributed over available NUMA nodes.
// This may not be true, but there isn't currently a better way of
// determining the number of CPUs specific to the requested node.
return NumSchedulableCPUs() / port::NUMANumNodes();
}
return NumSchedulableCPUs();
}
int NumTotalCPUs() { int NumTotalCPUs() {
int count = absl::base_internal::NumCPUs(); int count = absl::base_internal::NumCPUs();
return (count <= 0) ? kUnknownCPU : count; return (count <= 0) ? kUnknownCPU : count;