Remove global allowscalar from testing suite attempt 2

src/OutputWriters/fetch_output.jl

@@ -1,6 +1,6 @@
 using CUDA
 
-using Oceananigans.Fields: AbstractField, compute_at!, ZeroField
+using Oceananigans.Fields: AbstractField, compute_at!, ZeroField, FunctionField
 using Oceananigans.Models.LagrangianParticleTracking: LagrangianParticles
 
 # TODO: figure out how to support this
@@ -19,6 +19,9 @@ function fetch_output(field::AbstractField, model)
     return parent(field)
 end
 
+# Make sure a function field is evaluated on the CPU
+fetch_output(field::FunctionField, model) = on_architecture(CPU(), field)
+
 function fetch_output(lagrangian_particles::LagrangianParticles, model)
     particle_properties = lagrangian_particles.properties
     names = propertynames(particle_properties)

test/runtests.jl

@@ -15,8 +15,6 @@ end
 ##### Run tests
 #####
 
-CUDA.allowscalar() do
-
 @testset "Oceananigans" begin
 
     if test_file != :none
@@ -246,6 +244,3 @@ CUDA.allowscalar() do
         include("test_convergence.jl")
     end
 end
-
-end #CUDA.allowscalar()
-

test/test_abstract_operations.jl

@@ -366,4 +366,3 @@ for arch in archs
         end
     end
 end
-

test/test_distributed_models.jl

@@ -340,14 +340,16 @@ function test_triply_periodic_halo_communication_with_411_ranks(halo, child_arch
         fill!(field, arch.local_rank)
         fill_halo_regions!(field)
 
-        @test all(east_halo(field, include_corners=false) .== arch.connectivity.east)
-        @test all(west_halo(field, include_corners=false) .== arch.connectivity.west)
-
-        @test all(interior(field) .== arch.local_rank)
-        @test all(north_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(south_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(top_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(bottom_halo(field, include_corners=false) .== arch.local_rank)
+        CUDA.@allowscalar begin
+            @test all(east_halo(field, include_corners=false) .== arch.connectivity.east)
+            @test all(west_halo(field, include_corners=false) .== arch.connectivity.west)
+
+            @test all(interior(field) .== arch.local_rank)
+            @test all(north_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(south_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(top_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(bottom_halo(field, include_corners=false) .== arch.local_rank)
+        end # CUDA.@allowscalar
     end
 
     return nothing
@@ -362,14 +364,16 @@ function test_triply_periodic_halo_communication_with_141_ranks(halo, child_arch
         fill!(field, arch.local_rank)
         fill_halo_regions!(field)
 
-        @test all(north_halo(field, include_corners=false) .== arch.connectivity.north)
-        @test all(south_halo(field, include_corners=false) .== arch.connectivity.south)
+        CUDA.@allowscalar begin
+            @test all(north_halo(field, include_corners=false) .== arch.connectivity.north)
+            @test all(south_halo(field, include_corners=false) .== arch.connectivity.south)
 
-        @test all(interior(field) .== arch.local_rank)
-        @test all(east_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(west_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(top_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(bottom_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(interior(field) .== arch.local_rank)
+            @test all(east_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(west_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(top_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(bottom_halo(field, include_corners=false) .== arch.local_rank)
+        end # CUDA.@allowscalar
     end
 
     return nothing
@@ -384,19 +388,21 @@ function test_triply_periodic_halo_communication_with_221_ranks(halo, child_arch
         fill!(field, arch.local_rank)
         fill_halo_regions!(field)
 
-        @test all(interior(field) .== arch.local_rank)
-
-        @test all(east_halo(field, include_corners=false)  .== arch.connectivity.east)
-        @test all(west_halo(field, include_corners=false)  .== arch.connectivity.west)
-        @test all(north_halo(field, include_corners=false) .== arch.connectivity.north)
-        @test all(south_halo(field, include_corners=false) .== arch.connectivity.south)
-
-        @test all(top_halo(field, include_corners=false)    .== arch.local_rank)
-        @test all(bottom_halo(field, include_corners=false) .== arch.local_rank)
-        @test all(southwest_halo(field) .== arch.connectivity.southwest) 
-        @test all(southeast_halo(field) .== arch.connectivity.southeast) 
-        @test all(northwest_halo(field) .== arch.connectivity.northwest) 
-        @test all(northeast_halo(field) .== arch.connectivity.northeast) 
+        CUDA.@allowscalar begin
+            @test all(interior(field) .== arch.local_rank)
+
+            @test all(east_halo(field, include_corners=false)  .== arch.connectivity.east)
+            @test all(west_halo(field, include_corners=false)  .== arch.connectivity.west)
+            @test all(north_halo(field, include_corners=false) .== arch.connectivity.north)
+            @test all(south_halo(field, include_corners=false) .== arch.connectivity.south)
+
+            @test all(top_halo(field, include_corners=false)    .== arch.local_rank)
+            @test all(bottom_halo(field, include_corners=false) .== arch.local_rank)
+            @test all(southwest_halo(field) .== arch.connectivity.southwest)
+            @test all(southeast_halo(field) .== arch.connectivity.southeast)
+            @test all(northwest_halo(field) .== arch.connectivity.northwest)
+            @test all(northeast_halo(field) .== arch.connectivity.northeast)
+        end # CUDA.@allowscalar
     end
 
     return nothing

test/test_grids.jl

@@ -208,6 +208,8 @@ function test_regular_rectilinear_xnode_ynode_znode_and_spacings(arch, FT)
     for (grid_type, grid) in zip(grids_types, grids)
         @info "        Testing grid utils on $grid_type grid...."
 
+        CUDA.@allowscalar begin
+
         @test xnode(2, grid, Center()) ≈ FT(π/2)
         @test ynode(2, grid, Center()) ≈ FT(π/2)
         @test znode(2, grid, Center()) ≈ FT(π/2)
@@ -235,6 +237,8 @@ function test_regular_rectilinear_xnode_ynode_znode_and_spacings(arch, FT)
         @test Δx(1, 1, 1, grid, Face(), Center(), Center()) ≈ FT(π/N)
         @test Δy(1, 1, 1, grid, Center(), Face(), Center()) ≈ FT(π/N)
         @test Δz(1, 1, 1, grid, Center(), Center(), Face()) ≈ FT(π/N)
+
+        end # CUDA.@allowscalar
     end
 
     return nothing
@@ -349,7 +353,7 @@ function test_grid_equality(arch)
     grid2 = RectilinearGrid(arch, topology=topo, size=(Nx, Ny, Nz), x=(0, 1), y=(-1, 1), z=0:Nz)
     grid3 = RectilinearGrid(arch, topology=topo, size=(Nx, Ny, Nz), x=(0, 1), y=(-1, 1), z=0:Nz)
 
-    return grid1==grid1 && grid2 == grid3 && grid1 !== grid3
+    return grid1 == grid1 && grid2 == grid3 && grid1 !== grid3
 end
 
 function test_grid_equality_over_architectures()
@@ -708,7 +712,7 @@ function test_lat_lon_precomputed_metrics(FT, arch)
     longitude = (lonreg, lonstr, lonregB, lonstrB)
     zcoord    = (zreg,   zstr)
 
-    CUDA.allowscalar() do
+    CUDA.@allowscalar begin
 
     # grid with pre computed metrics vs metrics computed on the fly
     for lat in latitude
@@ -982,6 +986,9 @@ end
 
         for arch in archs
             for FT in float_types
+
+                CUDA.@allowscalar begin
+
                 ccc = (Center(), Center(), Center())
                 grid = RectilinearGrid(arch, FT, size=4, z=(-1, 0), topology=(Flat, Flat, Bounded))
                 x = xnodes(grid, ccc...)
@@ -1030,6 +1037,8 @@ end
                     gpu_grid = on_architecture(GPU(), grid)
                     @test gpu_grid == grid
                 end
+
+                end # CUDA.@allowscalar
             end
         end
     end

test/test_hydrostatic_free_surface_immersed_boundaries.jl

@@ -37,16 +37,16 @@ using Oceananigans.TurbulenceClosures
                 set!(model, u = 1, b = (x, y, z) -> 4z)
 
                 # Inside the bump
-                @test b[4, 4, 2] == 0 
-                @test u[4, 4, 2] == 0
+                CUDA.@allowscalar @test b[4, 4, 2] == 0 
+                CUDA.@allowscalar @test u[4, 4, 2] == 0
 
                 simulation = Simulation(model, Δt = 1e-3, stop_iteration=2)
 
                 run!(simulation)
 
                 # Inside the bump
-                @test b[4, 4, 2] == 0
-                @test u[4, 4, 2] == 0
+                CUDA.@allowscalar @test b[4, 4, 2] == 0
+                CUDA.@allowscalar @test u[4, 4, 2] == 0
             end
         end
 

test/test_implicit_free_surface_solver.jl

@@ -28,7 +28,7 @@ function set_simple_divergent_velocity!(model)
 
     # pick a surface cell at the middle of the domain
     i, j, k = Int(floor(grid.Nx / 2)) + 1, Int(floor(grid.Ny / 2)) + 1, grid.Nz
-    inactive_cell(i, j, k, grid) && error("The nudged cell at ($i, $j, $k) is inactive.")
+    CUDA.@allowscalar inactive_cell(i, j, k, grid) && error("The nudged cell at ($i, $j, $k) is inactive.")
 
     Δy = CUDA.@allowscalar Δyᶜᶠᶜ(i, j, k, grid)
     Δz = CUDA.@allowscalar Δzᶜᶠᶜ(i, j, k, grid)

test/test_jld2_output_writer.jl

@@ -382,8 +382,6 @@ for arch in archs
                                                                           filename = "sliced_func_fields_jld2_test",
                                                                           overwrite_existing = true)
 
-
-
         u₀ = CUDA.@allowscalar model.velocities.u[3, 3, 3]
         v₀ = CUDA.@allowscalar model.velocities.v[3, 3, 3]
         w₀ = CUDA.@allowscalar model.velocities.w[3, 3, 3]

test/test_shallow_water_models.jl

@@ -78,7 +78,7 @@ function test_shallow_water_diffusion_cosine(grid, formulation, fieldname, ξ)
     analytical_solution = Field(location(field), grid)
     analytical_solution .= diffusing_cosine.(ξ, model.clock.time, ν, m)
 
-    return isapprox(field, analytical_solution, atol=1e-6, rtol=1e-6)
+    return CUDA.@allowscalar isapprox(field, analytical_solution, atol=1e-6, rtol=1e-6)
 end
 
 @testset "Shallow Water Models" begin

test/test_simulations.jl

@@ -21,7 +21,7 @@ function wall_time_step_wizard_tests(arch)
     CFL = 0.45
     u₀ = 7
     Δt = 2.5
-    model.velocities.u[1, 1, 1] = u₀
+    CUDA.@allowscalar model.velocities.u[1, 1, 1] = u₀
 
     wizard = TimeStepWizard(cfl=CFL, max_change=Inf, min_change=0)
     Δt = new_time_step(Δt, wizard, model)
@@ -35,7 +35,7 @@ function wall_time_step_wizard_tests(arch)
     Δt = new_time_step(Δt, wizard, model)
     @test Δt ≈ 1.99
 
-    model.velocities.u[1, 1, 1] = u₀/100
+    CUDA.@allowscalar model.velocities.u[1, 1, 1] = u₀/100
 
     wizard = TimeStepWizard(cfl=CFL, max_change=1.1, min_change=0)
     Δt = new_time_step(1.0, wizard, model)
@@ -223,7 +223,7 @@ function run_nan_checker_test(arch; erroring)
     grid = RectilinearGrid(arch, size=(4, 2, 1), extent=(1, 1, 1))
     model = NonhydrostaticModel(grid=grid)
     simulation = Simulation(model, Δt=1, stop_iteration=2)
-    model.velocities.u[1, 1, 1] = NaN
+    CUDA.@allowscalar model.velocities.u[1, 1, 1] = NaN
     erroring && erroring_NaNChecker!(simulation)
 
     if erroring

test/test_turbulence_closures.jl

@@ -185,7 +185,7 @@ function run_catke_tke_substepping_tests(arch, closure)
     time_step!(model, 1)
 
     # Check that eⁿ⁺¹ == Δt * Gⁿ.e with Δt = 1 (euler step)
-    @test model.tracers.e ≈ model.timestepper.G⁻.e
+    CUDA.@allowscalar @test model.tracers.e ≈ model.timestepper.G⁻.e
 
     eⁿ  = deepcopy(model.tracers.e)
     G⁻⁻ = deepcopy(model.timestepper.G⁻.e)
@@ -199,8 +199,8 @@ function run_catke_tke_substepping_tests(arch, closure)
 
     eⁿ⁺¹ = compute!(Field(eⁿ + C₁ * G⁻ - C₂ * G⁻⁻))
 
-    # Check that eⁿ⁺¹ == eⁿ + Δt * (C₁ Gⁿ.e - C₂ G⁻.e)
-    @test model.tracers.e ≈ eⁿ⁺¹
+    # Check that eⁿ⁺¹ == eⁿ + Δt * (C₁ Gⁿ.e - C₂ G⁻.e) 
+    CUDA.@allowscalar @test model.tracers.e ≈ eⁿ⁺¹
 
     return model
 end
@@ -320,8 +320,8 @@ end
                     ν = viscosity(model.closure, model.diffusivity_fields)
                     @test viscosity(model) == viscosity(model.closure, model.diffusivity_fields)
                     ν_dx_u = ν * ∂x(u)
-                    @test ν_dx_u[1, 1, 1] == 0
-                    @test κ_dx_c[1, 1, 1] == 0
+                    CUDA.@allowscalar @test ν_dx_u[1, 1, 1] == 0
+                    CUDA.@allowscalar @test κ_dx_c[1, 1, 1] == 0
                 end
             end
         end

test/test_vertical_vorticity_field.jl

@@ -20,7 +20,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField,
 
             compute!(ζ)
 
-            @test all(isfinite.(ζ.data))
+            CUDA.@allowscalar @test all(isfinite.(ζ.data))
         end
     end
 end