Efficiently compute parallel matrix-vector product for block vectors with FFTs?

Question

Assume I have $P$ processors, each having a different vector $v_p$ of size $N$, $p=1, ..., P$. I learned in this question/answer that for computing the matrix-vector product $$w = (E\otimes I_N)v$$ in parallel, where $\otimes$ is the Kronecker product, $v = (v_1, ..., v_P)^T$ is the stacked vector consisting of all local vectors $v_p$, $E = (e_{pj})_{p,j}\in\mathbb{C}^{P\times P}$ some dense matrix and $I_N$ is the identity matrix of size $N\times N$, I can basically do $P$ parallel sums using binary trees.

$N$ is pretty large (say, $10^8$), in particular much larger than $P$ (which is only 10-100) and so large that $NP$ (the size of $w$ and/or all vectors $v_p$ together) does not fit into each processor's memory.

Depending on the application, the matrix $E$ could be a Fourier or DFT matrix, i.e. computing $$w_p = \sum_{j=1}^P e_{pj} v_j$$ should be possible using an FFT. For each of the components of the vector $v_j\in\mathbb{R}^N$, the FFT would not change. Do I still have to do $N$ FFTs or is there a "block-wise" FFT idea/implementation?

Answer 1

Let's look at the case where $P=2$, for starters. The FFT can be computed in a divide-and-conquer fashion, where first we compute the FFT of the even-numbered indices, then the FFT of the odd-numbered indices, and then combine them in a simple step (with a single layer of butterfly operations).

So, if the even-numbered indices are stored on one processor and the odd-numbered indices are stored on a second processor, then you could do this in parallel in a natural way, with the first processor doing the FFT on its own data, and the second processor doing the FFT on its own data; the final butterfly step requires only $\Theta(NP)$ data sent over the network.

If the data isn't partitioned that way, then the first step could be to send all the values at even-numbered indices to the first processor and the values at odd-numbered indices to the second processor, then do the above process. In total it will require $\Theta(NP)$ data sent over the network and $\Theta(N \lg N)$ computation ($\Theta(N \lg N)$ computation at each processor to compute the FFT locally).

More generally, if $P$ is a power of 2, you can generalize these ideas in a similar fashion. You send the data to the appropriate processor ($\Theta(NP)$ data sent over the network), then compute the FFT locally at each processor ($\Theta(N \lg N)$ steps of computation at each processor), then combine them with $\lg P$ butterfly layers (requires $\Theta(NP \lg P)$ data sent over the network).

If $P$ isn't a power of two, you can pick the largest power of two that is smaller than $P$, and use that many processor (the remainder are left idle).