numba.get_num_threads/set_num_threads resets the value of torch.get_num_threads about numba HOT 1 OPEN

Thanks for the report @itzsimpl. I've not tried to reproduce as the reported output is sufficiently detailed and it is also what I would guess would happen based on the description.

Some explanation as to what is happening... when a call is made to numba.get_num_threads() or numba.set_num_threads() if Numba hasn't launched a threading layer already, the first thing it will do is try and launch a threading layer based on availability and user preference (there are 3 to choose from). The threading layer will "read" the number of threads required indirectly from the environment variable NUMBA_NUM_THREADS and if that is not present it will typically use max(1, len(os.sched_getaffinity(0))). At this point the threading layer will launch and initialize, in the case of the OpenMP threading layer, as demonstrabily used in the example above, this will include setting the number of threads available to the OpenMP library to NUMBA_NUM_THREADS i.e. the OpenMP ICV nthreads-var ends up with the value NUMBA_NUM_THREADS. It so happens in the example above that torch is also using OpenMP and will pick up this change as a result of launching Numba's OpenMP based threading layer (they are clearly sharing the OpenMP ICV as only one OpenMP library is realistically ever going to be in process).

This explanation was discussed at the Numba maintainer's meeting on 2024-01-16 and it wasn't clear what the expectation is for torch and Numba interacting. Altering Numba's OpenMP threading layer to check the current thread-pool size, compare that to NUMBA_NUM_THREADS or sched_getaffinity[0] and then take some action based on it would lead to a deviation in behaviour of the Numba OpenMP threading layer in comparison to the TBB and workqueue threading layers. One thing that might work would be to set NUMBA_NUM_THREADS to the same value as whatever is being used to set the number of threads in torch, if that's the desired behaviour i.e.
NUMBA_NUM_THREADS = HOW_TORCH_SETS_MAX_THREADS = <value>.

It should also be noted that the reason for Numba having its own method for setting a thread count (NUMBA_NUM_THREADS env var) is that Numba has three threading layers and for ease of use it was decided that the method of setting a thread count needed to be invariant of the layer used (e.g. NUMBA_NUM_THREADS works for all of Numba's threading layers contrary to OpenMP specific env var OMP_NUM_THREADS which has no impact on the TBB or workqueue threading layer).

I hope this explanation helps, please do share your thoughts. Many thanks.

from numba.