GPU Kernel

The imaging computation in epic++ is performed by a single CUDA kernel, BlockEpicImager (src/ex/fft_dx.cuh). The design was published by Reddy et al. (2023), and the analysis that motivated it is given in Memo 009.

Memory-bandwidth limitation of multi-kernel implementations

Earlier implementations executed EPIC as a sequence of separate GPU kernels: clearing of the grid, gridding of the voltages, a cuFFT transform, cross-multiplication and accumulation (Figure 2, left).

Each kernel reads the complete grid from global memory and writes it back. The arithmetic intensity of every stage is low:

• gridding is a scatter-add;
• the FFT performs $\mathcal{O}({\log }_{2}N)$ operations per element;
• multiplication performs one operation per element.

Kent et al. (2019) showed by roofline analysis that such an implementation is bounded by memory bandwidth. A profile of the Bifrost implementation on an RTX 2080 Ti (Memo 009) gave the following per-gulp times for 64 × 64 images, 90 channels and a single polarization:

The total exceeds the duration of the gulp. Every stage operated close to the peak memory bandwidth, and clearing the grid alone accounted for approximately 19% of the execution time.

Fused on-chip design

A 128 × 128 complex grid in half precision occupies 64 KiB, which can be held in the shared memory and registers of a single streaming multiprocessor (SM). The complete per-channel computation can therefore be performed within one thread block. Global memory is accessed only to read the 4-bit voltages and to write the accumulated products. The kernel is organized as follows:

• Parallelization. Each thread block processes one frequency channel. On an RTX 4090, 128 SMs process 128 channels concurrently.
• Time loop. Each block iterates over the 1000 spectra of a gulp, imaging each spectrum and accumulating the polarization products in registers.
• Output. Results are written to global memory once per gulp, as one float4 per pixel.

Figure 2 contrasts the two organizations.

for each spectrum t in the gulp (1000 iterations):
    clear half-grid in shared memory
    grid antennas onto upper half  → registers
    grid antennas onto lower half  → registers
    accumulate per-antenna autocorrelations (shared memory)
    row FFTs (registers) → triangular transpose (shared memory) → column FFTs (registers)
    accumulate |X|², |Y|², X·Y* per pixel (registers)
after the loop:
    transform gridded autocorrelations; subtract from XX*, YY*
    apply grid correction
    write or add (XX*, YY*, Re XY*, Im XY*) to global memory

Implementation techniques

Dual-polarization transforms. The X and Y grids are packed into a single complex<half2> value per cell in the "RRII" layout. The real parts of both polarizations occupy one 32-bit register and the imaginary parts a second. The block-level FFTs of cuFFTDx (NVIDIA MathDx) transform both polarizations in a single pass. For a 128² image, the thread block has 8 × 128 = 1024 threads, each holding 16 elements of one row.

Half-grid gridding. A full 128 × 128 dual-polarization grid would require 128 KiB of shared memory, which exceeds the capacity of one SM. The upper half of the aperture is therefore gridded first and moved into registers, and then the lower half is gridded. Shared memory holds at most half of the grid at any time.

Gridding by cooperative thread groups.

• The threads of a block are partitioned into tiles of $S\times S$ threads, where $S$ is the kernel support. Each tile grids one antenna at a time.
• Each thread computes the contribution ${\varphi }_a{E}_a{w}_{a,ij}$ to one cell and adds it to shared memory with an atomic __half2 addition.
• Antenna positions and phase corrections are loaded into shared memory once per gulp.
• The kernel weights $w_{a,ij}$ are precomputed for each tuning.

Triangular transpose. The two-dimensional FFT is computed as row transforms in registers, followed by a transpose through shared memory and a second set of row transforms. The transpose exchanges the upper and lower triangles of the matrix, so only half of the matrix need reside in shared memory.

Numerical range. Two successive unnormalized 128-point transforms of 4-bit inputs can exceed the largest representable half-precision value, 65 504. A factor of $1/N$ is therefore applied during the transpose. Products are accumulated in single precision, or optionally in bfloat16 (--accum_16bit). Half-precision arithmetic biases accumulated power low by 0.6–0.7% on average, with a maximum of approximately 2%. Memo 011 adopts 1.5% as a conservative error bound.

Output compression. Only $X{X}^{\ast }$, $Y{Y}^{\ast }$ and the complex $X{Y}^{\ast }$ are retained, since $Y{X}^{\ast }=(X{Y}^{\ast }{)}^{\ast }$. Each pixel is thus represented by four real values rather than six.

Performance

The performance of the kernel on an NVIDIA RTX 4090 (Reddy et al. 2023) is summarized below.

Limits on scaling.

• Channels per GPU. The kernel requires one SM per channel at 128². The 128 SMs of an RTX 4090 therefore process 128 of the 132 channels in each F-engine sub-band, so six GPUs process 19.2 MHz of the 19.8 MHz available.
• Image size. Images of 256 × 256 pixels would match the angular resolution of the station at the upper end of the band. They require approximately four times the on-chip memory per channel, which is attainable only with distributed shared memory across a cluster of SMs (compute capability 9.0) or by partitioning the computation across GPUs.

Auxiliary kernels

Two auxiliary kernels are executed once per tuning:

• ComputeGcfElements integrates the prolate spheroidal kernel over each grid cell, for each antenna and channel, using a 5 × 5 quadrature. It then normalizes the weights of each antenna to unit sum.
• ComputeAvgGriddingKernel averages the weights over antennas. The result is used to compute the image-plane grid correction, $1/|\mathcal{F}[\bar{w}]{|}^2$.