Feature Proposal: Numba parallel permutation kernel and JIT caching in crand.py

[samay2504]

Currently, every new Python session recompiles the Numba JIT functions in crand.py from scratch, which takes around 80 seconds on a typical machine. On top of that, the inner computation loop in compute_chunk processes all observations one at a time, even though the work across observations is completely independent and could run in parallel.

This proposal is to add cache=True to the existing @njit decorators (vec_permutations and compute_chunk), which brings the cold start cost down to under 2 milliseconds after the first run. It also introduces a lightweight helper function called _wloc_offsets that pre-computes per-observation offsets into the flat CSR weights array using np.cumsum, removing the sequential accumulator that was blocking parallelism. A new compute_chunk_parallel function using @njit(parallel=True) and prange then distributes the per-observation work across all available CPU threads.

Benchmarked on a Ryzen 9 5900HX with 16 threads, this reduces the inner kernel time from 88.8ms to 19.3ms, a 4.59x speedup, with full end-to-end wall time for Moran_Local at n=2500 and 199 permutations sitting at a median of 28.7ms. At n=10000 it comes in at 127ms.

The changes are confined entirely to crand.py and pyproject.toml. No public API is touched, no other files change, and the n_jobs > 1 joblib path is completely unmodified. A proof-of-concept implementation with tests is available at samay2504:feat/numba-parallel-crand.

[sjsrey]

While the proposal is interesting and the reported benchmarks are promising, it is difficult to comment definitively without reviewing the actual code changes. In particular, a few concerns come to mind: we would need explicit confirmation about the behavior under nested parallelism to ensure there is no interference or CPU oversubscription when combined with the existing joblib-based n_jobs > 1 path. It would also be important to add a regression test verifying numerical equivalence between the serial and parallel implementations, as well as a test confirming that enabling cache=True does not alter computational results. Finally, a brief developer note should document that the crand kernels are now JIT-cached and that inner parallelism is enabled in the single-process execution path.

[samay2504]

Thanks for the feedback @sjsrey , All four concerns have been addressed directly in the updated diff. Nested parallelism is handled by making the two execution paths mutually exclusive: n_jobs=1 routes through the new compute_chunk_parallel kernel which uses Numba's own thread pool via prange, while n_jobs>1 continues to use joblib with inner_max_num_threads=1 passed to each worker, explicitly suppressing any Numba threading inside the subprocesses so there's no CPU oversubscription. A regression test test_njobs1_prange_equals_njobs2_joblib confirms numerical equivalence between the two paths by running Moran_Local with the same seed under both n_jobs=1 and n_jobs=2 and asserting the p_sim arrays are identical. A second test test_cache_true_determinism runs two back-to-back calls with the same seed and verifies the results are bit-for-bit identical, confirming that cache=True has no effect on output. Finally, a developer note block has been added directly after __all__ in crand.py documenting the JIT caching behaviour, the Numba thread pool parallelism active on the single-process path, and the nested parallelism safety guarantee on the multi-process path.