Poisson eigenvalues should have the same float type as the grid

[loisbaker]

The FFT Poisson solvers should have eigenvalues with the same float type as the grid:

using Oceananigans 
using Oceananigans.Solvers: FFTBasedPoissonSolver, FourierTridiagonalPoissonSolver

# The float type of the Poisson operator eigenvalues should be the same as that of the grid
grid = RectilinearGrid(CPU(), Float32, size = (2,2), extent=(1,1), topology = (Periodic, Flat, Bounded))
@show eltype(grid)
                                    
solver = FFTBasedPoissonSolver(grid)
@show eltype(solver.eigenvalues.λx)

solver = FourierTridiagonalPoissonSolver(grid)
@show eltype(solver.poisson_eigenvalues[1]) 

Output:

eltype(grid) = Float32
eltype(solver.eigenvalues.λx) = Float64
eltype(solver.poisson_eigenvalues[1]) = Float64

This will become a problem once Metal supports FFTs, since it doesn't support Float64:

using Metal
metal = Metal.MetalBackend()
arch = GPU(metal)

grid_on_GPU = RectilinearGrid(arch, Float32, size = (2,2), extent=(1,1), topology = (Periodic, Flat, Bounded))
solver = FFTBasedPoissonSolver(grid_on_GPU)

Output:

ERROR: LoadError: Metal does not support Float64 values, try using Float32 instead
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] check_eltype(T::Type)
    @ Metal ~/.julia/packages/Metal/EFzAO/src/array.jl:32
  [3] MtlArray{Float64, 3, Metal.PrivateStorage}(::UndefInitializer, dims::Tuple{Int64, Int64, Int64})
    @ Metal ~/.julia/packages/Metal/EFzAO/src/array.jl:54
  [4] (MtlArray{Float64, 3})(::UndefInitializer, dims::Tuple{Int64, Int64, Int64})
    @ Metal ~/.julia/packages/Metal/EFzAO/src/array.jl:204
  [5] MtlArray
    @ ~/.julia/packages/Metal/EFzAO/src/array.jl:291 [inlined]
  [6] MtlArray(A::Array{Float64, 3})
    @ Metal ~/.julia/packages/Metal/EFzAO/src/array.jl:307
  [7] on_architecture(::GPU{MetalBackend}, a::Array{Float64, 3})
    @ OceananigansMetalExt ~/Documents/Edinburgh/Projects/Oceananigans.jl/ext/OceananigansMetalExt.jl:31
  [8] FFTBasedPoissonSolver(grid::RectilinearGrid{…}, planner_flag::UInt32)
    @ Oceananigans.Solvers ~/Documents/Edinburgh/Projects/Oceananigans.jl/src/Solvers/fft_based_poisson_solver.jl:60
  [9] FFTBasedPoissonSolver(grid::RectilinearGrid{…})
    @ Oceananigans.Solvers ~/Documents/Edinburgh/Projects/Oceananigans.jl/src/Solvers/fft_based_poisson_solver.jl:52
 [10] top-level scope
    @ ~/Documents/Edinburgh/Projects/Oceananigans.jl/examples/MWE_FT_bug.jl:21
 [11] include(fname::String)
    @ Base.MainInclude ./client.jl:494
 [12] top-level scope
    @ REPL[2]:1
in expression starting at /Users/lbaker/Documents/Edinburgh/Projects/Oceananigans.jl/examples/MWE_FT_bug.jl:21
Some type information was truncated. Use `show(err)` to see complete types.

I'll add a PR to fix this by passing float type to poisson_eigenvalues() @simone-silvestri

[glwagner]

@loisbaker we use Float64 for FFTs even for grids with Float32 because we found that Float32 pressure correction produce much more unstable simulations than Float64 pressure correction.

These studies were done a while ago and not documented...

I would suggest that we take this approach:

1. Add a comment to the code so that future investigators understand why we are using Float64 for the pressure solver even when the grid is Float32
2. Create the possibility of a Metal extension of the Poisson solver that does use Float32 on Metal, simply because this is the only option.

I am also curious, are FFTs available on Metal?

[glwagner]

Another possibility is to revisit Float32 LES systematically, and establish with a new and documented study whether or not we should be using Float64 or Float32 for the pressure solver (when Float64 is available).

If the results of that study could show that Float32 pressure solver is adequate, then indeed I would support changing the behavior for all architectures, not just Metal.

[glwagner]

@ali-ramadhan was involved in that early work if I recall and might have more to say (although it is ancient history at this point)

[simone-silvestri]

I am also curious, are FFTs available on Metal?

Looks like there is a PR in Metal.jl implementing FFT transforms. JuliaGPU/Metal.jl#713 So as soon as that passes we should be ready to use the non hydrostatic model on a mac!

[glwagner]

Great, so I suggest trying to do something in the Metal extension that will override the Float64 default in src. You could introduce something like default_fft_floattype(grid) for example, and then extend this for a grid on Metal architecture. Make sense @loisbaker @simone-silvestri ?

[simone-silvestri]

I was thinking passing the architecture to the poisson_eigenvalues so that we can extend the function itself in Metal.

[glwagner]

The arch does not know the float type that is being used by the grid.

[loisbaker]

Would poisson_eigenvalues need to know the grid float type? Or could it just default to Float32 for Metal?

[glwagner]

Here is a bit of philosophy...

Not using the grid requires invoking two fragile, not-necessarily-true assumptions:

• There are only two float types of interest, f32 and f64
• When we are not using metal, we always want to use f64 and never the eltype of the grid

Inserting the grid is future proof, because you do not have to make these assumptions. Instead you implement a design that generalizes to cases with more number types, and also a potential future situation where we deduce (because of numerics changes or other knowledge) that we do want to use the eltype of the grid for the solver.

All else being equal, therefore, I would encourage using a future proof design (this is also general advice, not specific here). It is tempting to take the absolute minimum number of steps towards an objective, but if this path requires fragile assumptions to hold forevermore, it is not the best path.

That said, I haven't seen the code, so if there is some major refactor that passing the grid would trigger, that is a reason to take a shortcut. In this case I don't see how that could be, because we need to know the grid in order to build the solver (indeed, the arch must be extracted from the grid in the first place, so arguably passing the arch is more steps, not fewer).

[loisbaker]

Thanks @glwagner ! The philosophy makes sense :)

Is it correct that in order to extend poisson_eigenvalues in OceananigansMetalExt to use Float32, we would need to pass in arch and not just grid? So is the suggestion to pass in both, and not use either for the other architectures (right now)?

[glwagner]

The suggestion is to pass just grid, and define a type MetalGrid (from AbstractGrid, which knows the type of the architecture). Then use that to extend poisson_eigenvalues for MetalGrid`

[glwagner]

We do the same for Clock:

Oceananigans.jl/ext/OceananigansMetalExt.jl

Lines 71 to 72 in e4ebdcb

	const MetalGrid = GD.AbstractGrid{<:Any, <:Any, <:Any, <:Any, <:MetalGPU}
	Clock(grid::MetalGrid) = Clock{Float32}(time=0)

[simone-silvestri]

I am not sure if it is better passing the grid or the architecture.
If we need to force a float32 then the architecture can be enough.

We can probably do it in two different ways:

poisson_eigenvalues(grid, dim)

ancd dispatch on the grid, where N, L and T are inferred by dim and grid so that we do not duplicate the information passed to the function, or

poisson_eigenvalues(arch, N, L, dim, T)

or dispatch on the architecture

[glwagner]

If we need to force a float32 then the architecture can be enough.

This is only true if we assume that we always use float64 otherwise.