Support SPMD sharding of fft ops.

[hawkinsp]

In jax-ml/jax#13081 we found that XLA doesn't support SPMD sharding of fast-fourier transform ops. It should!

[cheshire]

Tracked internally in b/263023739

[jon-chuang]

1. Dim 1 FFT: To simplify the problem, consider when the number of shards and FFT size are powers of 2, i.e. num_shards=2^k. - map-reduce style: Do a size `N/num_shards` FFT on each device. Subsequently, there is a shuffle per stage. 2. For dim > 1 FFT, The problem actually becomes simpler: if N % num_shards == 0, one can perform the FFTs along the row dim, and then do a shuffle, performing the FFTs in parallel along the column dim. (see: http://dsp-book.narod.ru/FFTBB/0270_PDF_C23.pdf)

So we would have to restrict to certain choices of $N$ and num_shards

Oh wait, this is parallelizing single FFT... smpd sharding (pmap) should be a lot simpler.

[jon-chuang]

Wait, there seems to be previous work on SPMD for FFT: 180dbd6.
Seems it only supports rank >=3:

xla/xla/service/spmd/fft_handler.cc

Line 350 in 7b562aa

if (hlo->operand(0)->shape().rank() < 3 || hlo->fft_type() != FftType::FFT) {

And it does indeed do a sharded parallel FFT...

[jon-chuang]

Could you clarify what is meant by SPMD in the case of FFTs @hawkinsp @cheshire ?

Does this mean running multiple FFTs on tensors with rank >= 2, along a batch dimension? Or does it mean parallelising a single FFT across the FFT's ranks?