Dynamically disable CuArray scalar operations

src/Fields/field.jl

@@ -52,7 +52,7 @@ end
                                  data = zeros(arch, grid, (X, Y, Z)) ] )
 
 Construct a `Field` on `grid` with `data` on architecture `arch` with
-boundary conditions `bcs`. Each of `(X, Y, Z)` is either `Cell` or `Face` and determines 
+boundary conditions `bcs`. Each of `(X, Y, Z)` is either `Cell` or `Face` and determines
 the field's location in `(x, y, z)`.
 
 Example
@@ -61,7 +61,7 @@ Example
 julia> ω = Field(Face, Face, Cell, CPU(), RegularCartesianmodel.grid)
 
 """
-function Field(X, Y, Z, arch, grid, 
+function Field(X, Y, Z, arch, grid,
                 bcs = FieldBoundaryConditions(grid, (X, Y, Z)),
                data = zeros(eltype(grid), arch, grid, (X, Y, Z)))
 
@@ -89,7 +89,7 @@ Field(L::Tuple, args...) = Field(destantiate.(L)..., args...)
 #####
 
 """
-    CellField([ FT=eltype(grid) ], arch::AbstractArchitecture, grid, 
+    CellField([ FT=eltype(grid) ], arch::AbstractArchitecture, grid,
               [  bcs = TracerBoundaryConditions(grid),
                 data = zeros(FT, arch, grid, (Cell, Cell, Cell) ] )
 
@@ -111,13 +111,13 @@ end
 Return a `Field{Face, Cell, Cell}` on architecture `arch` and `grid` containing `data`
 with field boundary conditions `bcs`.
 """
-function XFaceField(FT::DataType, arch, grid, 
+function XFaceField(FT::DataType, arch, grid,
                      bcs = UVelocityBoundaryConditions(grid),
                     data = zeros(FT, arch, grid, (Face, Cell, Cell)))
 
     return Field{Face, Cell, Cell}(data, grid, bcs)
 end
-  
+
 """
     YFaceField([ FT=eltype(grid) ], arch::AbstractArchitecture, grid,
                [  bcs = VVelocityBoundaryConditions(grid),
@@ -132,7 +132,7 @@ function YFaceField(FT::DataType, arch, grid,
 
     return Field{Cell, Face, Cell}(data, grid, bcs)
 end
-  
+
 """
     ZFaceField([ FT=eltype(grid) ], arch::AbstractArchitecture, grid,
                [  bcs = WVelocityBoundaryConditions(grid),
@@ -167,7 +167,7 @@ location(f, i) = location(f)[i]
 architecture(f::Field) = architecture(f.data)
 architecture(o::OffsetArray) = architecture(o.parent)
 
-"Returns the length of a field's `data`."    
+"Returns the length of a field's `data`."
 @inline length(f::Field) = length(f.data)
 
 
@@ -219,7 +219,9 @@ contained by `f` along `x, y, z`.
 
 @inline cpudata(a) = data(a)
 
-@hascuda @inline cpudata(f::Field{X, Y, Z, <:CuArray}) where {X, Y, Z} =
+const OffsetCuArray = OffsetArray{T, D, <:CuArray} where {T, D}
+
+@hascuda @inline cpudata(f::Field{X, Y, Z, <:OffsetCuArray}) where {X, Y, Z} =
     OffsetArray(Array(parent(f)), f.grid, location(f))
 
 # Endpoint for recursive `datatuple` function:
@@ -229,12 +231,12 @@ contained by `f` along `x, y, z`.
 @inline Base.parent(f::Field) = f.data.parent
 
 "Returns a view of `f` that excludes halo points."
-@inline interior(f::Field{X, Y, Z}) where {X, Y, Z} = view(f.data, interior_indices(X, topology(f, 1), f.grid.Nx), 
-                                                                   interior_indices(Y, topology(f, 2), f.grid.Ny), 
+@inline interior(f::Field{X, Y, Z}) where {X, Y, Z} = view(f.data, interior_indices(X, topology(f, 1), f.grid.Nx),
+                                                                   interior_indices(Y, topology(f, 2), f.grid.Ny),
                                                                    interior_indices(Z, topology(f, 3), f.grid.Nz))
 
 "Returns a reference (not a view) to the interior points of `field.data.parent.`"
-@inline interiorparent(f::Field{X, Y, Z}) where {X, Y, Z} = 
+@inline interiorparent(f::Field{X, Y, Z}) where {X, Y, Z} =
     @inbounds f.data.parent[interior_parent_indices(X, topology(f, 1), f.grid.Nx, f.grid.Hx),
                             interior_parent_indices(Y, topology(f, 2), f.grid.Ny, f.grid.Hy),
                             interior_parent_indices(Z, topology(f, 3), f.grid.Nz, f.grid.Hz)]
@@ -271,7 +273,7 @@ offset_indices(::Type{Face}, ::Type{Bounded}, N, H=0) = 1 - H : N + H + 1
     OffsetArray(underlying_data, grid, loc)
 
 Returns an `OffsetArray` that maps to `underlying_data` in memory,
-with offset indices appropriate for the `data` of a field on 
+with offset indices appropriate for the `data` of a field on
 a `grid` of `size(grid)` and located at `loc`.
 """
 function OffsetArray(underlying_data, grid, loc)
@@ -286,7 +288,7 @@ end
     zeros([FT=Float64], ::CPU, grid, loc)
 
 Returns an `OffsetArray` of zeros of float type `FT`, with
-parent data in CPU memory and indices corresponding to a field on a 
+parent data in CPU memory and indices corresponding to a field on a
 `grid` of `size(grid)` and located at `loc`.
 """
 function Base.zeros(FT, ::CPU, grid, loc=(Cell, Cell, Cell))

src/Fields/set!.jl

@@ -52,7 +52,7 @@ function set!(Φ::NamedTuple; kwargs...)
     return nothing
 end
 
-set!(u::Field, v::Number) = @. u.data = v
+set!(u::Field, v::Number) = @. u.data.parent = v
 
 set!(u::Field{X, Y, Z, A}, v::Field{X, Y, Z, A}) where {X, Y, Z, A} =
     @. u.data.parent = v.data.parent

src/Oceananigans.jl

@@ -116,15 +116,6 @@ function write_output end
 #####
 
 include("Architectures.jl")
-
-using Oceananigans.Architectures: @hascuda
-@hascuda begin
-    @debug "CUDA-enabled GPU(s) detected:"
-    for (gpu, dev) in enumerate(CUDA.devices())
-        @debug "$dev: $(CUDA.name(dev))"
-    end
-end
-
 include("Utils/Utils.jl")
 include("Logger.jl")
 include("Grids/Grids.jl")
@@ -149,6 +140,7 @@ include("AbstractOperations/AbstractOperations.jl")
 ##### Re-export stuff from submodules
 #####
 
+using .Logger
 using .Architectures
 using .Utils
 using .Grids
@@ -164,4 +156,15 @@ using .Models
 using .TimeSteppers
 using .Simulations
 
+function __init__()
+    Logging.global_logger(ModelLogger())
+    @hascuda begin
+        @debug "CUDA-enabled GPU(s) detected:"
+        for (gpu, dev) in enumerate(CUDA.devices())
+            @debug "$dev: $(CUDA.name(dev))"
+        end
+	CUDA.allowscalar(false)
+    end
+end
+
 end # module

src/OutputWriters/OutputWriters.jl

@@ -8,12 +8,12 @@ export
 
 using CUDA
 
-using Oceananigans
+using Oceananigans.Architectures
 using Oceananigans.Grids
 using Oceananigans.Fields
-using Oceananigans.Architectures
+using Oceananigans.Models
 
-using Oceananigans: AbstractOutputWriter, @hascuda
+using Oceananigans: AbstractOutputWriter
 using Oceananigans.Fields: OffsetArray
 
 Base.open(ow::AbstractOutputWriter) = nothing

src/Solvers/channel_pressure_solver.jl

@@ -85,8 +85,8 @@ function ChannelPressureSolver(::GPU, grid, pressure_bcs, no_args...)
     M_ky = ones(1, Ny, 1) |> CuArray
     M_kz = ones(1, 1, Nz) |> CuArray
 
-    M_ky[1] = 0
-    M_kz[1] = 0
+    CUDA.@allowscalar M_ky[1] = 0
+    CUDA.@allowscalar M_kz[1] = 0
 
     constants = (ω_4Ny⁺ = ω_4Ny⁺, ω_4Nz⁺ = ω_4Nz⁺, ω_4Ny⁻ = ω_4Ny⁻, ω_4Nz⁻ = ω_4Nz⁻,
                  r_y_inds = r_y_inds, r_z_inds = r_z_inds,
@@ -142,7 +142,7 @@ function solve_poisson_equation!(solver::PressureSolver{Channel, GPU}, grid)
 
     @. B = -B / (kx² + ky² + kz²)
 
-    B[1, 1, 1] = 0  # Setting DC component of the solution (the mean) to be zero.
+    CUDA.@allowscalar B[1, 1, 1] = 0  # Setting DC component of the solution (the mean) to be zero.
 
     solver.transforms.IFFTx! * B  # Calculate IFFTˣ(ϕ̂) in place.
 

src/Solvers/horizontally_periodic_pressure_solver.jl

@@ -84,7 +84,7 @@ function HorizontallyPeriodicPressureSolver(::GPU, grid, pressure_bcs, no_args..
     # multiplied by 2. For some reason, we only need to account for this when
     # doing a 1D IDCT (for PPN boundary conditions) but not for a 2D IDCT. It's
     # possible that the masks are effectively doing this job.
-    ω_4Nz⁻[1] *= 1/2
+    CUDA.@allowscalar ω_4Nz⁻[1] *= 1/2
 
     constants = (ω_4Nz⁺ = ω_4Nz⁺, ω_4Nz⁻ = ω_4Nz⁻)
 
@@ -145,7 +145,7 @@ function solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, GP
     # Setting DC component of the solution (the mean) to be zero. This is also
     # necessary because the source term to the Poisson equation has zero mean
     # and so the DC component comes out to be ∞.
-    ϕ[1, 1, 1] = 0
+    CUDA.@allowscalar ϕ[1, 1, 1] = 0
 
     solver.transforms.IFFTxy! * ϕ  # Calculate IFFTˣʸ(ϕ̂) in place.
 

test/runtests.jl

@@ -72,31 +72,30 @@ closures = (
 
 include("runtests_utils.jl")
 
-with_logger(ModelLogger()) do
-    @testset "Oceananigans" begin
-        include("test_grids.jl")
-        include("test_operators.jl")
-        include("test_boundary_conditions.jl")
-        include("test_fields.jl")
-        include("test_halo_regions.jl")
-        include("test_solvers.jl")
-        include("test_pressure_solvers.jl")
-        include("test_coriolis.jl")
-        include("test_buoyancy.jl")
-        include("test_surface_waves.jl")
-        include("test_models.jl")
-        include("test_simulations.jl")
-        include("test_time_stepping.jl")
-        include("test_time_stepping_bcs.jl")
-        include("test_forcings.jl")
-        include("test_turbulence_closures.jl")
-        include("test_dynamics.jl")
-        include("test_diagnostics.jl")
-        include("test_output_writers.jl")
-        include("test_abstract_operations.jl")
-        include("test_regression.jl")
-        include("test_examples.jl")
-        include("test_verification.jl")
-        include("test_benchmarks.jl")
-    end
+@testset "Oceananigans" begin
+    include("test_grids.jl")
+    include("test_operators.jl")
+    include("test_boundary_conditions.jl")
+    include("test_fields.jl")
+    include("test_halo_regions.jl")
+    include("test_solvers.jl")
+    include("test_pressure_solvers.jl")
+    include("test_coriolis.jl")
+    include("test_buoyancy.jl")
+    include("test_surface_waves.jl")
+    include("test_models.jl")
+    include("test_simulations.jl")
+    include("test_time_stepping.jl")
+    include("test_time_stepping_bcs.jl")
+    include("test_forcings.jl")
+    include("test_turbulence_closures.jl")
+    include("test_dynamics.jl")
+    include("test_diagnostics.jl")
+    include("test_output_writers.jl")
+    include("test_abstract_operations.jl")
+    include("test_regression.jl")
+    include("test_examples.jl")
+    include("test_verification.jl")
+    include("test_benchmarks.jl")
 end
+

test/test_fields.jl

@@ -1,11 +1,13 @@
+using Oceananigans.Fields: cpudata
+
 """
     correct_field_size(N, L, ftf)
 
 Test that the field initialized by the field type function `ftf` on the grid g
 has the correct size.
 """
 correct_field_size(a, g, fieldtype, Tx, Ty, Tz) = size(parent(fieldtype(a, g)))  == (Tx, Ty, Tz)
-    
+
 """
      correct_field_value_was_set(N, L, ftf, val)
 
@@ -16,7 +18,7 @@ function.
 function correct_field_value_was_set(arch, grid, fieldtype, val::Number)
     f = fieldtype(arch, grid)
     set!(f, val)
-    return interior(f) ≈ val * ones(size(f))
+    CUDA.@allowscalar return interior(f) ≈ val * ones(size(f))
 end
 
 @testset "Fields" begin
@@ -84,11 +86,21 @@ end
         end
     end
 
-    @testset "Miscellaneous field functionality" begin
-        @info "  Testing miscellaneous field functionality..."
+    @testset "Field utils" begin
+        @info "  Testing field utils..."
+
         @test Fields.has_velocities(()) == false
         @test Fields.has_velocities((:u,)) == false
         @test Fields.has_velocities((:u, :v)) == false
         @test Fields.has_velocities((:u, :v, :w)) == true
+
+		grid = RegularCartesianGrid(size=(4, 6, 8), extent=(1, 1, 1))
+		ϕ = CellField(CPU(), grid)
+		@test cpudata(ϕ).parent isa Array
+
+		@hascuda begin
+			ϕ = CellField(GPU(), grid)
+			@test cpudata(ϕ).parent isa Array
+		end
     end
 end

test/test_halo_regions.jl

@@ -26,7 +26,7 @@ function halo_regions_correctly_filled(arch, FT, Nx, Ny, Nz)
     grid = RegularCartesianGrid(FT, size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
     field = CellField(FT, arch, grid)
 
-    interior(field) .= rand(FT, Nx, Ny, Nz)
+    set!(field, rand(FT, Nx, Ny, Nz))
     fill_halo_regions!(field, arch)
 
     Hx, Hy, Hz = grid.Hx, grid.Hy, grid.Hz