Merge pull request #1403 from CliMA/ali/flesh-out-vs-grid

.buildkite/pipeline.yml

@@ -78,6 +78,33 @@ steps:
       architecture: CPU
     depends_on: "init_cpu"
 
+#####
+##### Solver tests
+#####
+
+  - label: "🦅 gpu solver tests"
+    env:
+      JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "solvers"
+    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 solver tests"
+    env:
+      JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "solvers"
+      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"
+
 #####
 ##### IncompressibleModel and time stepping (part 1)
 #####

Project.toml

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

benchmark/Manifest.toml

@@ -2,9 +2,9 @@
 
 [[AbstractFFTs]]
 deps = ["LinearAlgebra"]
-git-tree-sha1 = "051c95d6836228d120f5f4b984dd5aba1624f716"
+git-tree-sha1 = "8ed9de2f1b1a9b1dee48582ad477c6e67b83eb2c"
 uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
-version = "0.5.0"
+version = "1.0.0"
 
 [[AbstractTrees]]
 deps = ["Markdown"]
@@ -31,9 +31,9 @@ uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 version = "0.1.0"
 
 [[BSON]]
-git-tree-sha1 = "dd36d7cf3d185eeaaf64db902c15174b22f5dafb"
+git-tree-sha1 = "2878972c4bc17d9c8d26d48d9ef00fcfe1899e7a"
 uuid = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
-version = "0.2.6"
+version = "0.3.0"
 
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
@@ -57,9 +57,9 @@ version = "2.4.1"
 
 [[CategoricalArrays]]
 deps = ["DataAPI", "Future", "JSON", "Missings", "Printf", "Statistics", "StructTypes", "Unicode"]
-git-tree-sha1 = "99809999c8ee01fa89498480b147f7394ea5450f"
+git-tree-sha1 = "dbfddfafb75fae5356e00529ce67454125935945"
 uuid = "324d7699-5711-5eae-9e2f-1d82baa6b597"
-version = "0.9.2"
+version = "0.9.3"
 
 [[ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
@@ -125,9 +125,9 @@ version = "0.1.3"
 
 [[FFTW]]
 deps = ["AbstractFFTs", "FFTW_jll", "IntelOpenMP_jll", "Libdl", "LinearAlgebra", "MKL_jll", "Reexport"]
-git-tree-sha1 = "8fda0934cb99db617171f7296dc361f4d6fa5424"
+git-tree-sha1 = "1b48dbde42f307e48685fa9213d8b9f8c0d87594"
 uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-version = "1.3.0"
+version = "1.3.2"
 
 [[FFTW_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
@@ -195,6 +195,12 @@ git-tree-sha1 = "b616937c31337576360cb9fb872ec7633af7b194"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
 version = "3.6.0"
 
+[[LazyArtifacts]]
+deps = ["Pkg"]
+git-tree-sha1 = "4bb5499a1fc437342ea9ab7e319ede5a457c0968"
+uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+version = "1.3.0"
+
 [[LeftChildRightSiblingTrees]]
 deps = ["AbstractTrees"]
 git-tree-sha1 = "71be1eb5ad19cb4f61fa8c73395c0338fd092ae0"
@@ -216,10 +222,10 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
 
 [[MKL_jll]]
-deps = ["IntelOpenMP_jll", "Libdl", "Pkg"]
-git-tree-sha1 = "eb540ede3aabb8284cb482aa41d00d6ca850b1f8"
+deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"]
+git-tree-sha1 = "c253236b0ed414624b083e6b72bfe891fbd2c7af"
 uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
-version = "2020.2.254+0"
+version = "2021.1.1+1"
 
 [[MacroTools]]
 deps = ["Markdown", "Random"]
@@ -247,9 +253,9 @@ uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 version = "0.7.14"
 
 [[OrderedCollections]]
-git-tree-sha1 = "d45739abcfc03b51f6a42712894a593f74c80a23"
+git-tree-sha1 = "4fa2ba51070ec13fcc7517db714445b4ab986bdf"
 uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-version = "1.3.3"
+version = "1.4.0"
 
 [[Parsers]]
 deps = ["Dates"]
@@ -349,9 +355,9 @@ uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [[StructTypes]]
 deps = ["Dates", "UUIDs"]
-git-tree-sha1 = "65a43f5218197bc7091b76bc273a5e323a1d7b0d"
+git-tree-sha1 = "d7f4287dbc1e590265f50ceda1b40ed2bb31bbbb"
 uuid = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
-version = "1.2.3"
+version = "1.4.0"
 
 [[TableTraits]]
 deps = ["IteratorInterfaceExtensions"]

benchmark/benchmark_fourier_tridiagonal_poisson_solver.jl

@@ -0,0 +1,49 @@
+using BenchmarkTools
+using CUDA
+using Oceananigans
+using Benchmarks
+
+using Oceananigans.Solvers
+
+# Benchmark function
+
+function benchmark_fourier_tridiagonal_poisson_solver(Arch, N, topo)
+    grid = VerticallyStretchedRectilinearGrid(architecture=Arch(), topology=topo, size=(N, N, N), x=(0, 1), y=(0, 1), zF=collect(0:N))
+    solver = FourierTridiagonalPoissonSolver(Arch(), grid)
+
+    solve_poisson_equation!(solver) # warmup
+
+    trial = @benchmark begin
+        @sync_gpu solve_poisson_equation!($solver)
+    end samples=10
+
+    return trial
+end
+
+# Benchmark parameters
+
+Architectures = has_cuda() ? [CPU, GPU] : [CPU]
+Ns = [256]
+PB = (Periodic, Bounded)
+Topologies = collect(Iterators.product(PB, PB, (Bounded,)))[:]
+
+# Run and summarize benchmarks
+
+suite = run_benchmarks(benchmark_fourier_tridiagonal_poisson_solver; Architectures, Ns, Topologies)
+
+df = benchmarks_dataframe(suite)
+sort!(df, [:Architectures, :Topologies, :Ns], by=(string, string, identity))
+benchmarks_pretty_table(df, title="Fourier-tridiagonal Poisson solver benchmarks")
+
+if GPU in Architectures
+    df = gpu_speedups_suite(suite) |> speedups_dataframe
+    sort!(df, [:Topologies, :Ns], by=(string, identity))
+    benchmarks_pretty_table(df, title="Fourier-tridiagonal Poisson solver CPU -> GPU speedup")
+end
+
+for Arch in Architectures
+    suite_arch = speedups_suite(suite[@tagged Arch], base_case=(Arch, Ns[1], (Periodic, Periodic, Bounded)))
+    df_arch = speedups_dataframe(suite_arch, slowdown=true)
+    sort!(df_arch, [:Topologies, :Ns], by=string)
+    benchmarks_pretty_table(df_arch, title="Fourier-tridiagonal Poisson solver relative performance ($Arch)")
+end

benchmark/benchmark_vertically_stretched_incompressible_model.jl

@@ -0,0 +1,40 @@
+using BenchmarkTools
+using CUDA
+using Oceananigans
+using Benchmarks
+
+# Benchmark function
+
+function benchmark_vertically_stretched_incompressible_model(Arch, FT, N)
+    grid = VerticallyStretchedRectilinearGrid(architecture=Arch(), size=(N, N, N), x=(0, 1), y=(0, 1), zF=collect(0:N))
+    model = IncompressibleModel(architecture=Arch(), float_type=FT, grid=grid)
+
+    time_step!(model, 1) # warmup
+
+    trial = @benchmark begin
+        @sync_gpu time_step!($model, 1)
+    end samples=10
+
+    return trial
+end
+
+# Benchmark parameters
+
+Architectures = has_cuda() ? [CPU, GPU] : [CPU]
+Float_types = [Float32, Float64]
+Ns = [32, 64, 128, 256]
+
+# Run and summarize benchmarks
+
+print_system_info()
+suite = run_benchmarks(benchmark_vertically_stretched_incompressible_model; Architectures, Float_types, Ns)
+
+df = benchmarks_dataframe(suite)
+sort!(df, [:Architectures, :Float_types, :Ns], by=(string, string, identity))
+benchmarks_pretty_table(df, title="Vertically-stretched incompressible model benchmarks")
+
+if GPU in Architectures
+    df_Δ = gpu_speedups_suite(suite) |> speedups_dataframe
+    sort!(df_Δ, [:Float_types, :Ns], by=(string, identity))
+    benchmarks_pretty_table(df_Δ, title="Vertically-stretched incompressible model CPU -> GPU speedup")
+end

src/Buoyancy/nonlinear_equation_of_state.jl

@@ -1,14 +1,16 @@
 using Oceananigans.Fields: AbstractField
+using Oceananigans.Grids: zF, zC
+using Oceananigans.Operators: Δzᵃᵃᶠ, Δzᵃᵃᶜ
 
 """ Return the geopotential height at `i, j, k` at cell centers. """
-@inline function Zᵃᵃᶜ(i, j, k, grid::AbstractGrid{FT}) where FT 
+@inline function Zᵃᵃᶜ(i, j, k, grid::AbstractGrid{FT}) where FT
     @inbounds begin
         if k < 1
-            return grid.zC[1] + (1 - k) * grid.Δz
+            return zC(1, grid) + (1 - k) * Δzᵃᵃᶠ(i, j, 1, grid)
         elseif k > grid.Nz
-            return grid.zC[grid.Nz] - (k - grid.Nz) * grid.Δz
+            return zC(grid.Nz, grid) - (k - grid.Nz) * Δzᵃᵃᶠ(i, j, grid.Nz, grid)
         else
-            return grid.zC[k]
+            return zC(k, grid)
         end
     end
 end
@@ -17,11 +19,11 @@ end
 @inline function Zᵃᵃᶠ(i, j, k, grid::AbstractGrid{FT}) where FT
     @inbounds begin
         if k < 1
-            return grid.zF[1] + (1 - k) * grid.Δz
+            return zF(1, grid) + (1 - k) * Δzᵃᵃᶜ(i, j, 1, grid)
         elseif k > grid.Nz + 1
-            return grid.zF[grid.Nz+1] - (k - grid.Nz + 1) * grid.Δz
+            return zF(grid.Nz + 1, grid) - (k - grid.Nz + 1) * Δzᵃᵃᶜ(i, j, k, grid)
         else
-            return grid.zF[k]
+            return zF(k, grid)
         end
     end
 end

src/Grids/Grids.jl

@@ -8,14 +8,14 @@ export
     xnode, ynode, znode, xnodes, ynodes, znodes, nodes,
     xC, xF, yC, yF, zC, zF
 
-import Base: size, length, eltype, show
-
-import Oceananigans: short_show
+using Adapt
+using OffsetArrays
 
 using Oceananigans
 using Oceananigans.Architectures
 
-using OffsetArrays
+import Base: size, length, eltype, show
+import Oceananigans: short_show
 
 #####
 ##### Abstract types