Merge branch 'main' into kp/dtopo

.buildkite/pipeline.yml

@@ -512,17 +512,6 @@ steps:
       queue: central
       slurm_ntasks: 1
 
-  - label: ":rocket: Allocations analysis"
-    key: "cpu_allocations"
-    command:
-      - "julia --color=yes --project=perf perf/allocs.jl"
-    artifact_paths:
-      - "perf/allocations_output/*"
-    agents:
-      config: cpu
-      queue: central
-      slurm_ntasks: 1
-
   - label: ":computer: dry held-suarez rho etot implicit"
     key: "cpu_held_suarez_rho_etot_imp"
     command:
@@ -543,4 +532,30 @@ steps:
     agents:
       config: cpu
       queue: central
-      slurm_ntasks: 1
+      slurm_ntasks: 1
+
+  - label: ":rocket::computer: Allocations analysis"
+    key: "cpu_allocations"
+    command:
+      - "julia --color=yes --project=perf perf/allocs.jl"
+    artifact_paths:
+      - "perf/allocations_output/*"
+    agents:
+      config: cpu
+      queue: central
+      slurm_ntasks: 1
+
+  - label: ":rocket::computer: Flamegraph profile"
+    key: "cpu_flamegraph"
+    depends_on: "cpu_allocations"
+    allow_dependency_failure: true
+    env:
+      CI_PERF_CPUPROFILE: "true"
+    command:
+      - "julia --color=yes --project=perf perf/flame.jl"
+    artifact_paths:
+      - "perf/output/*.cpuprofile"
+    agents:
+      config: cpu
+      queue: central
+      slurm_ntasks: 1

.gitignore

@@ -7,6 +7,9 @@ Manifest.toml
 # Files generated by invoking Julia with --track-allocation
 *.jl.mem
 
+# Files generated by cpuprofile
+*.cpuprofile
+
 # System-specific files and directories generated by the BinaryProvider and BinDeps packages
 # They contain absolute paths specific to the host computer, and so should not be committed
 deps/deps.jl
@@ -32,5 +35,6 @@ docs/src/tutorials/
 
 
 *.gif
+*.gz
 *.mp4
 *.png

examples/3dsphere/baroclinic_wave.jl

@@ -316,11 +316,7 @@ end
 dt = 5
 prob = ODEProblem(rhs!, Y, (0.0, time_end))
 
-if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
-    throw(:exit_profile)
-end
-
-sol = @timev solve(
+integrator = OrdinaryDiffEq.init(
     prob,
     SSPRK33(),
     dt = dt,
@@ -330,6 +326,12 @@ sol = @timev solve(
     progress_message = (dt, u, p, t) -> t,
 )
 
+if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
+    throw(:exit_profile)
+end
+
+sol = @timev OrdinaryDiffEq.solve!(integrator)
+
 # visualization artifacts
 if test_name == "baroclinic_wave"
     @info "Solution L₂ norm at time t = 0: ", norm(Y.Yc.ρe)

examples/3dsphere/deformation_flow.jl

@@ -40,14 +40,15 @@ const ϕ_c = 0.0 # initial latitude of tracers
 const z_c = 5.0e3 # initial altitude of tracers
 const R_t = R / 2 # horizontal half-width of tracers
 const Z_t = 1000.0 # vertical half-width of tracers
+const κ₄ = 1.0e16 # hyperviscosity
 
 # set up function space
 function sphere_3D(
     R = 6.37122e6,
     zlim = (0, 12.0e3),
-    helem = 6,
+    helem = 4,
     zelem = 12,
-    npoly = 3,
+    npoly = 4,
 )
     FT = Float64
     vertdomain = Domains.IntervalDomain(
@@ -81,7 +82,7 @@ r1(λ, ϕ) = R * acos(sind(ϕ_c) * sind(ϕ) + cosd(ϕ_c) * cosd(ϕ) * cosd(λ -
 r2(λ, ϕ) = R * acos(sind(ϕ_c) * sind(ϕ) + cosd(ϕ_c) * cosd(ϕ) * cosd(λ - λ_c2))
 
 p(z) = p_0 * exp(-z / H)
-ρ(z) = p(z) / R_d / T_0
+ρ_ref(z) = p(z) / R_d / T_0
 
 y0 = map(coords) do coord
     z = coord.z
@@ -108,12 +109,12 @@ y0 = map(coords) do coord
     end
     q4 = 1 - 3 / 10 * (q1 + q2 + q3)
 
-    ρq1 = ρ(z) * q1
-    ρq2 = ρ(z) * q2
-    ρq3 = ρ(z) * q3
-    ρq4 = ρ(z) * q4
+    ρq1 = ρ_ref(z) * q1
+    ρq2 = ρ_ref(z) * q2
+    ρq3 = ρ_ref(z) * q3
+    ρq4 = ρ_ref(z) * q4
 
-    return (ρq1 = ρq1, ρq2 = ρq2, ρq3 = ρq3, ρq4 = ρq4)
+    return (ρ = ρ_ref(z), ρq1 = ρq1, ρq2 = ρq2, ρq3 = ρq3, ρq4 = ρq4)
 end
 
 function rhs!(dydt, y, (coords, face_coords), t)
@@ -142,21 +143,25 @@ function rhs!(dydt, y, (coords, face_coords), t)
     uu = @. ua + ud
     uv = @. k * sind(2 * λp) * cosd(ϕ) * cos(pi * t / τ)
     ω = @. ω_0 * sind(λpf) * cosd(ϕf) * cos(2 * pi * t / τ) * sp
-    uw = @. -ω / ρ(zf) / grav
+    uw = @. -ω / ρ_ref(zf) / grav
 
     uₕ = Geometry.Covariant12Vector.(Geometry.UVVector.(uu, uv))
     w = Geometry.Covariant3Vector.(Geometry.WVector.(uw))
 
+    ρ = y.ρ
     ρq1 = y.ρq1
     ρq2 = y.ρq2
     ρq3 = y.ρq3
     ρq4 = y.ρq4
 
+    dρ = dydt.ρ
     dρq1 = dydt.ρq1
     dρq2 = dydt.ρq2
     dρq3 = dydt.ρq3
     dρq4 = dydt.ρq4
 
+    dρ .= 0 .* ρ
+
     If2c = Operators.InterpolateF2C()
     Ic2f = Operators.InterpolateC2F(
         bottom = Operators.Extrapolate(),
@@ -167,23 +172,43 @@ function rhs!(dydt, y, (coords, face_coords), t)
         bottom = Operators.SetValue(Geometry.Contravariant3Vector(0.0)),
     )
     hdiv = Operators.Divergence()
+    hwdiv = Operators.WeakDivergence()
+    hgrad = Operators.Gradient()
+
+    ### HYPERVISCOSITY
+
+    χq1 = @. dρq1 = hwdiv(hgrad(ρq1 / ρ))
+    Spaces.weighted_dss!(dρq1)
+    @. dρq1 = -κ₄ * hwdiv(ρ * hgrad(χq1))
+
+    χq2 = @. dρq2 = hwdiv(hgrad(ρq2 / ρ))
+    Spaces.weighted_dss!(dρq2)
+    @. dρq2 = -κ₄ * hwdiv(ρ * hgrad(χq2))
+
+    χq3 = @. dρq3 = hwdiv(hgrad(ρq3 / ρ))
+    Spaces.weighted_dss!(dρq3)
+    @. dρq3 = -κ₄ * hwdiv(ρ * hgrad(χq3))
+
+    χq4 = @. dρq4 = hwdiv(hgrad(ρq4 / ρ))
+    Spaces.weighted_dss!(dρq4)
+    @. dρq4 = -κ₄ * hwdiv(ρ * hgrad(χq4))
 
     cw = If2c.(w)
     cuvw = Geometry.Covariant123Vector.(uₕ) .+ Geometry.Covariant123Vector.(cw)
 
-    @. dρq1 = -hdiv(cuvw * ρq1)
+    @. dρq1 -= hdiv(cuvw * ρq1)
     @. dρq1 -= vdivf2c(w * Ic2f(ρq1))
     @. dρq1 -= vdivf2c(Ic2f(uₕ * ρq1))
 
-    @. dρq2 = -hdiv(cuvw * ρq2)
+    @. dρq2 -= hdiv(cuvw * ρq2)
     @. dρq2 -= vdivf2c(w * Ic2f(ρq2))
     @. dρq2 -= vdivf2c(Ic2f(uₕ * ρq2))
 
-    @. dρq3 = -hdiv(cuvw * ρq3)
+    @. dρq3 -= hdiv(cuvw * ρq3)
     @. dρq3 -= vdivf2c(w * Ic2f(ρq3))
     @. dρq3 -= vdivf2c(Ic2f(uₕ * ρq3))
 
-    @. dρq4 = -hdiv(cuvw * ρq4)
+    @. dρq4 -= hdiv(cuvw * ρq4)
     @. dρq4 -= vdivf2c(w * Ic2f(ρq4))
     @. dρq4 -= vdivf2c(Ic2f(uₕ * ρq4))
 
@@ -210,24 +235,24 @@ sol = solve(
 )
 
 q1_error =
-    norm(sol.u[end].ρq1 ./ ρ.(coords.z) .- y0.ρq1 ./ ρ.(coords.z)) /
-    norm(y0.ρq1 ./ ρ.(coords.z))
-@test q1_error ≈ 0.0 atol = 0.9
+    norm(sol.u[end].ρq1 ./ ρ_ref.(coords.z) .- y0.ρq1 ./ ρ_ref.(coords.z)) /
+    norm(y0.ρq1 ./ ρ_ref.(coords.z))
+@test q1_error ≈ 0.0 atol = 0.7
 
 q2_error =
-    norm(sol.u[end].ρq2 ./ ρ.(coords.z) .- y0.ρq2 ./ ρ.(coords.z)) /
-    norm(y0.ρq2 ./ ρ.(coords.z))
-@test q2_error ≈ 0.0 atol = 0.04
+    norm(sol.u[end].ρq2 ./ ρ_ref.(coords.z) .- y0.ρq2 ./ ρ_ref.(coords.z)) /
+    norm(y0.ρq2 ./ ρ_ref.(coords.z))
+@test q2_error ≈ 0.0 atol = 0.03
 
 q3_error =
-    norm(sol.u[end].ρq3 ./ ρ.(coords.z) .- y0.ρq3 ./ ρ.(coords.z)) /
-    norm(y0.ρq3 ./ ρ.(coords.z))
-@test q3_error ≈ 0.0 atol = 0.6
+    norm(sol.u[end].ρq3 ./ ρ_ref.(coords.z) .- y0.ρq3 ./ ρ_ref.(coords.z)) /
+    norm(y0.ρq3 ./ ρ_ref.(coords.z))
+@test q3_error ≈ 0.0 atol = 0.4
 
 q4_error =
-    norm(sol.u[end].ρq4 ./ ρ.(coords.z) .- y0.ρq4 ./ ρ.(coords.z)) /
-    norm(y0.ρq4 ./ ρ.(coords.z))
-@test q4_error ≈ 0.0 atol = 0.04
+    norm(sol.u[end].ρq4 ./ ρ_ref.(coords.z) .- y0.ρq4 ./ ρ_ref.(coords.z)) /
+    norm(y0.ρq4 ./ ρ_ref.(coords.z))
+@test q4_error ≈ 0.0 atol = 0.03
 
 # visualization artifacts
 ENV["GKSwstype"] = "nul"
@@ -248,14 +273,14 @@ end
 
 Plots.png(
     Plots.plot(
-        sol.u[trunc(Int, end / 2)].ρq3 ./ ρ.(coords.z),
+        sol.u[trunc(Int, end / 2)].ρq3 ./ ρ_ref.(coords.z),
         level = 5,
         clim = (-1, 1),
     ),
     joinpath(path, "q3_6day.png"),
 )
 
 Plots.png(
-    Plots.plot(sol.u[end].ρq3 ./ ρ.(coords.z), level = 5, clim = (-1, 1)),
+    Plots.plot(sol.u[end].ρq3 ./ ρ_ref.(coords.z), level = 5, clim = (-1, 1)),
     joinpath(path, "q3_12day.png"),
 )

examples/3dsphere/implicit_3d_sphere_utils.jl

@@ -456,30 +456,6 @@ function sphere_3D(
     return (hv_center_space, hv_face_space)
 end
 
-
-
-# temporary FieldVector broadcast and fill patches that speeds up solves by 2-3x
-import Base: copyto!, fill!
-using Base.Broadcast: Broadcasted, broadcasted, BroadcastStyle
-transform_broadcasted(bc::Broadcasted{Fields.FieldVectorStyle}, symb, axes) =
-    Broadcasted(
-        bc.f,
-        map(arg -> transform_broadcasted(arg, symb, axes), bc.args),
-        axes,
-    )
-transform_broadcasted(fv::Fields.FieldVector, symb, axes) =
-    parent(getproperty(fv, symb))
-transform_broadcasted(x, symb, axes) = x
-@inline function Base.copyto!(
-    dest::Fields.FieldVector,
-    bc::Broadcasted{Fields.FieldVectorStyle},
-)
-    for symb in propertynames(dest)
-        p = parent(getproperty(dest, symb))
-        copyto!(p, transform_broadcasted(bc, symb, axes(p)))
-    end
-    return dest
-end
 function Base.fill!(a::Fields.FieldVector, x)
     for symb in propertynames(a)
         fill!(parent(getproperty(a, symb)), x)

examples/3dsphere/solid_body_rotation_3d.jl

@@ -264,12 +264,7 @@ T = 3600
 dt = 5
 prob = ODEProblem(rhs!, Y, (0.0, T))
 
-if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
-    throw(:exit_profile)
-end
-
-# solve ode
-sol = solve(
+integrator = OrdinaryDiffEq.init(
     prob,
     SSPRK33(),
     dt = dt,
@@ -279,6 +274,13 @@ sol = solve(
     progress_message = (dt, u, p, t) -> t,
 )
 
+if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
+    throw(:exit_profile)
+end
+
+# solve ode
+sol = @timev OrdinaryDiffEq.solve!(integrator)
+
 uₕ_phy = Geometry.transform.(Ref(Geometry.UVAxis()), sol.u[end].uₕ)
 w_phy = Geometry.transform.(Ref(Geometry.WAxis()), sol.u[end].w)
 

examples/bickleyjet/bickleyjet_cg_invariant_hypervisc.jl

@@ -91,7 +91,7 @@ function rhs!(dydt, y, _, t)
     @. dydt.u =
         wgrad(sdiv(y.u)) -
         Geometry.Covariant12Vector(wcurl(Geometry.Covariant3Vector(curl(y.u))))
-    @. dydt.ρθ = wdiv(grad(y.ρθ))
+    @. dydt.ρθ = wdiv(grad(y.ρθ / y.ρ))
 
     Spaces.weighted_dss!(dydt)
 
@@ -101,7 +101,7 @@ function rhs!(dydt, y, _, t)
                 wcurl(Geometry.Covariant3Vector(curl(dydt.u))),
             )
         )
-    @. dydt.ρθ = -D₄ * wdiv(grad(dydt.ρθ))
+    @. dydt.ρθ = -D₄ * wdiv(y.ρ * grad(dydt.ρθ))
 
     # add in pieces
     @. begin

examples/hybrid/bubble_2d.jl

@@ -295,11 +295,7 @@ using OrdinaryDiffEq
 Δt = 0.03
 prob = ODEProblem(rhs!, Y, (0.0, 500.0))
 
-if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
-    throw(:exit_profile)
-end
-
-sol = @timev solve(
+integrator = OrdinaryDiffEq.init(
     prob,
     SSPRK33(),
     dt = Δt,
@@ -308,6 +304,12 @@ sol = @timev solve(
     progress_message = (dt, u, p, t) -> t,
 );
 
+if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
+    throw(:exit_profile)
+end
+
+sol = @timev OrdinaryDiffEq.solve!(integrator)
+
 ENV["GKSwstype"] = "nul"
 using ClimaCorePlots, Plots
 Plots.GRBackend()

examples/hybrid/bubble_3d.jl

@@ -303,12 +303,7 @@ rhs!(dYdt, Y, nothing, 0.0);
 using OrdinaryDiffEq
 Δt = 0.05
 prob = ODEProblem(rhs!, Y, (0.0, 1.0))
-
-if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
-    throw(:exit_profile)
-end
-
-sol = @timev solve(
+integrator = OrdinaryDiffEq.init(
     prob,
     SSPRK33(),
     dt = Δt,
@@ -317,6 +312,12 @@ sol = @timev solve(
     progress_message = (dt, u, p, t) -> t,
 );
 
+if haskey(ENV, "CI_PERF_SKIP_RUN") # for performance analysis
+    throw(:exit_profile)
+end
+
+sol = @timev OrdinaryDiffEq.solve!(integrator)
+
 ENV["GKSwstype"] = "nul"
 import Plots
 Plots.GRBackend()