Try using Krylov to solve the Poisson equation

src/Fields/abstract_field.jl

@@ -125,3 +125,21 @@ for f in (:+, :-)
     @eval Base.$f(ϕ::AbstractField, ψ::AbstractArray) = $f(interior(ϕ), ψ)
 end
 
+# TODO: support dims? test
+function Statistics.norm(a::AbstractField; condition = nothing)
+    conditional_a = condition_operand(a, condition, 0)
+    result = zeros(a.grid, 1)
+    Base.mapreducedim!(x -> x * x, +, result, conditional_a)
+    return CUDA.@allowscalar sqrt(first(result))
+end
+
+# TODO: needs test
+function Statistics.dot(a::AbstractField, b::AbstractField)
+    conditional_a = condition_operand(a, condition, 0)
+    conditional_b = condition_operand(b, condition, 0)
+    result = zeros(a.grid, 1)
+    Base.mapreducedim!((x, y) -> x * y, +, result, conditional_a, conditional_b)
+    return CUDA.@allowscalar first(result)
+end
+
+

src/Fields/field.jl

@@ -9,7 +9,7 @@ using Base: @propagate_inbounds
 import Oceananigans: boundary_conditions
 import Oceananigans.Architectures: on_architecture
 import Oceananigans.BoundaryConditions: fill_halo_regions!, getbc
-import Statistics: norm, mean, mean!
+import Statistics: mean, mean!
 import Base: ==
 
 #####
@@ -219,6 +219,8 @@ ZFaceField(grid::AbstractGrid, T::DataType=eltype(grid); kw...) = Field((Center,
 ##### Field utils
 #####
 
+Base.copyto!(a::Field, b::Field) = copyto!(parent(a), parent(b))
+
 # Canonical `similar` for Field (doesn't transfer boundary conditions)
 function Base.similar(f::Field, grid=f.grid)
     loc = location(f)
@@ -590,9 +592,6 @@ const ReducedAbstractField = Union{XReducedAbstractField,
                                    XYReducedAbstractField,
                                    XYZReducedAbstractField}
 
-# TODO: needs test
-Statistics.dot(a::Field, b::Field) = mapreduce((x, y) -> x * y, +, interior(a), interior(b))
-
 # TODO: in-place allocations with function mappings need to be fixed in Julia Base...
 const SumReduction     = typeof(Base.sum!)
 const MeanReduction    = typeof(Statistics.mean!)
@@ -695,9 +694,9 @@ for reduction in (:sum, :maximum, :minimum, :all, :any, :prod)
             T = filltype(Base.$(reduction!), c)
             loc = reduced_location(location(c); dims)
             r = Field(loc, c.grid, T; indices=indices(c))
-            conditioned_c = condition_operand(f, c, condition, mask)
-            initialize_reduced_field!(Base.$(reduction!), identity, r, conditioned_c)
-            Base.$(reduction!)(identity, r, conditioned_c, init=false)
+            conditional_c = condition_operand(f, c, condition, mask)
+            initialize_reduced_field!(Base.$(reduction!), identity, r, conditional_c)
+            Base.$(reduction!)(identity, r, conditional_c, init=false)
 
             if dims isa Colon
                 return CUDA.@allowscalar first(r)
@@ -734,19 +733,14 @@ function Statistics.mean!(f::Function, r::ReducedAbstractField, a::AbstractField
     return r
 end
 
-Statistics.mean!(r::ReducedAbstractField, a::AbstractArray; kwargs...) = Statistics.mean!(identity, r, a; kwargs...)
-
-function Statistics.norm(a::AbstractField; condition = nothing)
-    r = zeros(a.grid, 1)
-    Base.mapreducedim!(x -> x * x, +, r, condition_operand(a, condition, 0))
-    return CUDA.@allowscalar sqrt(r[1])
-end
+Statistics.mean!(r::ReducedAbstractField, a::AbstractArray; kwargs...) =
+    Statistics.mean!(identity, r, a; kwargs...)
 
 function Base.isapprox(a::AbstractField, b::AbstractField; kw...)
-    conditioned_a = condition_operand(a, nothing, one(eltype(a)))
-    conditioned_b = condition_operand(b, nothing, one(eltype(b)))
+    conditional_a = condition_operand(a, nothing, one(eltype(a)))
+    conditional_b = condition_operand(b, nothing, one(eltype(b)))
     # TODO: Make this non-allocating?
-    return all(isapprox.(conditioned_a, conditioned_b; kw...))
+    return all(isapprox.(conditional_a, conditional_b; kw...))
 end
 
 #####

src/Solvers/Solvers.jl

@@ -43,6 +43,7 @@ include("plan_transforms.jl")
 include("fft_based_poisson_solver.jl")
 include("fourier_tridiagonal_poisson_solver.jl")
 include("conjugate_gradient_poisson_solver.jl")
+include("krylov.jl")
 include("sparse_approximate_inverse.jl")
 include("matrix_solver_utils.jl")
 include("sparse_preconditioners.jl")

src/Solvers/krylov.jl

@@ -0,0 +1,32 @@
+import Krylov: kaxpby!, kaxpby!, kdot, knrm2
+
+using KernelAbstractions: @index, @kernel
+
+kdot(n::Integer, x::Field, dx::Integer, y::Field, dy::Integer) = dot(x, y)
+knrm2(n::Integer, x::Field, dx::Integer) = norm(x)
+kcopy!(n::Integer, x::Field, dx::Integer, y::Field, dy :: Integer) = copyto!(y, x)
+
+function kaxpy!(n::Integer, s::T, a::Field, dx::Integer, y::Field, dy::Integer) where T<:AbstractFloat
+    grid = a.grid
+    arch = architecture(grid)
+    launch!(arch, grid, size(a), _axpy!, y, s, x)
+    return nothing
+end
+
+@kernel function _axpy!(y, s, x)
+    i, j, k = @index(Global, NTuple)
+    @inbounds y[i, j, k] += s * x[i, j, k]
+end
+
+function kaxpby!(n::Integer, s::T, a::Field, dx::Integer, t::T, y::Field, dy::Integer) where T<:AbstractFloat
+    grid = a.grid
+    arch = architecture(grid)
+    launch!(arch, grid, size(a), _axpy!, y, t, s, x)
+    return nothing
+end
+
+@kernel function _axpby!(y, t, s, x)
+    i, j, k = @index(Global, NTuple)
+    @inbounds y[i, j, k] = s * x[i, j, k] + t * y[i, j, k]
+end
+