Repository: pytorch
Issue: Pytorch openmp thread number tuning option for CPU trainning
Labels: module: performance,module: cpu,triaged,enhancement,module: multithreading

## 🚀 Feature
Allow an option to user so that user can choose the openmp thread either be 1 or max thread number.
## Motivation
We found that the tuning of paralel_for() makes the openmp thread number changing across the layers and we suspect this is the reason why it causes the performance difference. On the optimized DLRM version, , we have found that the performance runs best when all the operations uses max thread number.  

In the CPU trainning scenario where we assume all cpus dedicated for training, changing the thread number for some ops may hurt the locality, and also the transient low thread number may not really help the CPU utilization because usually we want the working thread spining a while to expect something coming soon.

## Additional context

We reported DLRM perf regression (https://github.com/pytorch/pytorch/issues/28198) due to PR https://github.com/pytorch/pytorch/pull/26963, which choose number of openmp threads for parallel_for according to the grain size and task size. When the task size is X times bigger than grain size, the number of openmp threads is set to Min(X, max thread number). The max thread number is specified by the user OMP_NUM_THREAD (or default as physical core#).  Priori to the change, the number of openmp threads is either 1 or max thread number.

To address the regression, a further tuning was made on PR28770 (https://github.com/pytorch/pytorch/pull/28770). The PR identified THFloatTensor_fill (used by cpu_zero) not using all threads. It reported the THxxx uses TH_OMP_OVERHEAD_THRESHOLD 100K which is about 3x higher than at::internal::GRAIN_SIZE.  The PR changed TH_OMP_OVERHEAD_THRESHOLD to be the same as at::internal::GRAIN_SIZE and reported the performance is largely back (back from 3505.85ms/iter to 3453.23ms/iter), since it allows more threads being used for THFloatTensor_fill. 

But the tuning doesn't solve all the regression. When we further optimize the DLRM to the level of ~36ms/iter (yes ~100x better than original version on SKX 8180), we observed that a changing number of openmp threads among different layers hurt performance.   The tuning of openmp threads for 
at::parallel_for() impacts lots of operations, including these implemented on
TensorIterator::for_each() and cpu_kernel_vec() since these build on top of at::parallel_for(). For example,  index_put() uses TensorIterator::for_each() and relu() uses cpu_kernel_vec().  

We found that the tuning of paralel_for() makes the openmp thread number changing across the layers and we suspect this is the reason why it causes the performance difference. On the optimized DLRM version, we have found that the performance runs best when all the operations uses max thread number. As an experiment, we change the cpu_kernel_vec() and cpu_index_kernel() interface so that it effectively launch max thread number (PR 29705 https://github.com/pytorch/pytorch/pull/29705),  and we found the performance regression is fully back (38.25 ms/iter vs. 36.44 ms/iter). This is consistent to our performance tuning BKM for using openmp in training scenario. 

With various tuning, we doubt whether it makes sense for Pytorch to tune the openmp thread number per the grain size in the trainning scenario where we assume all cpus dedicated for training. Changing the thread number for some ops may hurt the locality, and also the transient low thread number may not really help the CPU utilization because usually we want the free working thread spining a while to expect something coming soon. Maybe we allow an option to user so that user can choose the openmp thread either be 1 or max thread number. 



cc @VitalyFedyunin @ngimel @mruberry