I am currently working on a more efficient implementation of a pulse propagation algorithm. The propagation is done in the spectral domain, but several evaluations (such as energy calculations) are done in the time domain. The field is either defined in cylindrical coordinates (i.e. $E=E(r, t)$) or in cartesian coordinates (i.e. $E=E(x, y, t)$).
To convert the field from the temporal domain into the spectral domain, I can either apply the hankel transformation for converting the field from $r$-space to $k$-space, and afterwards a 1d-FFT for converting the field from the $t$-space into the $\omega$-space (for the cylindrical coordinate system), or I can first apply a 2d-FFT for converting the field from the $x, y$-space into the $k_x,\,k_y$-space, and afterwards a 1d-FFT for going from the $t$-space into the $\omega$-space. The Hankel transformation can be reduced into a matrix-matrix-multiplication.
In my current implementation I am calculating everything on a GPU, while using
ArrayFire as library in the backend. Still, for higher resolutions and longer pulses the memory of my GPU is not enough, and I have to switch to using the CPU instead. According to benchmarks, running the calculations on a single node this might increase the runtime by a factor of 5-20, depending on parameters. Therefore, I was intending to use several nodes together (using MPI) to reduce that factor at least a bit. Would that make sense?
ArrayFire does not support such calculations (zgemm/n-FFT) via MPI, but I also do not want to use several different libraries just depending on the target device. Are there libraries which support those operations both for GPUs, local threads on a CPU node and MPI for several nodes? I have seen that PETSc might support that, but I am not sure about that. Are there others? Or would another approach be easier here?