Parallel numba rime #186
Parallel numba rime #186
Conversation
|
@JSKenyon Could you take a quick look at this? I'd appreciate general feedback on the approach. Briefly, the user can instruct codex-africanus to create a parallel implementation for numba RIME functions using the donfig configuration package. { 'rime.feed_rotation.parallel': True }
# or
{
'rime.feed_rotation.parallel': {
# number of threads
'threads': 2,
# prange axes
'axes': ['source', 'row'],
}
}
# or
from africanus.config import config
with config.set({"rime.feed_rotation.parallel": True}):
from africanus.rime import feed_rotation
...One can place the above config in a YAML file or even set options via the command line $ AFRICANUS_RIME__FEED_ROTATION__PARALLEL="{'threads':2}" python script.pyOne issue is brittleness around imports (3rd python approach above): The config import and set needs to occur before the feed_rotation import otherwise the module import code may not be given the appropriate config. |
|
That all looks really good/cool to me. Are you intending to nest pranges? Just wondering based on the ability to specify two axes in the above example. One thing which might be useful is the ability to set this in a more global fashion if required. What I mean is that it might get tedious to specify this manually per term. Having some sensible default for the axes and the ability to just throw nthreads at the problem would be handy. |
I thought I'd try to support it, although it may not be possible in all cases. I know that in order to collapse two loops in OpenMP, there can't be any code between them and the same may apply in the case of numba prange.
Yeah, I guess a Any thoughts on the use of |
|
I am not sure if I have understood it in sufficient detail, but I believe that by default |
…us into parallel-numba-rime
|
@JSKenyon. In general I've set up the RIME function generated_jit's as follows: @generated_jit(nopython=True, nogil=True, cache=not parallel, parallel=parallel)
def fn(...):
passThere's something a bit off in the There are a couple of scratch numpy arrays allocated outside the prange but used inside. Moving them inside doesn't fix the above recursion error. If I track it down with pdb, it looks like its getting very confused by the |
|
To invoke the parallelism in this PR, the following structure can be used: from africanus.config import config
cfg= {
'rime.feed_rotation.parallel': True,
'rime.predict_vis.parallel':True,
'rime.phase_delay.parallel':True,
'model.spectral_model.parallel':True,
'model.shape.gaussian.parallel':True
}
with config.set(cfg):
# All imports **must** happen after the config.set
from africanus.rime.dask import phase_delay, feed_rotation, predict_vis
from africanus.model.dask import spectral_model
from africanus.model.shape import gaussianI've typed the above from memory, so some of it may be incorrect, but the general idea holds. |
|
I have played with this PR briefly and the components I have used seem to work. It is worth noting that the parallelism is still bound by certain non-Numba tasks such as |
Thanks for trying this out.
I take it this is in the nested parallelism case where there are a mix of dask and numba threads?
This PR still needs some work:
I hope to find time for this later this week. |
Tests added / passed
If the pep8 tests fail, the quickest way to correct
this is to run
autopep8and thenflake8andpycodestyleto fix the remaining issues.Fully documented, including
HISTORY.rstfor all changesand one of the
docs/*-api.rstfiles for new APITo build the docs locally: