[wgridder] handle a coherencies dimension

[Joshuaalbert]

Currently, as far as I know, if you want to degrid a multi-coherency image or grid to one, the only way is to run wgridder once per coherency.
However, that's wasteful. It's the same convolutional operation applied to each coherency, so all those extra kernel evaluations are redundant. I imagine there are other overheads that would also be saved by adding a coherency dimension. The best place for this dimension would be the most memory-contiguous, so in C ordering dirty would be [num_l, num_m, num_coh], and visibilities would be [num_row, num_chan, num_coh].

[landmanbester]

See also the discussion here. It's not that simple a change to make since the wgridder assumes the input image to be real valued. This is not necessarily true for the off-diagonal coherencies. Thus you either need to convert to "Stokes visibilities" before gridding or use the trick here if you want to grid the data directly

[mreineck]

Thank you for chiming in, @landmanbester! (I was on vacation.)
I'm still quite convinced that the potential gains from gridding several coherencies together will be modest. At least for actual w-gridding, the overall cost is dominated by the multiplication of kernel weights with grid values (when degridding) and adding weighted visibilities to the grid (when gridding), since this part requires more accesses to cache and/or memory, in contrast to the kernel evaluation where most operands sit in CPU registers all the time.
Note that the current implementation is, on paper, already very wasteful: the 2D grid weights that need to be applied to a visibility are identical in each of its w planes, yet the wgridder recomputes them over and over for every plane. In other words, for a kernel support of, say, 8, wgridder spends eight time as much time on weight calculation as would be strictly necessary, and it is still quite fast compared to other implementations. I'm doing it this way, because storing the weights to main memory is too slow (and potentially requires too much memory), and switching to an algorithm that computes several w planes at the same time would increase overall memory consumption as well as cache misses.

[Joshuaalbert]

Thanks @landmanbester. It's a good point. Also @mreineck I understand the kernel is cheap enough that recomputing them makes more sense compared to more memory accesses. That surely must have a limit though. We're recomputng them potentially 4x more. And for gather-mul-sum (degridding) and mul-scatter-add (gridding) wouldn't the fact that the coherencies are contiguously lined up in memory save cost in both directions? Maybe the compiler can make good use of vectorized instructions.

[mreineck]

I'm already using vector instructions very aggressively in the current code, Kernel weights are evaluated in parallel with the maximum vector length supported by the CPU, as well as application of weights to visibilities, and grid read/write operations.
Also, if on-the-fly computation is faster than the reading precomputed values from memory (yes, I know how counter-intuitive that feels!), doing it more often will not change the situation.
That said, I'm totally sure that with enough hand-tweaking you may very well get an implementation which, for specific hardware and problem configuration, is 50% faster than the vanilla wgridder. But that will then lock you to

• a specific set of CPU features (like AVX2 with FMA)
• a given precision (single or double)
• a fixed number of coherencies

All that may make perfect sense if you do number crunching for a specific experiment, where these parameters won't change for a decade or so. For a general purpose code like ducc, implementing and maintaining a set of such specialized solutions is not really feasible.

[Joshuaalbert]

@mreineck speaking of problem configurations, do you have a sense for how to optimise for very uv-filled large-N arrays, such as the DSA? Our entire UV grid (> 10m baselines) is filled after a few channels of freq synthesis and a few minutes of time synthesis. I've been wondering if there is a tweak that could take advantage of that to speed up gridding and degridding. I'd be most curious to hear your thoughts.

[mreineck]

I'm not sure I understand the described situation correctly. "A few channels of freq synthesis and a few minutes of time synthesis" sounds to me like "very few visibilities to be gridded compared to the grid size. On the other hand you mention "very uv-filled" arrays, which to me suggests the opposite.
In the former case tweaking the gridding/degridding algorithm will not help in a meaningful way, since (if you are using something like wgridder) the computations will be totally dominated by FFT operations on the uv grid. To do this efficiently, algorithms like IDG or a variant of the approach in https://arxiv.org/abs/2108.10720v2 (possibly with a wgridder-like gridding component inside) might be your best bet.
In the latter case, tuning the wgridder could indeed help. To get a feeling for that, I'd need typical numbers for vis2dirty parameters, in particular the dimensions of uvw, freq, and dirty.

[Joshuaalbert]

What I meant with 'filled after a few channels of freq synthesis' is that we have such a well sampled aperture that we only need a small amount of both types of synthesis to fill each UV grid cell with counts. The number of antennas is 1650, the number of channels is 10k, the number of time steps is 840. there are is too much data to store, so gridding and degridding is done in on sequential chunks of time and in parallel over chunks of channels. Typically uvw.shape = [5.5E6 (sometimes 21E6), 3] freq.shape = [20], and dirty.shape = [17.5k, 17.5k].

A typical UV coverage of observation for slice of 20 channels:

[mreineck]

Thanks, that was very helpful!

I have created a small benchmark you can use to explore the wgridder scalings for this kind of setup:

import ducc0
import numpy as np

SPEEDOFLIGHT = 299792458.

def run_vis2dirty(nx, ny, nrow, nchan, epsilon, singleprec, nthreads):
    use_wgt = True  # set this to False if you don't need to supply an external array of weights
    use_mask = True  # we probably need a mask
    wstacking = True  # not sure whether that's needed, probably yes
    fov = 3.  # field of view in degrees; please adjust!
    nxdirty, nydirty = nx, ny
    rng = np.random.default_rng(42)
    pixsizex = fov*np.pi/180/nxdirty
    pixsizey = fov*np.pi/180/nydirty*1.1
    f0 = 0.7e9  # 0.7 GHz
    df = 134e3  # 134 kHz, taken from web site
    freq = f0 + np.arange(nchan)*df
    # generate a totally random uvw coverage; unrealistic but should be OK for this test
    uvw = (rng.random((nrow, 3))-0.5)
    # scale u and v properly
    uvw[:,0:2] /= (pixsizex*f0/SPEEDOFLIGHT)
    # scale w. I don't know the proper factor, please adjust!
    uvw[:,2] *= 1000
    # random visibility values
    ms = rng.random((nrow, nchan))-0.5 + 1j*(rng.random((nrow, nchan))-0.5)
    # random visibility weights
    wgt = rng.uniform(0.9, 1.1, (nrow, nchan)) if use_wgt else None
    # no idea how many visibilities are masked, so I'm assuming that all
    # visibilities are used
    mask = np.ones((nrow,nchan), dtype=np.uint8) if use_mask else None
    if singleprec:
        ms = ms.astype("c8")
        if wgt is not None:
            wgt = wgt.astype("f4")
    dirty = ducc0.wgridder.vis2dirty(
        uvw=uvw, freq=freq, vis=ms, wgt=wgt, npix_x=nxdirty, npix_y=nydirty,
        pixsize_x=pixsizex, pixsize_y=pixsizey, epsilon=epsilon,
        do_wgridding=wstacking, nthreads=nthreads, verbosity=1, mask=mask,
        divide_by_n=False)

run_vis2dirty(17500, 17500, 5000000, 20, 1e-5, True, 8)

Note that you'll probably need to adjust the scaling of the w range (I don't know the proper numbers), as well as the FOV.
For the choice of parameters you mention, the gridder is definitely in the "sparse" regime, i.e. there are only few visibilities to grid in comparison to the number of pixels on the grid itself; consequently a lot of time is spent on FFTs, and also cache re-use is not as good as it could be.
The easiest way to improve this is to give the gridder more channels to work on in a single call. If I increase nchan from 20 to 100 in my test, transform time increased from 25s to just 38s, which is a massive performance increase.
Alternatively you can also increase nrows, but there the effects will most likely not be as pronounced.
Memory-wise it should be possible to imcrease the nchan or nrow quite a bit if you are working on typical compute nodes; I can easily go to several hundred channels on my laptop.

[mreineck]

Slightly related question: you mention "UV grid cell" ... how does DSA define its size? Given the coverage above, I'm absolutely sure that most of the calls in the uv grid internally used by the wgridder will not contain a visibility when the number of visibilities is 5.5e6*20, since the uv grid size is dirty image size oversampled by roughly 1.5 in both dimensions, which s already larger than the total number of visibilities. Additionally, the cell count varies over orders of magnitude, which makes the effect even worse.

[Joshuaalbert]

Our grid parameters are determined by an oversampling of 1.2 in the UV space and 1.5 in the LM space (which we accomplish with padding before fft). This works out to UV cellsize of roughly 2.3m (a little under half the dish diameter).

def determine_grid_params(nu_MHz, fluxcut=0.1, maxUV=20070.75, oversampling_uv=1.2, oversampling_lm=1.5):
    fwhm = 2.5 * (1350 / nu_MHz)  # degrees
    fov_lm = fwhm * np.sqrt(-np.log(fluxcut) / np.log(2.0)) * np.pi / 180  # radians
    wavelength = 299792458 / (nu_MHz * 1e6)  # meters
    fov_uv = 2 * maxUV / wavelength  # lambdas
    dl = 1 / (oversampling_lm * fov_uv)  # radians
    du = 1 / (oversampling_uv * fov_lm)  # lambdas
    n_lm = int(2 * np.sin(fov_lm / 2) / dl) // 2 * 2
    n_uv = int(fov_uv / du) // 2 * 2
    return n_lm, dl, fov_lm, n_uv, du, fov_uv

[mreineck]

Thanks!
At 1350MHz that gives me an n_uv of 17250, i.e. quite close to the 17500 resolution of the dirty image you mentioned above.
Just to be sure: does your above statement "dirty.shape = [17.5k, 17.5k]" refer to the un-clipped dirty image? If so, I should stress that the output of wgridder is always the clipped dirty image; in other words, you can assume the specified accuracy everywhere in the image, even at the borders.
(I have to do it this way, since wgridder chooses the oversampling factor on the fly, and it cannot really be predicted how large the un-trimmed dirty image will be.)
Coming back to the coverage density of the uv grid; using 5 million rows and 20 channels we have 1e8 visibilities, vs. 17250**2 (roughly 3e8) uv grid cells. Even with a fully homogeneous distribution this means that the majority of cells are empty. I cannot figure out how this matches with your coverage plot ... does it cover a longer time span?
In any case, this ratio of visibilities vs. uv grid size puts wgridder into a fairly inefficient mode; it will do much better if you can provide more visibilities per call.

[Joshuaalbert]

Thanks for the performance test script. Yes my plot above is for a longer timespan. I extended your test, and also ran for degridding too. I see degridding efficiency gains are even greater:

gridding:

done: nrow=1360425 nchan=1 npix=17244 median=13.5s  vis/s=1e+05
done: nrow=1360425 nchan=20 npix=17244 median=17.2s  vis/s=1.58e+06
done: nrow=1360425 nchan=40 npix=17244 median=19.7s  vis/s=2.76e+06
done: nrow=4081275 nchan=1 npix=17244 median=19.2s  vis/s=2.12e+05
done: nrow=4081275 nchan=20 npix=17244 median=23.3s  vis/s=3.5e+06
done: nrow=4081275 nchan=40 npix=17244 median=24.5s  vis/s=6.65e+06
done: nrow=6802125 nchan=1 npix=17244 median=22.4s  vis/s=3.04e+05
done: nrow=6802125 nchan=20 npix=17244 median=25.8s  vis/s=5.27e+06
done: nrow=6802125 nchan=40 npix=17244 median=26.8s  vis/s=1.02e+07
done: nrow=9522975 nchan=1 npix=17244 median=24s  vis/s=3.97e+05
done: nrow=9522975 nchan=20 npix=17244 median=26.7s  vis/s=7.13e+06
done: nrow=9522975 nchan=40 npix=17244 median=28.2s  vis/s=1.35e+07
done: nrow=12243825 nchan=1 npix=17244 median=25.1s  vis/s=4.88e+05
done: nrow=12243825 nchan=20 npix=17244 median=27.7s  vis/s=8.84e+06
done: nrow=12243825 nchan=40 npix=17244 median=29.2s  vis/s=1.68e+07
done: nrow=14964675 nchan=1 npix=17244 median=25.8s  vis/s=5.8e+05
done: nrow=14964675 nchan=20 npix=17244 median=28.2s  vis/s=1.06e+07
done: nrow=964675 nchan=40 npix=17244 median=29.9s  vis/s=2e+07

degridding:

done: nrow=1360425 nchan=1 npix=17244 median=6.62s  vis/s=2.06e+05
done: nrow=1360425 nchan=20 npix=17244 median=6.94s  vis/s=3.92e+06
done: nrow=1360425 nchan=40 npix=17244 median=7.39s  vis/s=7.37e+06
done: nrow=4081275 nchan=1 npix=17244 median=7.12s  vis/s=5.73e+05
done: nrow=4081275 nchan=20 npix=17244 median=8.24s  vis/s=9.9e+06
done: nrow=4081275 nchan=40 npix=17244 median=8.95s  vis/s=1.82e+07
done: nrow=6802125 nchan=1 npix=17244 median=7.63s  vis/s=8.91e+05
done: nrow=6802125 nchan=20 npix=17244 median=9.28s  vis/s=1.47e+07
done: nrow=6802125 nchan=40 npix=17244 median=10.3s  vis/s=2.64e+07
done: nrow=9522975 nchan=1 npix=17244 median=7.86s  vis/s=1.21e+06
done: nrow=9522975 nchan=20 npix=17244 median=9.91s  vis/s=1.92e+07
done: nrow=9522975 nchan=40 npix=17244 median=11.9s  vis/s=3.19e+07
done: nrow=12243825 nchan=1 npix=17244 median=8.01s  vis/s=1.53e+06
done: nrow=12243825 nchan=20 npix=17244 median=10.7s  vis/s=2.3e+07
done: nrow=12243825 nchan=40 npix=17244 median=12.9s  vis/s=3.8e+07
done: nrow=14964675 nchan=1 npix=17244 median=8.12s  vis/s=1.84e+06
done: nrow=14964675 nchan=20 npix=17244 median=11.2s  vis/s=2.68e+07
done: nrow=14964675 nchan=40 npix=17244 median=14s  vis/s=4.27e+07

System:
  Python:   3.11.8
  NumPy:    2.4.0
  ducc0:    0.40.0
  Machine:  aarch64  (Linux 6.2.0-1015-nvidia-64k)

[mreineck]

These numbers are very interesting ... it seems that the FFT is even more of a bottleneck on your machine than it is on mine. Are you using very many threads?

[mreineck]

In case you are using as many threads as your machine has virtual cores, please also try just using a thread per physical core; this may improve things.

[Joshuaalbert]

I'm using 72 threads, and have 72 physical cores:

  Model name:           Neoverse-V2
    Thread(s) per core: 1
    Core(s) per socket: 72