Avoid undefined behavior in Eigen::TensorCostModel::numThreads.

If the cost is large enough then the thread count can be larger than the maximum
representable int, so just casting it to an int is undefined behavior.

Contributed by phurst@google.com.
This commit is contained in:
Rasmus Munk Larsen 2017-09-08 15:49:55 -07:00
parent 6d42309f13
commit 94e2213b38

View File

@ -174,8 +174,9 @@ class TensorCostModel {
static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int numThreads(
double output_size, const TensorOpCost& cost_per_coeff, int max_threads) {
double cost = totalCost(output_size, cost_per_coeff);
int threads = (cost - kStartupCycles) / kPerThreadCycles + 0.9;
return numext::mini(max_threads, numext::maxi(1, threads));
// Make sure we don't invoke undefined behavior when we convert to an int.
threads = numext::mini<double>(threads, GenericNumTraits<int>::highest());
return numext::mini(max_threads, numext::maxi<int>(1, threads));
}
// taskSize assesses parallel task size.