Revert "Improved fill_halo_regions"

src/BoundaryConditions/fill_halo_regions.jl

@@ -1,33 +1,29 @@
-using Oceananigans.Architectures: architecture
-
 #####
 ##### General halo filling functions
 #####
 
 fill_halo_regions!(::Nothing, args...) = []
 
 """
-    fill_halo_regions!(fields)
+    fill_halo_regions!(fields, arch)
 
 Fill halo regions for each field in the tuple `fields` according to their boundary
 conditions, possibly recursing into `fields` if it is a nested tuple-of-tuples.
 """
-function fill_halo_regions!(fields::Union{Tuple, NamedTuple}, args...)
+function fill_halo_regions!(fields::Union{Tuple, NamedTuple}, arch, args...)
 
     for field in fields
-        fill_halo_regions!(field, args...)
+        fill_halo_regions!(field, arch, args...)
     end
 
     return nothing
 end
 
-fill_halo_regions!(field, args...) =
-    fill_halo_regions!(field.data, field.boundary_conditions, field.grid, args...)
+fill_halo_regions!(field, arch, args...) = fill_halo_regions!(field.data, field.boundary_conditions, arch, field.grid, args...)
 
 "Fill halo regions in x, y, and z for a given field."
-function fill_halo_regions!(c::AbstractArray, fieldbcs, grid, args...; kwargs...)
+function fill_halo_regions!(c::AbstractArray, fieldbcs, arch, grid, args...; kwargs...)
 
-    arch = architecture(c)
     barrier = Event(device(arch))
 
       west_event =   fill_west_halo!(c, fieldbcs.west,   arch, barrier, grid, args...; kwargs...)

src/BoundaryConditions/fill_halo_regions_normal_flow.jl

@@ -2,19 +2,19 @@
 ##### Outer functions for setting velocity on boundary and filling halo beyond boundary.
 #####
 
-@kernel function set_east_west_u_velocity!(u, i_boundary, bc, grid, args...)
+@kernel function set_east_west_u_velocity!(u, i_boundary, bc, grid, clock, model_fields)
     j, k = @index(Global, NTuple)
-    @inbounds u[i_boundary, j, k] = getbc(bc, j, k, grid, args...)
+    @inbounds u[i_boundary, j, k] = getbc(bc, j, k, grid, clock, model_fields)
 end
 
-@kernel function set_north_south_v_velocity!(v, j_boundary, bc, grid, args...)
+@kernel function set_north_south_v_velocity!(v, j_boundary, bc, grid, clock, model_fields)
     i, k = @index(Global, NTuple)
-    @inbounds v[i, j_boundary, k] = getbc(bc, i, k, grid, args...)
+    @inbounds v[i, j_boundary, k] = getbc(bc, i, k, grid, clock, model_fields)
 end
 
-@kernel function set_top_bottom_w_velocity!(w, k_boundary, bc, grid, args...)
+@kernel function set_top_bottom_w_velocity!(w, k_boundary, bc, grid, clock, model_fields)
     i, j = @index(Global, NTuple)
-    @inbounds w[i, j, k_boundary] = getbc(bc, i, j, grid, args...)
+    @inbounds w[i, j, k_boundary] = getbc(bc, i, j, grid, clock, model_fields)
 end
 
   fill_west_halo!(u, bc::NFBC, arch, dep, grid, args...; kwargs...) = launch!(arch, grid, :yz, set_east_west_u_velocity!,   u,           1, bc, grid, args...; dependencies=dep, kwargs...)

src/Fields/computed_field.jl

@@ -100,7 +100,7 @@ function compute!(comp::ComputedField{X, Y, Z}, time=nothing) where {X, Y, Z}
 
     wait(device(arch), event)
 
-    fill_halo_regions!(comp)
+    fill_halo_regions!(comp, arch)
 
     return nothing
 end

src/Fields/function_field.jl

@@ -1,5 +1,3 @@
-import Oceananigans.Architectures: architecture
-
 """
     FunctionField{X, Y, Z, C, F, G} <: AbstractField{X, Y, Z, F, G}
 

src/Fields/kernel_computed_field.jl

@@ -124,7 +124,7 @@ function compute!(kcf::KernelComputedField{X, Y, Z}) where {X, Y, Z}
 
     wait(device(arch), event)
 
-    fill_halo_regions!(kcf)
+    fill_halo_regions!(kcf, arch)
 
     return nothing
 end

src/Fields/reduced_field.jl

@@ -1,5 +1,3 @@
-using Oceananigans.Architectures: architecture
-
 using Adapt
 
 import Oceananigans.BoundaryConditions: fill_halo_regions!
@@ -33,8 +31,8 @@ const ARF = AbstractReducedField
 @inline Base.getindex( r::ARF{Nothing, Nothing, Nothing},    i, j, k) = @inbounds r.data[1, 1, 1]
 @inline Base.setindex!(r::ARF{Nothing, Nothing, Nothing}, d, i, j, k) = @inbounds r.data[1, 1, 1] = d
 
-fill_halo_regions!(field::AbstractReducedField, args...) =
-    fill_halo_regions!(field.data, field.boundary_conditions, field.grid, args...; reduced_dimensions=field.dims)
+fill_halo_regions!(field::AbstractReducedField, arch, args...) =
+    fill_halo_regions!(field.data, field.boundary_conditions, arch, field.grid, args...; reduced_dimensions=field.dims)
 
 const DimsType = NTuple{N, Int} where N
 

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -13,17 +13,17 @@ Update peripheral aspects of the model (halo regions, diffusivities, hydrostatic
 function update_state!(model::HydrostaticFreeSurfaceModel)
 
     # Fill halos for velocities and tracers
-    fill_halo_regions!(fields(model), model.clock, fields(model))
+    fill_halo_regions!(fields(model), model.architecture, model.clock, fields(model))
 
     compute_w_from_continuity!(model)
 
-    fill_halo_regions!(model.velocities.w, model.clock, fields(model))
+    fill_halo_regions!(model.velocities.w, model.architecture, model.clock, fields(model))
 
     # Calculate diffusivities
     calculate_diffusivities!(model.diffusivities, model.architecture, model.grid, model.closure,
                              model.buoyancy, model.velocities, model.tracers)
 
-    fill_halo_regions!(model.diffusivities, model.clock, fields(model))
+    fill_halo_regions!(model.diffusivities, model.architecture, model.clock, fields(model))
 
     # Calculate hydrostatic pressure
     pressure_calculation = launch!(model.architecture, model.grid, :xy, update_hydrostatic_pressure!,
@@ -33,7 +33,7 @@ function update_state!(model::HydrostaticFreeSurfaceModel)
     # Fill halo regions for pressure
     wait(device(model.architecture), pressure_calculation)
 
-    fill_halo_regions!(model.pressure.pHY′)
+    fill_halo_regions!(model.pressure.pHY′, model.architecture)
 
     return nothing
 end

src/Models/IncompressibleModels/pressure_correction.jl

@@ -9,11 +9,11 @@ Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `vel
 """
 function calculate_pressure_correction!(model::IncompressibleModel, Δt)
 
-    fill_halo_regions!(model.velocities, model.clock, fields(model))
+    fill_halo_regions!(model.velocities, model.architecture, model.clock, fields(model))
 
     solve_for_pressure!(model.pressures.pNHS, model.pressure_solver, model.architecture, model.grid, Δt, model.velocities)
 
-    fill_halo_regions!(model.pressures.pNHS)
+    fill_halo_regions!(model.pressures.pNHS, model.architecture)
 
     return nothing
 end

src/Models/IncompressibleModels/update_incompressible_model_state.jl

@@ -12,13 +12,14 @@ Update peripheral aspects of the model (halo regions, diffusivities, hydrostatic
 function update_state!(model::IncompressibleModel)
 
     # Fill halos for velocities and tracers
-    fill_halo_regions!(merge(model.velocities, model.tracers), model.clock, fields(model))
+    fill_halo_regions!(merge(model.velocities, model.tracers), model.architecture, 
+                       model.clock, fields(model))
 
     # Calculate diffusivities
     calculate_diffusivities!(model.diffusivities, model.architecture, model.grid, model.closure,
                              model.buoyancy, model.velocities, model.tracers)
 
-    fill_halo_regions!(model.diffusivities, model.clock, fields(model))
+    fill_halo_regions!(model.diffusivities, model.architecture, model.clock, fields(model))
 
     # Calculate hydrostatic pressure
     pressure_calculation = launch!(model.architecture, model.grid, :xy, update_hydrostatic_pressure!,
@@ -28,7 +29,7 @@ function update_state!(model::IncompressibleModel)
     # Fill halo regions for pressure
     wait(device(model.architecture), pressure_calculation)
 
-    fill_halo_regions!(model.pressures.pHY′)
+    fill_halo_regions!(model.pressures.pHY′, model.architecture)
 
     return nothing
 end

src/Models/ShallowWaterModels/update_shallow_water_state.jl

@@ -11,6 +11,7 @@ function update_state!(model::ShallowWaterModel)
 
     # Fill halos for velocities and tracers
     fill_halo_regions!(merge(model.solution, model.tracers),
+                       model.architecture,
                        model.clock,
                        fields(model))
 

src/Solvers/preconditioned_conjugate_gradient_solver.jl

@@ -214,8 +214,7 @@ function solve_poisson_equation!(solver::PreconditionedConjugateGradientSolver,
 ==#
       println("PreconditionedConjugateGradientSolver ", i," ", norm(r.parent) )
 
-    fill_halo_regions!(x, sset.bcs, sset.grid)
-
+    fill_halo_regions!(x, sset.bcs, sset.arch, sset.grid)
     return x, norm(r.parent)
 end
 

test/test_averaged_field.jl

@@ -27,10 +27,10 @@ using Oceananigans.Grids: halo_size
 
                 # Note: halo regions must be *filled* prior to computing an average
                 # if the average within halo regions is to be correct.
-                fill_halo_regions!(T)
+                fill_halo_regions!(T, arch)
                 @compute T̅ = AveragedField(T, dims=(1, 2))
 
-                fill_halo_regions!(T)
+                fill_halo_regions!(T, arch)
                 @compute T̂ = AveragedField(T, dims=1)
 
                 @compute w̃ = AveragedField(w, dims=(1, 2, 3))

test/test_boundary_conditions_integration.jl

@@ -125,7 +125,6 @@ test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType),
 @testset "Boundary condition integration tests" begin
     @info "Testing boundary condition integration into IncompressibleModel..."
 
-    #=
     @testset "Boundary condition regularization" begin
         @info "  Testing boundary condition regularization in IncompressibleModel constructor..."
 
@@ -191,7 +190,6 @@ test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType),
         @test location(model.tracers.T.boundary_conditions.north.condition)  == (Center, Nothing, Center)
         @test location(model.tracers.T.boundary_conditions.south.condition)  == (Center, Nothing, Center)
     end
-    =#
 
     @testset "Boudnary condition time-stepping works" begin
         for arch in archs, FT in (Float64,) #float_types
@@ -215,7 +213,6 @@ test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType),
         end
     end
 
-    #=
     @testset "Budgets with Flux boundary conditions" begin
         for arch in archs, FT in float_types
             @info "  Testing budgets with Flux boundary conditions on u, v, T [$(typeof(arch)), $FT]..."
@@ -231,5 +228,4 @@ test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType),
             @test fluxes_with_diffusivity_boundary_conditions_are_correct(arch, FT)
         end
     end
-    =#
 end

test/test_halo_regions.jl

@@ -27,7 +27,7 @@ function halo_regions_correctly_filled(arch, FT, Nx, Ny, Nz)
     field = CenterField(FT, arch, grid)
 
     set!(field, rand(FT, Nx, Ny, Nz))
-    fill_halo_regions!(field)
+    fill_halo_regions!(field, arch)
 
     Hx, Hy, Hz = grid.Hx, grid.Hy, grid.Hz
     data = field.data

test/test_operators.jl

@@ -138,7 +138,7 @@ end
 
         A3 = OffsetArray(zeros(Tx, Ty, Tz), 1-Hx:Nx+Hx, 1-Hy:Ny+Hy, 1-Hz:Nz+Hz)
         @. @views A3[1:Nx, 1:Ny, 1:Nz] = rand()
-        fill_halo_regions!(A3, TracerBoundaryConditions(grid), grid)
+        fill_halo_regions!(A3, TracerBoundaryConditions(grid), arch, grid)
 
         # A yz-slice with Nx==1.
         A2yz = OffsetArray(zeros(1+2Hx, Ty, Tz), 1-Hx:1+Hx, 1-Hy:Ny+Hy, 1-Hz:Nz+Hz)

test/test_poisson_solvers.jl

@@ -20,9 +20,9 @@ function random_divergent_source_term(FT, arch, grid)
     set!(Rw, rand(Nx, Ny, Nz))
 
     # Adding (nothing, nothing) in case we need to dispatch on ::NFBC
-    fill_halo_regions!(Ru)
-    fill_halo_regions!(Rv)
-    fill_halo_regions!(Rw)
+    fill_halo_regions!(Ru, arch, nothing, nothing)
+    fill_halo_regions!(Rv, arch, nothing, nothing)
+    fill_halo_regions!(Rw, arch, nothing, nothing)
 
     # Compute the right hand side R = ∇⋅U
     ArrayType = array_type(arch)
@@ -46,15 +46,15 @@ function random_divergence_free_source_term(FT, arch, grid)
     set!(Rv, rand(Nx, Ny, Nz))
     set!(Rw, zeros(Nx, Ny, Nz))
 
-    fill_halo_regions!(Ru)
-    fill_halo_regions!(Rv)
-    fill_halo_regions!(Rw)
+    fill_halo_regions!(Ru, arch, nothing, nothing)
+    fill_halo_regions!(Rv, arch, nothing, nothing)
+    fill_halo_regions!(Rw, arch, nothing, nothing)
 
     event = launch!(arch, grid, :xy, _compute_w_from_continuity!, U, grid,
                     dependencies=Event(device(arch)))
     wait(device(arch), event)
 
-    fill_halo_regions!(Rw)
+    fill_halo_regions!(Rw, arch, nothing, nothing)
 
     # Compute the right hand side R = ∇⋅U
     ArrayType = array_type(arch)
@@ -67,10 +67,10 @@ function random_divergence_free_source_term(FT, arch, grid)
 end
 
 function compute_∇²!(∇²ϕ, ϕ, arch, grid)
-    fill_halo_regions!(ϕ)
+    fill_halo_regions!(ϕ, arch)
     event = launch!(arch, grid, :xyz, ∇²!, grid, ϕ.data, ∇²ϕ.data, dependencies=Event(device(arch)))
     wait(device(arch), event)
-    fill_halo_regions!(∇²ϕ)
+    fill_halo_regions!(∇²ϕ, arch)
     return nothing
 end
 

test/test_preconditioned_conjugate_gradient_solver.jl

@@ -39,7 +39,7 @@ function run_pcg_solver_tests(arch)
     function Amatrix_function!(result, x, arch, grid, bcs; args...)
         event = launch!(arch, grid, :xyz, ∇²!, grid, x, result, dependencies=Event(device(arch)))
         wait(device(arch), event)
-        fill_halo_regions!(result, bcs, grid)
+        fill_halo_regions!(result, bcs, arch, grid)
         return nothing
     end
 
@@ -55,13 +55,13 @@ function run_pcg_solver_tests(arch)
     jmid = Int(floor(grid.Ny / 2)) + 1
     CUDA.@allowscalar u.data[imid, jmid, 1] = 1
 
-    fill_halo_regions!(u.data, u.boundary_conditions, grid)
+    fill_halo_regions!(u.data, u.boundary_conditions, arch, grid)
 
     event = launch!(arch, grid, :xyz, divergence!, grid, u.data, v.data, w.data, RHS.data,
                     dependencies=Event(device(arch)))
     wait(device(arch), event)
 
-    fill_halo_regions!(RHS.data, RHS.boundary_conditions, grid)
+    fill_halo_regions!(RHS.data, RHS.boundary_conditions, arch, grid)
 
     pcg_params = (
         PCmatrix_function = nothing,
@@ -83,7 +83,7 @@ function run_pcg_solver_tests(arch)
     event = launch!(arch, grid, :xyz, ∇²!, grid, ϕ.data, result, dependencies=Event(device(arch)))
     wait(device(arch), event)
 
-    fill_halo_regions!(result, ϕ.boundary_conditions, grid)
+    fill_halo_regions!(result, ϕ.boundary_conditions, arch, grid)
 
     CUDA.@allowscalar begin
         @test abs(minimum(result[1:Nx, 1:Ny, 1] .- RHS.data[1:Nx, 1:Ny, 1])) < 1e-12

test/test_reduced_field.jl

@@ -65,7 +65,7 @@ end
                 set!(reduced_field, A)
                 @test reduced_field[1, 1, 1] == A[1, 1, 1]
 
-                fill_halo_regions!(reduced_field)
+                fill_halo_regions!(reduced_field, arch)
 
                 # No-flux boundary conditions at top and bottom
                 @test reduced_field[1, 1, 0] == A[1, 1, 1]

test/test_turbulence_closures.jl

@@ -41,7 +41,7 @@ function constant_isotropic_diffusivity_fluxdiv(FT=Float64; ν=FT(0.3), κ=FT(0.
     end
 
     model_fields = merge(datatuple(velocities), datatuple(tracers))
-    fill_halo_regions!(merge(velocities, tracers), nothing, model_fields)
+    fill_halo_regions!(merge(velocities, tracers), arch, nothing, model_fields)
 
     U, C = datatuples(velocities, tracers)
 
@@ -80,7 +80,7 @@ function anisotropic_diffusivity_fluxdiv(FT=Float64; νh=FT(0.3), κh=FT(0.7), 
     interior(T)[:, 1, 4] .= [0,  1, 0]
 
     model_fields = merge(datatuple(velocities), datatuple(tracers))
-    fill_halo_regions!(merge(velocities, tracers), nothing, model_fields)
+    fill_halo_regions!(merge(velocities, tracers), arch, nothing, model_fields)
 
     U, C = datatuples(velocities, tracers)