General matrix constructor for Multigrid solver on GPU uses scalar operations

[navidcy]

At the moment, the construction of a matrix from a linear operation for the multigrid solver on the GPU involves scalar operations:

https://github.com/CliMA/Oceananigans.jl/blob/60470f4d716026f161dfd303a30d9d14fbfb97a1/src/Solvers/multigrid_solver.jl#L246-L279

This matrix construction is only called once upon initialization of the solver.

Furthermore, this particular method that uses scalar operations is not used by the implicit free-surface solvers since in that case we provide the MultigridSolver constructor with the matrix itself, see:

https://github.com/CliMA/Oceananigans.jl/blob/60470f4d716026f161dfd303a30d9d14fbfb97a1/src/Models/HydrostaticFreeSurfaceModels/mg_implicit_free_surface_solver.jl#L77-L104

For general matrix construction that works for any linear operation, it would be nice to avoid using scalar operations on the GPU or construct the matrix on the CPU and convert it to the GPU afterwards.

cc @elise-palethorpe

[navidcy]

these comments on #2688 are relevant:

@elise-palethorpe -> https://github.com/CliMA/Oceananigans.jl/pull/2688#discussion_r967800941

@simone-silvestri -> https://github.com/CliMA/Oceananigans.jl/pull/2688#discussion_r970861258

[glwagner]

Shoudn't we write a kernel? Presumably that explicit 3D loop is incredibly slow?

[navidcy]

Ideally we should!

[elise-palethorpe]

Shoudn't we write a kernel? Presumably that explicit 3D loop is incredibly slow?

Indeed! Unfortunately it's not embarrassingly parallel so my kernel skills are no longer sophisticated enough :) I don't know how to gather count, rowval or colval. The conditional preseumably also leads to poor performance (but no where near as poor as it is now!).

[glwagner]

@simone-silvestri any tips?

For conditionals, we should use ifelse; then there is no issue for performance.