Numba

[prisae]

There are currently two branches which try to implement numba:

• https://github.com/empymod/empymod/tree/numba
• https://github.com/empymod/empymod/tree/numba-newtry

Roadmap:

• Properly implement the four kernel functions.
  • wavenumber()
  • greenfct()
  • reflections()
  • fields()-
• Further Benchmarks and Profiling
  • Test with parallel=True and nb.prange
  • https://github.com/empymod/empymod-examples/blob/master/6b_Profiling.ipynb
  • https://github.com/empymod/empymod-examples/blob/master/6a_Benchmark.ipynb
  • => Decide if this is enough, or if also parts of transform.? have to be jitted.
  • Adjusts pytest for numba
  • Ensure documentation on RTFD for numba builds properly
• Remove completely numexpr
  • from code
  • from documentation etc
  • requirements
  • examples
  • README
• In other repos:
  • Adjust website (only once v2.0 is released)
  • Adjust benchmarks (empymod-asv)

---

Below was the original commit from 2018.

Currently 13 of the heaviest computations are implemented twice, once with NumPy, once as strings for numexpr (switch opt=None or opt='parallel'). These are all in empymod.kernel, and are

• 1x sqrt of complex matrix of size (nfreq, noff, nlay, nwvnr)
• 11x exp of complex matrix of size (nfreq, noff, nwvnr)
• 1x division of complex matrices of size (nfreq, noff, nwvnr)

Maybe these two implementation should be replaced by one single implementation using numba. See the notebooks

• Complex-Sqrt_Numpy-Numexpr-Numba-Cython.ipynb
• Sqrt_Numpy-Numba.ipynb
• Exp_Numpy-Numba.ipynb

in https://github.com/prisae/tmp-share.

It looks like a speed-up of up to 70% could be achieved. At the cost of readability. But then, given the first list, 3 numba-functions might be enough; sqrt, exp, and the division.

Something to think about.

[prisae]

An implementation happens here in the branch numba2.

There is an older branch, but kind of stale: numba.

[prisae]

It is not a straight forward issue, and I think a lot more profiling would be required, to even decide whether it is worth it or not. The numba2 branch has a very basic implementation of numba, there would be lots of room for improvement. But with the current state, a very rough comparison looks like this:

Except for single offsets spline is the fastest method, but because it interpolates, so it is out of the category of this comparison. For very small model, nothing beats plain NumPy. For larger models, numexpr is the fastest, however, numba is also faster than NumPy. (And remember, numba is not yet implemented in as many calculations as numexpr, so it could catch up there.)

However, I think it boils down to how well the different methods implement parallelization. In this method, numexpr always used the 4 available cores, while numba somehow at times, but not as consequent (just from looking at the system monitor). NumPy, on the other hand, just run on 1 core.

The heavy calculation or bottlenecks boil down to sqrt and exp functions. So it might be that there is actually not much room for improvement, just better parallelization. For instance, the whole empymod.wavenumber could be run in parallel, and then all results gathered before the transforms.

So for the moment I leave it as it is, NumPy and numexpr, where the latter can be used to run it on as many cores as desired. I keep the numba-branches, if I or someone else decides to pick that up again.

[prisae]

For future reference, if I ever pick that up again:

The issue is complex. The comparison shows various interesting things:

• The implementation to calculate Gamma is pretty good in numba.
• The implementation of exp(-Gam*d) is not so good in numba.
• spline is so fast that other things (dlf, and generally just everything) becomes important. Not much to improve there without rewriting the entire empymod in C or Fortran.
• For anything but the splined version, almost the entire time of empymod is spent in the kernel (everything in and below wavenumber in the figure). Not even the Hankel transform exceeds 1% of total calculation.

The last point shows that if there is any speed-up to be achieved, the only important functions are:

• empymod.kernel.greenfct,
• empymod.kernel.reflections, and
• empymod.kernel.fields.
  These three functions do the number crunching, and it boils down to complex square roots of massive matrices (Gamma) and exponentials and operations with this Gamma. Gamma can reach millions if not billions of elements (if the memory cannot hold it than one has to loop over offsets or frequencies).

Now, two things could be done:

• Mixture of numba and numexpr (don't think it is a good idea).
• Improve numba for the exp(-Gam*d) calculations.

However, given that spline is accurate enough for any real-life calculation, it is probably not worth it at all. For the exact solution time shouldn't matter that much, or numexpr can be used. So I think empymod is just fine as it is at the moment.

[prisae]

Running the above again but using only 1 thread for numba and for numexpr results that there is (almost) no gain at all over the pure numpy implementation. So the benefit is just from parallelization.

The best way to achieve this is, however, to parallelize the whole kernel, hence the calls to kernel.wavenumber, which would have to be done in every Hankel transform method within transform, or the call to kernel.greenfct within kernel.wavenumber, which might be the better approach.

I am closing this issue. I won't implement numba. I keep numexpr in for now, because it is an easy way to run it in parallel. But ideally the whole kernel, and not just a few of the calculations, could be run in parallel.

[prisae]

After all the experience with emg3d I am sure that numba could replace numexpr, leaving only one method (no duplication numpy/numexpr). Re-opening to keep it on the list...

[prisae]

Note that 'parallel', hence the use of numexpr, really is only faster if it can use several threads. Otherwise the plain numpy is faster. Maybe just get rid of it altogether to simplify things?

[prisae]

Replacing the whole kernel with numba should be the approach...

[prisae]

As a note

The difficulty comes from the many different scenarios, which yield different things:

• E.g., if there are dozens of layers, the recursion becomes expensive, and just using numexpr or numba for a few functions does not help anything, as the whole reflection function would have to be rewritten.
• If time-domain data is wanted, then the spline version is not used, which throws another light on the above comments.
• If spline can be used, then so much stuff becomes important (DLF etc) that it is a whole other issue.

So this is a problem that depends a lot on the use case, which makes it difficult. There might be absolutely now gain for some cases, and there might be a lot of gain for other cases.

[prisae]

Will be closed by #77