Exploring the possibility of expanding the support to multi-CPU/GPU

[doraemonho]

Hi,

I am exploring the possible way to implement the multi-XPU computation under the FourierFlows framework. It seems to me that the current bottleneck is on the native FFT transform, which only supports the single CPU/GPU. I see a few Julia packages trying to tackle this problem (like PencilFFTs.jl ). I wonder if we can utilize those packages and expand the support of FourierFlows to multi-CPU/GPU.