Add single-threaded njit Xphi kernel for single-process path#26
Open
sbwells22 wants to merge 1 commit into
Open
Add single-threaded njit Xphi kernel for single-process path#26sbwells22 wants to merge 1 commit into
sbwells22 wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The single_process=True fallback routes the Xphi computation to compute_Xphi_data_numpy, which carries no numba decorator. That leaves both ways of running multiple trials compromised: running parallel-numba trials concurrently oversubscribes cores (each worker spins up an all-core threadpool), while the single_process alternative avoids the oversubscription only by dropping JIT on the hottest kernel. This PR adds compute_Xphi_data_serial. which is the compute_Xphi_data body with prange -> range, and points the single_process branch in _fit at it. compute_Xphi_data_numpy is left in place as the reference. It's additive and no existing function renamed or changed, no public API touched. The serial kernel is bit-identical to the parallel one with the side benefit that single_process=True now matches single_process=False exactly, where they previously diverged slightly (the numpy fallback derives e_logx and normalizes the softmax differently).
Scope—basically, what this does not fix. This does not address the oversubscription in scHPF_consensus.py. With -j N, consensus launches N concurrent scHPF train subprocesses via subprocess.Popen, and scHPF train calls run_trials with single_process=False where each child spins up its own all-core numba threadpool and you get ~N×cores threads on cores cores, independent of this change. The serial kernel only affects the single_process=True path, which consensus never takes. Fixing that is a separate change, such as capping NUMBA_NUM_THREADS per child, or routing consensus through train-pool/run_trials_pool (which uses single_process=True and would then pick up this kernel).