diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index f2ce92160a..857ccc9b7c 100644 --- a/src/Fields/Fields.jl +++ b/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") diff --git a/src/Fields/kernel_computed_field.jl b/src/Fields/kernel_computed_field.jl new file mode 100644 index 0000000000..44b669324a --- /dev/null +++ b/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_variance!(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_variance!, model; field_dependencies=(u, U,)) + v′² = KernelComputedField(Cell, Face, Cell, compute_variance!, model; field_dependencies=(v, V,)) + w′² = KernelComputedField(Cell, Cell, Face, compute_variance!, 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) diff --git a/test/runtests.jl b/test/runtests.jl index cc2855aee6..982c81d142 100644 --- a/test/runtests.jl +++ b/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 diff --git a/test/test_abstract_operations.jl b/test/test_abstract_operations_computed_field.jl similarity index 100% rename from test/test_abstract_operations.jl rename to test/test_abstract_operations_computed_field.jl diff --git a/test/test_kernel_computed_field.jl b/test/test_kernel_computed_field.jl new file mode 100644 index 0000000000..afb32d0cab --- /dev/null +++ b/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