Adds kernel computed field

src/Fields/Fields.jl

@@ -4,7 +4,7 @@ export
     Face, Cell,
     AbstractField, Field,
     CellField, XFaceField, YFaceField, ZFaceField,
-    ReducedField, AveragedField, ComputedField, BackgroundField,
+    ReducedField, AveragedField, ComputedField, KernelComputedField, BackgroundField,
     interior, interiorparent, data,
     xnode, ynode, znode, location,
     set!, compute!, @compute,
@@ -26,6 +26,7 @@ include("zero_field.jl")
 include("reduced_field.jl")
 include("averaged_field.jl")
 include("computed_field.jl")
+include("kernel_computed_field.jl")
 include("pressure_field.jl")
 include("function_field.jl")
 include("set!.jl")

src/Fields/kernel_computed_field.jl

@@ -0,0 +1,103 @@
+using Oceananigans: AbstractModel
+using Oceananigans.Grids
+using Oceananigans.Utils: tupleit
+
+struct KernelComputedField{X, Y, Z, S, A, G, K, F, P} <: AbstractField{X, Y, Z, A, G}
+                  data :: A
+                  grid :: G
+                kernel :: K
+    field_dependencies :: F
+            parameters :: P
+                status :: S
+
+    """
+        KernelComputedField(loc, kernel, grid)
+
+    Example
+    =======
+
+    ```julia
+    using KernelAbstractions: @index, @kernel
+    using Oceananigans.Fields: AveragedField, KernelComputedField, compute!
+    using Oceananigans.Grids: Cell, Face
+
+    @inline ψ²(i, j, k, grid, ψ, Ψ) = @inbounds (ψ[i, j, k] - Ψ[i, j, k])^2
+    @kernel function compute_var!(var, grid, ϕ, Φ)
+        i, j, k = @index(Global, NTuple)
+
+        @inbounds var[i, j, k] = ψ′²(i, j, k, grid, ϕ, Φ)
+    end
+
+    u, v, w = model.velocities
+
+    U = AveragedField(u, dims=(1, 2))
+    V = AveragedField(V, dims=(1, 2))
+
+    u′² = KernelComputedField(Face, Cell, Cell, compute_var!, model; field_dependencies=(u, U,))
+    v′² = KernelComputedField(Cell, Face, Cell, compute_var!, model; field_dependencies=(v, V,))
+    w′² = KernelComputedField(Cell, Cell, Face, compute_var!, model; field_dependencies=(w, 0,))
+
+    compute!(u′²)
+    compute!(v′²)
+    compute!(w′²)
+    ```
+    """
+    function KernelComputedField{X, Y, Z}(kernel::K, arch, grid;
+                                          field_dependencies = (),
+                                          parameters::P = nothing,
+                                          data = nothing,
+                                          recompute_safely = true) where {X, Y, Z, K, P}
+
+        field_dependencies = tupleit(field_dependencies)
+
+        if isnothing(data)
+            data = new_data(arch, grid, (X, Y, Z))
+        end
+
+        # Use FieldStatus if we want to avoid always recomputing
+        status = recompute_safely ? nothing : FieldStatus(0.0)
+
+        G = typeof(grid)
+        A = typeof(data)
+        S = typeof(status)
+        F = typeof(field_dependencies)
+
+        return new{X, Y, Z, S,
+                   A, G, K, F, P}(data, grid, kernel, field_dependencies, parameters)
+    end
+end
+
+KernelComputedField(X, Y, Z, kernel, model::AbstractModel; kwargs...) =
+    KernelComputedField{X, Y, Z}(kernel, model.architecture, model.grid; kwargs...)
+
+KernelComputedField(X, Y, Z, kernel, arch::AbstractArchitecture, grid::AbstractGrid; kwargs...) =
+    KernelComputedField{X, Y, Z}(kernel, arch, grid; kwargs...)
+
+function compute!(kcf::KernelComputedField{X, Y, Z}) where {X, Y, Z}
+
+    for dependency in kcf.field_dependencies
+        compute!(dependency)
+    end
+
+    arch = architecture(kcf.data)
+
+    workgroup, worksize = work_layout(kcf.grid,
+                                      :xyz,
+                                      location=(X, Y, Z),
+                                      include_right_boundaries=true)
+
+    compute_kernel! = kcf.kernel(device(arch), workgroup, worksize)
+
+    event = isnothing(kcf.parameters) ?
+        compute_kernel!(kcf.data, kcf.grid, kcf.field_dependencies...) :
+        compute_kernel!(kcf.data, kcf.grid, kcf.field_dependencies..., kcf.parameters)
+
+    wait(device(arch), event)
+
+    return nothing
+end
+
+compute!(field::KernelComputedField{X, Y, Z, <:FieldStatus}, time) where {X, Y, Z} =
+    conditional_compute!(field, time)
+
+Adapt.adapt_structure(to, kcf::KernelComputedField) = Adapt.adapt(to, kcf.data)

test/runtests.jl

@@ -82,6 +82,7 @@ group = get(ENV, "TEST_GROUP", :all) |> Symbol
             include("test_boundary_conditions.jl")
             include("test_fields.jl")
             include("test_averaged_field.jl")
+            include("test_kernel_computed_field.jl")
             include("test_halo_regions.jl")
             include("test_solvers.jl")
             include("test_pressure_solvers.jl")
@@ -113,7 +114,7 @@ group = get(ENV, "TEST_GROUP", :all) |> Symbol
             include("test_simulations.jl")
             include("test_diagnostics.jl")
             include("test_output_writers.jl")
-            include("test_abstract_operations.jl")
+            include("test_abstract_operations_computed_field.jl")
         end
     end
 

test/test_kernel_computed_field.jl

@@ -0,0 +1,50 @@
+using Oceananigans.Fields: KernelComputedField
+using KernelAbstractions: @kernel, @index
+
+grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1),
+                            topology=(Periodic, Periodic, Bounded))
+
+@kernel function double!(doubled_field, grid, single_field)
+    i, j, k = @index(Global, NTuple)
+    doubled_field[i, j, k] = 2 * single_field[i, j, k]
+end
+
+@kernel function multiply!(multiplied_field, grid, field, multiple)
+    i, j, k = @index(Global, NTuple)
+    multiplied_field[i, j, k] = multiple * field[i, j, k]
+end
+
+for arch in archs
+    @testset "KernelComputedField [$(typeof(arch))]" begin
+        @info "  Testing KernelComputedField..."
+
+        single_field = Field(Cell, Cell, Cell, arch, grid)
+
+        doubled_field = KernelComputedField(Cell, Cell, Cell,
+                                            double!, arch, grid,
+                                            field_dependencies = single_field)
+
+        # Test that the constructor worked
+        @test doubled_field isa KernelComputedField
+
+        multiple = 3
+        multiplied_field = KernelComputedField(Cell, Cell, Cell,
+                                               multiply!, arch, grid,
+                                               field_dependencies = doubled_field,
+                                               parameters = multiple)
+
+        # Test that the constructor worked
+        @test multiplied_field isa KernelComputedField
+
+        set!(single_field, π)
+
+        compute!(doubled_field)
+        @test doubled_field.data[1, 1, 1] == 2π
+
+        set!(doubled_field, 0)
+        compute!(multiplied_field)
+
+        @test doubled_field.data[1, 1, 1] == 2π
+        @test multiplied_field.data[1, 1, 1] == multiple * 2 * π
+    end
+end