Adapt Field, AveragedField, and ComputedField for GPU, round 2

src/AbstractOperations/AbstractOperations.jl

@@ -15,7 +15,6 @@ using Oceananigans.Grids
 using Oceananigans.Operators
 using Oceananigans.BoundaryConditions
 using Oceananigans.Fields
-using Oceananigans.Fields: gpufriendly
 
 using Oceananigans.Operators: interpolation_operator
 using Oceananigans.Architectures: device

src/AbstractOperations/binary_operations.jl

@@ -43,7 +43,7 @@ function _binary_operation(Lc, op, a, b, La, Lb, Lab, grid)
      ▶a = interpolation_operator(La, Lab)
      ▶b = interpolation_operator(Lb, Lab)
     ▶op = interpolation_operator(Lab, Lc)
-    return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, gpufriendly(a), gpufriendly(b), ▶a, ▶b, ▶op, grid)
+    return BinaryOperation{Lc[1], Lc[2], Lc[3]}(op, a, b, ▶a, ▶b, ▶op, grid)
 end
 
 """Return an expression that defines an abstract `BinaryOperator` named `op` for `AbstractField`."""
@@ -105,6 +105,8 @@ Example
 =======
 
 ```jldoctest
+julia> using Oceananigans, Oceananigans.AbstractOperations, Oceananigans.Grids
+
 julia> plus_or_times(x, y) = x < 0 ? x + y : x * y
 plus_or_times (generic function with 1 method)
 
@@ -117,18 +119,16 @@ julia> @binary plus_or_times
  :*
  :plus_or_times
 
-julia> c, d = (Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid((1, 1, 16), (1, 1, 1))) for i = 1:2);
+julia> c, d = (Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:2);
 
 julia> plus_or_times(c, d)
 BinaryOperation at (Cell, Cell, Cell)
-├── grid: RegularCartesianGrid{Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
-│   ├── size: (1, 1, 16)
-│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [0.0, -1.0]
+├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
+│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
 └── tree:
-
-plus_or_times at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
-├── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
-└── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
+    plus_or_times at (Cell, Cell, Cell) via identity
+    ├── Field located at (Cell, Cell, Cell)
+    └── Field located at (Cell, Cell, Cell)
 """
 macro binary(ops...)
     expr = Expr(:block)

src/AbstractOperations/derivatives.jl

@@ -32,7 +32,7 @@ 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]}(∂, gpufriendly(arg), ▶, grid)
+    return Derivative{L[1], L[2], L[3]}(∂, arg, ▶, grid)
 end
 
 """Return `Cell` if given `Face` or `Face` if given `Cell`."""

src/AbstractOperations/multiary_operations.jl

@@ -20,7 +20,7 @@ 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(gpufriendly(a) for a in args), ▶, grid)
+    return MultiaryOperation{L[1], L[2], L[3]}(op, Tuple(a for a in args), ▶, grid)
 end
 
 """Return an expression that defines an abstract `MultiaryOperator` named `op` for `AbstractField`."""
@@ -64,49 +64,46 @@ Example
 =======
 
 ```jldoctest
+julia> using Oceananigans, Oceananigans.Grids, Oceananigans.AbstractOperations
+
 julia> harmonic_plus(a, b, c) = 1/3 * (1/a + 1/b + 1/c)
 harmonic_plus(generic function with 1 method)
 
-julia> @multiary harmonic_plus
-3-element Array{Any,1}:
- :+
- :*
- :harmonic_plus
-
-julia> c, d, e = Tuple(Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid((1, 1, 16), (1, 1, 1))) for i = 1:3);
+julia> c, d, e = Tuple(Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:3);
 
-julia> harmonic_plus(c, d, e) # this calls the original function, which in turn returns a (correct) operation tree
+julia> harmonic_plus(c, d, e) # before magic @multiary transformation
 BinaryOperation at (Cell, Cell, Cell)
-├── grid: RegularCartesianGrid{Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
-│   ├── size: (1, 1, 16)
-│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [0.0, -1.0]
+├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
+│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
 └── tree:
-
-* at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
-├── 0.3333333333333333
-└── + at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
-    ├── + at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
-    │   ├── / at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
-    │   │   ├── 1
-    │   │   └── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
-    │   └── / at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
+    * at (Cell, Cell, Cell) via identity
+    ├── 0.3333333333333333
+    └── + at (Cell, Cell, Cell)
+        ├── / at (Cell, Cell, Cell) via identity
         │   ├── 1
-        │   └── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
-    └── / at (Cell, Cell, Cell) via Oceananigans.AbstractOperations.identity
-        ├── 1
-        └── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
+        │   └── Field located at (Cell, Cell, Cell)
+        ├── / at (Cell, Cell, Cell) via identity
+        │   ├── 1
+        │   └── Field located at (Cell, Cell, Cell)
+        └── / at (Cell, Cell, Cell) via identity
+            ├── 1
+            └── Field located at (Cell, Cell, Cell)
+
+julia> @multiary harmonic_plus
+Set{Any} with 3 elements:
+  :+
+  :harmonic_plus
+  :*
 
-julia> @at (Cell, Cell, Cell) harmonic_plus(c, d, e) # this returns a `MultiaryOperation` as expected
+julia> harmonic_plus(c, d, e)
 MultiaryOperation at (Cell, Cell, Cell)
-├── grid: RegularCartesianGrid{Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
-│   ├── size: (1, 1, 16)
-│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [0.0, -1.0]
+├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
+│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
 └── tree:
-
-harmonic_plus at (Cell, Cell, Cell)
-├── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
-├── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
-└── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
+    harmonic_plus at (Cell, Cell, Cell)
+    ├── Field located at (Cell, Cell, Cell)
+    ├── Field located at (Cell, Cell, Cell)
+    └── Field located at (Cell, Cell, Cell)
 """
 macro multiary(ops...)
     expr = Expr(:block)

src/AbstractOperations/show_abstract_operations.jl

@@ -1,4 +1,4 @@
-using Oceananigans: short_show
+import Oceananigans: short_show
 using Oceananigans.Grids: domain_string
 using Oceananigans.Fields: show_location
 
@@ -9,11 +9,12 @@ for op_string in ("UnaryOperation", "BinaryOperation", "MultiaryOperation", "Der
     end
 end
 
+short_show(operation::AbstractOperation) = string(operation_name(operation), " at ", show_location(operation))
+
 Base.show(io::IO, operation::AbstractOperation) =
     print(io,
-          operation_name(operation), " at ", show_location(operation), '\n',
-          "├── grid: ", typeof(operation.grid), '\n',
-          "│   ├── size: ", size(operation.grid), '\n',
+          short_show(operation), '\n',
+          "├── grid: ", short_show(operation.grid), '\n',
           "│   └── domain: ", domain_string(operation.grid), '\n',
           "└── tree: ", "\n", "    ", tree_show(operation, 1, 0))
 

src/AbstractOperations/unary_operations.jl

@@ -33,7 +33,7 @@ end
 result from `Larg` to `L`."""
 function _unary_operation(L, operator, arg, Larg, grid)
     ▶ = interpolation_operator(Larg, L)
-    return UnaryOperation{L[1], L[2], L[3]}(operator, gpufriendly(arg), ▶, grid)
+    return UnaryOperation{L[1], L[2], L[3]}(operator, arg, ▶, grid)
 end
 
 """
@@ -49,30 +49,32 @@ Also note: a unary function in `Base` must be imported to be extended: use `impo
 Example
 =======
 
+```jldoctest
+julia> using Oceananigans, Oceananigans.Grids, Oceananigans.AbstractOperations
+
 julia> square_it(x) = x^2
 square_it (generic function with 1 method)
 
 julia> @unary square_it
-7-element Array{Any,1}:
- :sqrt
- :sin
- :cos
- :exp
- :tanh
- :-
- :square_it
+Set{Any} with 7 elements:
+  :sqrt
+  :square_it
+  :cos
+  :exp
+  :-
+  :tanh
+  :sin
 
-julia> c = Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid((1, 1, 16), (1, 1, 1)));
+julia> c = Field(Cell, Cell, Cell, CPU(), RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1)));
 
 julia> square_it(c)
 UnaryOperation at (Cell, Cell, Cell)
-├── grid: RegularCartesianGrid{Float64,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
-│   ├── size: (1, 1, 16)
-│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [0.0, -1.0]
+├── grid: RegularCartesianGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
+│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
 └── tree:
-
-square_it at (Cell, Cell, Cell) via identity
-└── OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}}
+    square_it at (Cell, Cell, Cell) via identity
+    └── Field located at (Cell, Cell, Cell)
+```
 """
 macro unary(ops...)
     expr = Expr(:block)

src/Buoyancy/buoyancy_field.jl

@@ -1,12 +1,14 @@
 using Adapt
 using KernelAbstractions
 using Oceananigans.Fields: AbstractField, FieldStatus, validate_field_data, new_data, conditional_compute!
-using Oceananigans.Fields: datatuple, architecture
+using Oceananigans.Fields: architecture, tracernames
 using Oceananigans.Architectures: device
 using Oceananigans.Utils: work_layout
 
 import Oceananigans.Fields: compute!
 
+import Oceananigans: short_show
+
 """
     struct BuoyancyField{B, A, G, T} <: AbstractField{X, Y, Z, A, G}
 
@@ -29,8 +31,6 @@ struct BuoyancyField{B, S, A, G, T} <: AbstractField{Cell, Cell, Cell, A, G}
 
         validate_field_data(Cell, Cell, Cell, data, grid)
 
-        tracers = datatuple(tracers)
-
         status = recompute_safely ? nothing : FieldStatus(zero(eltype(grid)))
 
         return new{typeof(buoyancy), typeof(status), typeof(data),
@@ -39,7 +39,6 @@ struct BuoyancyField{B, S, A, G, T} <: AbstractField{Cell, Cell, Cell, A, G}
 
     function BuoyancyField(data, grid, buoyancy, tracers, status)
         validate_field_data(Cell, Cell, Cell, data, grid)
-        tracers = datatuple(tracers)
         return new{typeof(buoyancy), typeof(status), typeof(data),
                    typeof(grid), typeof(tracers)}(data, grid, buoyancy, tracers, status)
     end
@@ -120,9 +119,18 @@ end
 ##### Adapt
 #####
 
-Adapt.adapt_structure(to, buoyancy_field::BuoyancyField) =
-    BuoyancyField(Adapt.adapt(to, buoyancy_field.data),
-                  Adapt.adapt(to, buoyancy_field.grid),
-                  Adapt.adapt(to, buoyancy_field.buoyancy),
-                  Adapt.adapt(to, buoyancy_field.tracers),
-                  Adapt.adapt(to, buoyancy_field.status))
+Adapt.adapt_structure(to, buoyancy_field::BuoyancyField) = Adapt.adapt(to, buoyancy_field.data)
+
+#####
+##### Show
+#####
+
+short_show(field::BuoyancyField) = string("BuoyancyField for ", typeof(field.buoyancy))
+
+show(io::IO, field::BuoyancyField) =
+    print(io, "$(short_show(field))\n",
+          "├── data: $(typeof(field.data)), size: $(size(field.data))\n",
+          "├── grid: $(short_show(field.grid))", '\n',
+          "├── buoyancy: $(typeof(field.buoyancy))", '\n',
+          "├── tracers: $(tracernames(field.tracers))", '\n',
+          "└── status: ", show_status(field.status), '\n')

src/Buoyancy/nonlinear_equation_of_state.jl

@@ -1,3 +1,5 @@
+using Oceananigans.Fields: AbstractField
+
 """ Return the geopotential height at `i, j, k` at cell centers. """
 @inline function Zᵃᵃᶜ(i, j, k, grid::AbstractGrid{FT}) where FT 
     @inbounds begin
@@ -24,11 +26,13 @@ end
     end
 end
 
+const ArrayOrField = Union{AbstractArray, AbstractField}
+
 # Dispatch shenanigans
-@inline θ_and_sᴬ(i, j, k, θ::AbstractArray, sᴬ::AbstractArray) = @inbounds θ[i, j, k], sᴬ[i, j, k]
-@inline θ_and_sᴬ(i, j, k, θ::Number,        sᴬ::AbstractArray) = @inbounds θ, sᴬ[i, j, k]
-@inline θ_and_sᴬ(i, j, k, θ::AbstractArray, sᴬ::Number)        = @inbounds θ[i, j, k], sᴬ
-@inline θ_and_sᴬ(i, j, k, θ::Number,        sᴬ::Number)        = @inbounds θ, sᴬ
+@inline θ_and_sᴬ(i, j, k, θ::ArrayOrField, sᴬ::ArrayOrField) = @inbounds θ[i, j, k], sᴬ[i, j, k]
+@inline θ_and_sᴬ(i, j, k, θ::Number,       sᴬ::ArrayOrField) = @inbounds θ, sᴬ[i, j, k]
+@inline θ_and_sᴬ(i, j, k, θ::ArrayOrField, sᴬ::Number)       = @inbounds θ[i, j, k], sᴬ
+@inline θ_and_sᴬ(i, j, k, θ::Number,       sᴬ::Number)       = @inbounds θ, sᴬ
 
 # Basic functionality
 @inline ρ′(i, j, k, grid, eos, θ, sᴬ) = @inbounds ρ′(θ_and_sᴬ(i, j, k, θ, sᴬ)..., Zᵃᵃᶜ(i, j, k, grid), eos)

src/Fields/abstract_field.jl

@@ -7,9 +7,9 @@ using Oceananigans.Utils
 using Oceananigans.Grids: interior_indices, interior_parent_indices
 
 import Oceananigans: location
-import Oceananigans.Utils: datatuple
 import Oceananigans.Architectures: architecture
 import Oceananigans.Grids: total_size, topology, nodes, xnodes, ynodes, znodes, xnode, ynode, znode
+import Oceananigans.Utils: datatuple
 
 """
     AbstractField{X, Y, Z, A, G}
@@ -31,6 +31,9 @@ function validate_field_data(X, Y, Z, data, grid)
     return nothing
 end
 
+# Endpoint for recursive `datatuple` function:
+@inline datatuple(obj::AbstractField) = data(obj)
+
 #####
 ##### Computing AbstractField
 #####
@@ -144,12 +147,6 @@ total_size(f::AbstractField) = total_size(location(f), f.grid)
 @hascuda @inline cpudata(f::AbstractField{X, Y, Z, <:OffsetCuArray}) where {X, Y, Z} =
     offset_data(Array(parent(f)), f.grid, location(f))
 
-# Endpoint for recursive `datatuple` function:
-@inline datatuple(obj::AbstractField) = data(obj)
-
-""" Converts a field into a GPU-friendly alternative if necessary. """
-@inline gpufriendly(a) = a # fallback
-
 "Returns `f.data.parent` for `f::Field`."
 @inline Base.parent(f::AbstractField) = parent(data(f))
 

src/Fields/averaged_field.jl

@@ -29,14 +29,8 @@ struct AveragedField{X, Y, Z, S, A, G, N, O} <: AbstractReducedField{X, Y, Z, A,
     end
 
     function AveragedField{X, Y, Z}(data, grid, dims, operand, status) where {X, Y, Z}
-
-        dims = validate_reduced_dims(dims)
-        validate_reduced_locations(X, Y, Z, dims)
-        validate_field_data(X, Y, Z, data, grid)
-
         return new{X, Y, Z, typeof(status), typeof(data),
-                   typeof(grid), length(dims), typeof(operand)}(data, grid, dims,
-                                                                operand, status)
+                   typeof(grid), length(dims), typeof(operand)}(data, grid, dims, operand, status)
     end
 end
 
@@ -94,7 +88,4 @@ Statistics.mean(ϕ::AbstractField; kwargs...) = AveragedField(ϕ; kwargs...)
 
 Adapt.adapt_structure(to, averaged_field::AveragedField{X, Y, Z}) where {X, Y, Z} = 
     AveragedField{X, Y, Z}(Adapt.adapt(to, averaged_field.data),
-                           Adapt.adapt(to, averaged_field.grid),
-                           Adapt.adapt(to, averaged_field.dims),
-                           Adapt.adapt(to, averaged_field.operand),
-                           Adapt.adapt(to, averaged_field.status))
+                           nothing, averaged_field.dims, nothing, nothing)

src/Fields/computed_field.jl

@@ -91,8 +91,4 @@ end
 ##### Adapt
 #####
 
-Adapt.adapt_structure(to, computed_field::ComputedField{X, Y, Z}) where {X, Y, Z} = 
-    ComputedField{X, Y, Z}(Adapt.adapt(to, computed_field.data),
-                           Adapt.adapt(to, computed_field.grid),
-                           Adapt.adapt(to, computed_field.operand),
-                           Adapt.adapt(to, computed_field.status))
+Adapt.adapt_structure(to, computed_field::ComputedField) = Adapt.adapt(to, computed_field.data)

src/Fields/field.jl

@@ -126,4 +126,4 @@ ZFaceField(arch::AbstractArchitecture, grid, args...) = ZFaceField(eltype(grid),
 
 @propagate_inbounds Base.setindex!(f::Field, v, inds...) = @inbounds setindex!(f.data, v, inds...)
 
-gpufriendly(f::Field) = data(f)
+Adapt.adapt_structure(to, field::Field) = Adapt.adapt(to, field.data)