Broadcasting for AbstractField, AbstractReducedField, and AbstractOpe…

…rations (#1596)

* Changes fourth type parameter of AbstractField to architecture

* Adds architecture property to all fields

* Updates set! for new AbstractField type parameters

* Update constructor for AveragedField

* Adds architecture property to BuoyancyField

* Fixes cpudata

* set! for GPU fields is identitcal for arrays and functions

* Technical debt *sigh*

* Subtype abstract array

* Refactors AbstractField infrastructure to subtype AbstractArray

* Implements copyto! for in-place broadcasting to dest::AbstractField

* Remove spurious setindex! for Field

* Adds back manual reduction functions (hopefully we can get rid of these eventually)

* Only reduce fields backed by concrete data

* Adds 2d, 1d, and 0d Cartesian indexing for ReducedFields, and more flexible broadcast_2d! kernel

* Expand definitions of Base.similar for AbstractField and AbstractReducedField

* Dont copy boundary conditions in similar

* Removes calls to parent() in conjugate gradient iteration

* Dont test that similar fields have the same boundary conditions

* Cleans up AveragedField test for different Float types

* Fields take eltype from grid

* Fixes typo in KernelComputedField constructor

* Update fill_halo_regions dispatch now that Field<:AbstractArray

* Move fill_halo_regions for fields to AbstractField

* Adds some rudimentary broadcasting test for Fields

* Dont convert arrays to tuples

* Cleans up KernelComputedField

* Adds show method for kernelComputedField

* Better docstring for Field

* Fixes bug in test for broadcasting with ReducedField

* Simplifies KernelComputedField tests

* Different kernels for different 2d configurations

* Converts broadcasted to abstract operations where appropriate

* Needed eval after all

* AbstractOperations lists are symbols?

* Disambiguate distributed Field constructor

* Get AbstractGrid into Distributed

* Provide rules for broadcasting Fields with Arrays

* Import missing functions and types from Base.Broadcast

* Adds a new function at(loc, op) for changing the location of AbstractOperations

* Its not a derivative _operation_..."

* Attempts to insert logic to insert locations into broadcasted field expressions

* Refactors AbstractOperations to be more broadly useful in broadcasting

* Improvements to braodcasting field operations

* Fixes bug in broadcasting test for ReducedField

* Dont use BroadcastStyle just yet until we can get it working

* Defines linear indexing and fill! for fields

* Redefine checkbounds for AbstractField

* Starts implementing new checkbounds pending appropriate parentaxes function for fields

* BuoyancyField is backed by data

* Banishes parent() from preconditioned_conjugate_gradient_solver

* Import AbstractDataField into Buoyancy

* Specifies fill_halo_regions to OffsetArray for distributed models

* More general location inference for binary operations

* Add identity to list of unary operators

* Adds additional type parameter to zero field

* Revert "Adds additional type parameter to zero field"

This reverts commit 1a77150.

* Dont restrict i, j, k to integer

* Field getindex over data getindex

* Introduces "stubborn" binary operations

* Converts Broadcasted to AbstractOperations when destination is a field

* Cleans up doc strings and adapt_structure for AbstractOperations

* Fix some scoping issues with interpolation wrapper for at macro

* Move interpolate_operation definition to at.jl

* No need to throw errors for operations at Nothing location

* Computations with AveragedField might work?

* Adds incremental build up of AveragedField operations testing

* Dont interpolate unecessarily

* Bugfix in abstract operations tests

* Cleans up merge issues in AbstractOperations

* Use custom binary op function for getindex

* No more test skip in abstrsct operations tests!

* Holy typo

* Fix merge-added typo;

* Move broadcasting tests to simulation group

* Add new test group for abstract operations and broadcasting

* Updates testset description to include broadcasting

* Fixes superfluous include(at.jl)

* Infer whether broadcast requires interpolation or not

* More broadcasting tests including for interpolation

* Creates new binary operator functions to help type inference

* Broadcasting shenanigans

* Fixes bugs in broadcasting tests

* Finds BroadcastStyle name

* Defines fallback for interpolate_operation

* Polish off interpolation for broadcasting

* Updates broadcasting tests

* Field takes eltype from grid

* Test a few unaries but not all

* interpolate_identity is officially a unary

* Test explicit set of unaries

* Update unary_operations.jl

* Update kelvin_helmholtz_instability.jl

* Removes wrong comments

* Bump version

* No need for needs_interpolation

.buildkite/pipeline.yml

@@ -243,6 +243,33 @@ steps:
       architecture: CPU
     depends_on: "init_cpu"
 
+#####
+##### AbstractOperations
+#####
+
+  - label: "👻 gpu abstract operations tests"
+    env:
+      JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "abstract_operations"
+    commands:
+      - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: GPU
+    depends_on: "init_gpu"
+
+  - label: "🤖 cpu abstract operations tests"
+    env:
+      JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "abstract_operations"
+      CUDA_VISIBLE_DEVICES: "-1"
+    commands:
+      - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: CPU
+    depends_on: "init_cpu"
+
 #####
 ##### Cubed sphere
 #####

Project.toml

@@ -1,6 +1,6 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
-version = "0.55.0"
+version = "0.55.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

examples/kelvin_helmholtz_instability.jl

@@ -332,7 +332,7 @@ nothing # hide
 
 using Random, Statistics
 
-mean_perturbation_kinetic_energy = mean(1/2 * (u^2 + w^2), dims=(1, 2, 3))
+mean_perturbation_kinetic_energy = AveragedField(1/2 * (u^2 + w^2), dims=(1, 2, 3))
 
 noise(x, y, z) = randn()
 

src/AbstractOperations/AbstractOperations.jl

@@ -25,12 +25,7 @@ import Oceananigans.Fields: data, compute_at!
 ##### Basic functionality
 #####
 
-"""
-    AbstractOperation{X, Y, Z, G} <: AbstractField{X, Y, Z, Nothing, G}
-
-Represents an operation performed on grid of type `G` at locations `X`, `Y`, and `Z`.
-"""
-abstract type AbstractOperation{X, Y, Z, G} <: AbstractField{X, Y, Z, Nothing, G} end
+abstract type AbstractOperation{X, Y, Z, A, G, T} <: AbstractField{X, Y, Z, A, G, T} end
 
 const AF = AbstractField
 
@@ -40,13 +35,20 @@ Base.parent(op::AbstractOperation) = op
 # AbstractOperation macros add their associated functions to this list
 const operators = Set()
 
-include("grid_validation.jl")
+"""
+    at(loc, abstract_operation)
+
+Returns `abstract_operation` relocated to `loc`ation.
+"""
+at(loc, f) = f # fallback
 
+include("grid_validation.jl")
 include("unary_operations.jl")
 include("binary_operations.jl")
 include("multiary_operations.jl")
 include("derivatives.jl")
-
+include("at.jl")
+include("broadcasting_abstract_operations.jl")
 include("show_abstract_operations.jl")
 include("averages_of_operations.jl")
 
@@ -77,7 +79,5 @@ eval(define_multiary_operator(:*))
 push!(operators, :*)
 push!(multiary_operators, :*)
 
-include("at.jl")
-
 end # module
 

src/AbstractOperations/at.jl

@@ -23,12 +23,14 @@ end
 insert_location!(anything, location) = nothing
 
 # A very special UnaryOperation
-@inbounds identity(i, j, k, grid, a::Number) = a
-@inbounds identity(i, j, k, grid, a::AbstractField) = @inbounds a[i, j, k]
+@inline interpolate_identity(x) = x
+@unary interpolate_identity
+
+interpolate_operation(L, x) = x
 
 function interpolate_operation(L, x::AbstractField)
     L == location(x) && return x # Don't interpolate unecessarily
-    return _unary_operation(L, identity, x, location(x), x.grid)
+    return interpolate_identity(L, x)
 end
 
 """

src/AbstractOperations/averages_of_operations.jl

@@ -25,10 +25,3 @@ function AveragedField(op::AbstractOperation; dims, data=nothing, operand_data=n
     return AveragedField(computed, dims=dims, data=data, recompute_safely=recompute_safely)
 end
 
-"""
-    mean(op::AbstractOperation; kwargs...)
-
-Returns an Oceananigans.AveragedField representing the an average over `op`eration.
-See `Oceananigans.Fields.AveragedField`.
-"""
-Statistics.mean(op::AbstractOperation; kwargs...) = AveragedField(op; kwargs...)

src/AbstractOperations/binary_operations.jl

@@ -1,22 +1,23 @@
 const binary_operators = Set()
 
-struct BinaryOperation{X, Y, Z, O, A, B, IA, IB, G} <: AbstractOperation{X, Y, Z, G}
-      op :: O
-       a :: A
-       b :: B
-      ▶a :: IA
-      ▶b :: IB
-    grid :: G
+struct BinaryOperation{X, Y, Z, O, A, B, IA, IB, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T}
+              op :: O
+               a :: A
+               b :: B
+              ▶a :: IA
+              ▶b :: IB
+    architecture :: R
+            grid :: G
 
     """
-        BinaryOperation{X, Y, Z}(op, a, b, ▶a, ▶b, grid)
+        BinaryOperation{X, Y, Z}(op, a, b, ▶a, ▶b, arch, grid)
 
     Returns an abstract representation of the binary operation `op(▶a(a), ▶b(b))`.
-    where `▶a` and `▶b` interpolate `a` and `b` to (X, Y, Z).
+    on `grid` and `arch`itecture, where `▶a` and `▶b` interpolate `a` and `b` to (X, Y, Z).
     """
-    function BinaryOperation{X, Y, Z}(op, a, b, ▶a, ▶b, grid) where {X, Y, Z}
-        return new{X, Y, Z, typeof(op), typeof(a), typeof(b), typeof(▶a), typeof(▶b),
-                   typeof(grid)}(op, a, b, ▶a, ▶b, grid)
+    function BinaryOperation{X, Y, Z}(op::O, a::A, b::B, ▶a::IA, ▶b::IB, arch::R, grid::G) where {X, Y, Z, O, A, B, IA, IB, R, G}
+        T = eltype(grid)
+        return new{X, Y, Z, O, A, B, IA, IB, R, G, T}(op, a, b, ▶a, ▶b, arch, grid)
     end
 end
 
@@ -26,12 +27,16 @@ end
 ##### BinaryOperation construction
 #####
 
-"""Create a binary operation for `op` acting on `a` and `b` with locations `La` and `Lb`.
-The operator acts at `Lab` and the result is interpolated to `Lc`."""
+# Recompute location of binary operation
+@inline at(loc, β::BinaryOperation) = β.op(loc, at(loc, β.a), at(loc, β.b))
+
+"""Create a binary operation for `op` acting on `a` and `b` at `Lc`, where
+`a` and `b` have location `La` and `Lb`."""
 function _binary_operation(Lc, op, a, b, La, Lb, grid)
      ▶a = interpolation_operator(La, Lc)
      ▶b = interpolation_operator(Lb, Lc)
-    return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, a, b, ▶a, ▶b, grid)
+    arch = architecture(a, b)
+    return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, a, b, ▶a, ▶b, arch, grid)
 end
 
 const ConcreteLocationType = Union{Type{Face}, Type{Center}}
@@ -56,6 +61,34 @@ function define_binary_operator(op)
         @inline $op(i, j, k, grid::AbstractGrid, ▶a, ▶b, a, b) =
             @inbounds $op(▶a(i, j, k, grid, a), ▶b(i, j, k, grid, b))
 
+        # These shenanigans seem to help / encourage the compiler to infer types of objects
+        # buried in deep AbstractOperations trees.
+        @inline function $op(i, j, k, grid::AbstractGrid, ▶a, ▶b, A::BinaryOperation, B::BinaryOperation)
+            @inline a(ii, jj, kk, grid) = A.op(A.▶a(ii, jj, kk, grid, A.a), A.▶b(ii, jj, kk, grid, A.b))
+            @inline b(ii, jj, kk, grid) = B.op(B.▶a(ii, jj, kk, grid, B.a), B.▶b(ii, jj, kk, grid, B.b))
+            return @inbounds $op(▶a(i, j, k, grid, a), ▶b(i, j, k, grid, b))
+        end
+
+        @inline function $op(i, j, k, grid::AbstractGrid, ▶a, ▶b, A::BinaryOperation, B::AbstractField)
+            @inline a(ii, jj, kk, grid) = A.op(A.▶a(ii, jj, kk, grid, A.a), A.▶b(ii, jj, kk, grid, A.b))
+            return @inbounds $op(▶a(i, j, k, grid, a), ▶b(i, j, k, grid, B))
+        end
+
+        @inline function $op(i, j, k, grid::AbstractGrid, ▶a, ▶b, A::AbstractField, B::BinaryOperation)
+            @inline b(ii, jj, kk, grid) = B.op(B.▶a(ii, jj, kk, grid, B.a), B.▶b(ii, jj, kk, grid, B.b))
+            return @inbounds $op(▶a(i, j, k, grid, A), ▶b(i, j, k, grid, b))
+        end
+
+        @inline function $op(i, j, k, grid::AbstractGrid, ▶a, ▶b, A::BinaryOperation, B::Number)
+            @inline a(ii, jj, kk, grid) = A.op(A.▶a(ii, jj, kk, grid, A.a), A.▶b(ii, jj, kk, grid, A.b))
+            return @inbounds $op(▶a(i, j, k, grid, a), B)
+        end
+
+        @inline function $op(i, j, k, grid::AbstractGrid, ▶a, ▶b, A::Number, B::BinaryOperation)
+            @inline b(ii, jj, kk, grid) = B.op(B.▶a(ii, jj, kk, grid, B.a), B.▶b(ii, jj, kk, grid, B.b))
+            return @inbounds $op(A, ▶b(i, j, k, grid, b))
+        end
+
         """
             $($op)(Lc, a, b)
 
@@ -152,7 +185,7 @@ end
 ##### Architecture inference for BinaryOperation
 #####
 
-architecture(β::BinaryOperation) = architecture(β.a, β.b)
+architecture(β::BinaryOperation) = β.architecture
 
 function architecture(a, b)
     arch_a = architecture(a)
@@ -184,5 +217,5 @@ end
 "Adapt `BinaryOperation` to work on the GPU via CUDAnative and CUDAdrv."
 Adapt.adapt_structure(to, binary::BinaryOperation{X, Y, Z}) where {X, Y, Z} =
     BinaryOperation{X, Y, Z}(Adapt.adapt(to, binary.op), Adapt.adapt(to, binary.a),  Adapt.adapt(to, binary.b),
-                             Adapt.adapt(to, binary.▶a), Adapt.adapt(to, binary.▶b), Adapt.adapt(to, binary.grid))
-                             
+                             Adapt.adapt(to, binary.▶a), Adapt.adapt(to, binary.▶b), nothing, Adapt.adapt(to, binary.grid))
+

src/AbstractOperations/broadcasting_abstract_operations.jl

@@ -0,0 +1,16 @@
+import Oceananigans.Fields: broadcasted_to_abstract_operation
+
+using Base.Broadcast: Broadcasted
+using Base: identity
+
+const BroadcastedIdentity = Broadcasted{<:Any, <:Any, typeof(identity), <:Any}
+
+broadcasted_to_abstract_operation(loc, grid, bc::BroadcastedIdentity) =
+    interpolate_operation(loc, Tuple(broadcasted_to_abstract_operation(loc, grid, a) for a in bc.args)...)
+
+broadcasted_to_abstract_operation(loc, grid, op::AbstractOperation) = at(loc, op)
+
+function broadcasted_to_abstract_operation(loc, grid, bc::Broadcasted{<:Any, <:Any, <:Any, <:Any})
+    abstract_op = bc.f(loc, Tuple(broadcasted_to_abstract_operation(loc, grid, a) for a in bc.args)...)
+    return interpolate_operation(loc, abstract_op) # For "stubborn" BinaryOperations
+end

src/AbstractOperations/derivatives.jl

@@ -1,19 +1,21 @@
 using Oceananigans.Operators: interpolation_code
 
-struct Derivative{X, Y, Z, D, A, I, G} <: AbstractOperation{X, Y, Z, G}
-       ∂ :: D
-     arg :: A
-       ▶ :: I
-    grid :: G
+struct Derivative{X, Y, Z, D, A, I, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T}
+               ∂ :: D
+             arg :: A
+               ▶ :: I
+    architecture :: R
+            grid :: G
 
     """
         Derivative{X, Y, Z}(∂, arg, ▶, grid)
 
     Returns an abstract representation of the derivative `∂` on `arg`,
     and subsequent interpolation by `▶` on `grid`.
     """
-    function Derivative{X, Y, Z}(∂, arg, ▶, grid) where {X, Y, Z}
-        return new{X, Y, Z, typeof(∂), typeof(arg), typeof(▶), typeof(grid)}(∂, arg, ▶, grid)
+    function Derivative{X, Y, Z}(∂::D, arg::A, ▶::I, arch::R, grid::G) where {X, Y, Z, D, A, I, R, G}
+        T = eltype(grid)
+        return new{X, Y, Z, D, A, I, R, G, T}(∂, arg, ▶, arch, grid)
     end
 end
 
@@ -27,9 +29,13 @@ end
 interpolation to `L` on `grid`."""
 function _derivative(L, ∂, arg, L∂, grid) where {X, Y, Z}
     ▶ = interpolation_operator(L∂, L)
-    return Derivative{L[1], L[2], L[3]}(∂, arg, ▶, grid)
+    arch = architecture(arg)
+    return Derivative{L[1], L[2], L[3]}(∂, arg, ▶, arch, grid)
 end
 
+# Recompute location of derivative
+@inline at(loc, d::Derivative) = d.∂(loc, at(loc, d.arg))
+
 """Return `Center` if given `Face` or `Face` if given `Center`."""
 flip(::Type{Face}) = Center
 flip(::Type{Center}) = Face
@@ -100,7 +106,7 @@ Return an abstract representation of a z-derivative acting on `a`.
 ##### Architecture inference for derivatives
 #####
 
-architecture(∂::Derivative) = architecture(∂.arg)
+architecture(∂::Derivative) = ∂.architecture
 
 #####
 ##### Nested computations
@@ -117,4 +123,5 @@ compute_at!(∂::Derivative, time) = compute_at!(∂.arg, time)
 "Adapt `Derivative` to work on the GPU via CUDAnative and CUDAdrv."
 Adapt.adapt_structure(to, deriv::Derivative{X, Y, Z}) where {X, Y, Z} =
     Derivative{X, Y, Z}(Adapt.adapt(to, deriv.∂), Adapt.adapt(to, deriv.arg),
-                        Adapt.adapt(to, deriv.▶), Adapt.adapt(to, deriv.grid))
+                        Adapt.adapt(to, deriv.▶), nothing, Adapt.adapt(to, deriv.grid))
+

src/AbstractOperations/multiary_operations.jl

@@ -1,13 +1,16 @@
 const multiary_operators = Set()
 
-struct MultiaryOperation{X, Y, Z, N, O, A, I, G} <: AbstractOperation{X, Y, Z, G}
-      op :: O
-    args :: A
-       ▶ :: I
-    grid :: G
-
-    function MultiaryOperation{X, Y, Z}(op, args, ▶, grid) where {X, Y, Z}
-        return new{X, Y, Z, length(args), typeof(op), typeof(args), typeof(▶), typeof(grid)}(op, args, ▶, grid)
+struct MultiaryOperation{X, Y, Z, N, O, A, I, R, G, T} <: AbstractOperation{X, Y, Z, R, G, T}
+              op :: O
+            args :: A
+               ▶ :: I
+    architecture :: R
+            grid :: G
+
+    function MultiaryOperation{X, Y, Z}(op::O, args::A, ▶::I, arch::R, grid::G) where {X, Y, Z, O, A, I, R, G}
+        T = eltype(grid)
+        N = length(args)
+        return new{X, Y, Z, N, O, A, I, R, G, T}(op, args, ▶, arch, grid)
     end
 end
 
@@ -20,17 +23,21 @@ end
 
 function _multiary_operation(L, op, args, Largs, grid) where {X, Y, Z}
     ▶ = Tuple(interpolation_operator(La, L) for La in Largs)
-    return MultiaryOperation{L[1], L[2], L[3]}(op, Tuple(a for a in args), ▶, grid)
+    arch = architecture(args...)
+    return MultiaryOperation{L[1], L[2], L[3]}(op, Tuple(a for a in args), ▶, arch, grid)
 end
 
+# Recompute location of multiary operation
+@inline at(loc, Π::MultiaryOperation) = Π.op(loc, Tuple(at(loc, a) for a in Π.args)...)
+
 """Return an expression that defines an abstract `MultiaryOperator` named `op` for `AbstractField`."""
 function define_multiary_operator(op)
     return quote
         function $op(Lop::Tuple,
-                     a::Union{Function, Oceananigans.Fields.AbstractField},
-                     b::Union{Function, Oceananigans.Fields.AbstractField},
-                     c::Union{Function, Oceananigans.Fields.AbstractField},
-                     d::Union{Function, Oceananigans.Fields.AbstractField}...)
+                     a::Union{Function, Number, Oceananigans.Fields.AbstractField},
+                     b::Union{Function, Number, Oceananigans.Fields.AbstractField},
+                     c::Union{Function, Number, Oceananigans.Fields.AbstractField},
+                     d::Union{Function, Number, Oceananigans.Fields.AbstractField}...)
 
             args = tuple(a, b, c, d...)
             grid = Oceananigans.AbstractOperations.validate_grid(args...)
@@ -42,7 +49,7 @@ function define_multiary_operator(op)
             return Oceananigans.AbstractOperations._multiary_operation(Lop, $op, args, Largs, grid)
         end
 
-        $op(a::Union{Function, Oceananigans.Fields.AbstractField},
+        $op(a::Oceananigans.Fields.AbstractField,
             b::Union{Function, Oceananigans.Fields.AbstractField},
             c::Union{Function, Oceananigans.Fields.AbstractField},
             d::Union{Function, Oceananigans.Fields.AbstractField}...) = $op(Oceananigans.Fields.location(a), a, b, c, d...)
@@ -127,7 +134,7 @@ end
 ##### Architecture inference for MultiaryOperation
 #####
 
-architecture(Π::MultiaryOperation) = architecture(Π.args...)
+architecture(Π::MultiaryOperation) = Π.architecture
 
 function architecture(a, b, c, d...)
 
@@ -162,4 +169,5 @@ end
 "Adapt `MultiaryOperation` to work on the GPU via CUDAnative and CUDAdrv."
 Adapt.adapt_structure(to, multiary::MultiaryOperation{X, Y, Z}) where {X, Y, Z} =
     MultiaryOperation{X, Y, Z}(Adapt.adapt(to, multiary.op), Adapt.adapt(to, multiary.args),
-                               Adapt.adapt(to, multiary.▶),  Adapt.adapt(to, multiary.grid))
+                               Adapt.adapt(to, multiary.▶),  nothing, Adapt.adapt(to, multiary.grid))
+