(0.91.6) Distributed FFTs using Oceananigans' inhouse DiscreteTransforms #3279
Conversation
There was a problem hiding this comment.
Thanks @simone-silvestri this is great and I'm happy with the changes, so I'm approving.
I will mention though, there are still quite a few comments from @glwagner from Jun 28th that seem to be unresolved, especially for files distributed_transpose.jl amd transposable_field.jl. But I see that he already approved the PR so maybe it doesn't matter :)
🚀
|
|
||
| zgrid = field_in.grid # We support only a 2D partition in X and Y | ||
| ygrid = twin_grid(zgrid; free_dimension = :y) | ||
| xgrid = twin_grid(zgrid; free_dimension = :x) |
src/DistributedComputations/distributed_fft_based_poisson_solver.jl
Outdated
Show resolved
Hide resolved
src/DistributedComputations/distributed_fft_based_poisson_solver.jl
Outdated
Show resolved
Hide resolved
| 1. `storage.zfield`, partitioned over ``(x, y)`` is initialized with the `rhs`. | ||
| 2. Transform along ``z``. | ||
| 3 Transpose + communicate to `storage[2]` in layout ``(x, z, y)``, | ||
| which is distributed into `(Rx, Ry)` processes in ``(z, y)``. | ||
| 4. Transform along ``x``. | ||
| 5. Transpose + communicate to `last(storage)` in layout ``(y, x, z)``, | ||
| which is distributed into `(Rx, Ry)` processes in ``(x, z)``. | ||
| 6. Transform in ``y``. | ||
| 3 Transpose + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``. | ||
| 4. Transform along ``y``. | ||
| 5. Transpose + communicate to `storage.xfield` partitioned into `(Rx, Ry)` processes in ``(y, z)``. | ||
| 6. Transform in ``x``. |
There was a problem hiding this comment.
Sounds same question as mine.. (the one by @tomchor egarding storage)
I don't understand this comment |
I have added a line in the docstring explaining what storage is |
|
@navidcy are you satisfied with the corrections? |
… into ss/distributed-fft
simone-silvestri
changed the title
|
@simone-silvestri you are now bumping to 0.92. Can you clarify what's breaking about this PR? |
|
Well it adds a new feature. I could just bump to 0.91.6 instead but it seems quite a major feature to justify the bump in version. |
|
Only bump to a new version for a breaking change, 91 versions is enough! |
simone-silvestri
changed the title
|
merge time? |
This PR removes the PencilFFT library from Oceananigans and builds a distributed FFT solver using Oceananigans' inhouse transforms. This allows us to run on GPUs both periodic and bounded domains.
No stretched mesh is supported at the moment (that will come in a later PR)
The transposition is performed through a custom
transposeroutine built for Oceananigans' fields that assumesAn additional assumption is that:
All these assumptions can be relaxed in following PRs