Merge #1228 #1263

1228: Fix dss_transform ambiguity aqua test r=charleskawczynski a=charleskawczynski

This PR fixes a method ambiguity for `dss_transform`-- perhaps not too important, but 🤷🏻.

1263: Speed up inertial gravity wave examples r=charleskawczynski a=charleskawczynski

CI has been slow lately, and failures due to unknown reasons (e.g., the recent [mac os timeout](https://github.com/CliMA/ClimaCore.jl/actions/runs/5071735121/jobs/9108493197)) is painful when we have long chains of serial tasks. I've been working on this branch for a while now, mostly while I'm waiting for CI to pass and I think it's finally ready. The main goal is to speed up the inertial gravity wave examples, which are very slow.

This PR attempts to speed up the inertial gravity wave examples by accelerating the Bretherton computations using a threaded mapreduce.

Supersedes #849.

Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

Author: bors[bot]

.buildkite/pipeline.yml

@@ -780,6 +780,7 @@ steps:
           TEST_NAME: "plane/inertial_gravity_wave"
         agents:
           slurm_cpus_per_task: 8
+          slurm_mem: 20GB
 
       - label: ":computer: stretched 2D plane inertial gravity wave"
         key: "cpu_stretch_inertial_gravity_wave"
@@ -792,6 +793,7 @@ steps:
           Z_STRETCH: "true"
         agents:
           slurm_cpus_per_task: 8
+          slurm_mem: 20GB
 
   - group: "Performance"
     steps:

examples/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.8.5"
 manifest_format = "2.0"
-project_hash = "7ec2fabea202c1fbfc0bf694c476ce42ac582098"
+project_hash = "1e5a4795567d9d355e43d11e9245fcf683ab6bb8"
 
 [[deps.AMD]]
 deps = ["Libdl", "LinearAlgebra", "SparseArrays", "Test"]
@@ -27,6 +27,11 @@ git-tree-sha1 = "cc37d689f599e8df4f464b2fa3870ff7db7492ef"
 uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 version = "3.6.1"
 
+[[deps.ArgCheck]]
+git-tree-sha1 = "a3a402a35a2f7e0b87828ccabbd5ebfbebe356b4"
+uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+version = "2.3.0"
+
 [[deps.ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
 version = "1.1.1"
@@ -88,9 +93,20 @@ git-tree-sha1 = "6ef8fc1d77b60f41041d59ce61ef9eb41ed97a83"
 uuid = "aae01518-5342-5314-be14-df237901396f"
 version = "0.17.18"
 
+[[deps.BangBang]]
+deps = ["Compat", "ConstructionBase", "InitialValues", "LinearAlgebra", "Requires", "Setfield", "Tables"]
+git-tree-sha1 = "54b00d1b93791f8e19e31584bd30f2cb6004614b"
+uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+version = "0.3.38"
+
 [[deps.Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
+[[deps.Baselet]]
+git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e"
+uuid = "9718e550-a3fa-408a-8086-8db961cd8217"
+version = "0.1.1"
+
 [[deps.BitFlags]]
 git-tree-sha1 = "43b1a4a8f797c1cddadf60499a8a077d4af2cd2d"
 uuid = "d1d4a3ce-64b1-5f1a-9ba4-7e7e69966f35"
@@ -195,7 +211,7 @@ version = "0.4.2"
 deps = ["Adapt", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DiffEqBase", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "Rotations", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"]
 path = ".."
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.10.33"
+version = "0.10.36"
 
 [[deps.ClimaCorePlots]]
 deps = ["ClimaCore", "RecipesBase", "StaticArrays", "TriplotBase"]
@@ -204,7 +220,7 @@ uuid = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 version = "0.2.5"
 
 [[deps.ClimaCoreTempestRemap]]
-deps = ["ClimaComms", "ClimaCore", "Dates", "LinearAlgebra", "MPI", "NCDatasets", "PkgVersion", "TempestRemap_jll", "Test"]
+deps = ["ClimaComms", "ClimaCore", "Dates", "LinearAlgebra", "NCDatasets", "PkgVersion", "TempestRemap_jll", "Test"]
 path = "../lib/ClimaCoreTempestRemap"
 uuid = "d934ef94-cdd4-4710-83d6-720549644b70"
 version = "0.3.8"
@@ -290,6 +306,11 @@ git-tree-sha1 = "02d2316b7ffceff992f3096ae48c7829a8aa0638"
 uuid = "b152e2b5-7a66-4b01-a709-34e65c35f657"
 version = "0.1.3"
 
+[[deps.CompositionsBase]]
+git-tree-sha1 = "802bb88cd69dfd1509f6670416bd4434015693ad"
+uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b"
+version = "0.1.2"
+
 [[deps.ConcurrentUtilities]]
 deps = ["Serialization", "Sockets"]
 git-tree-sha1 = "b306df2650947e9eb100ec125ff8c65ca2053d30"
@@ -345,6 +366,11 @@ version = "1.0.0"
 deps = ["Printf"]
 uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
+[[deps.DefineSingletons]]
+git-tree-sha1 = "0fba8b706d0178b4dc7fd44a96a92382c9065c2c"
+uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52"
+version = "0.1.2"
+
 [[deps.DelimitedFiles]]
 deps = ["Mmap"]
 uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
@@ -681,6 +707,11 @@ git-tree-sha1 = "5cd07aab533df5170988219191dfad0519391428"
 uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9"
 version = "0.1.3"
 
+[[deps.InitialValues]]
+git-tree-sha1 = "4da0f88e9a39111c2fa3add390ab15f3a44f3ca3"
+uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
+version = "0.3.1"
+
 [[deps.IntelOpenMP_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
 git-tree-sha1 = "0cb9352ef2e01574eeebdb102948a58740dcaf83"
@@ -1028,6 +1059,12 @@ git-tree-sha1 = "c13304c81eec1ed3af7fc20e75fb6b26092a1102"
 uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e"
 version = "0.3.2"
 
+[[deps.MicroCollections]]
+deps = ["BangBang", "InitialValues", "Setfield"]
+git-tree-sha1 = "629afd7d10dbc6935ec59b32daeb33bc4460a42e"
+uuid = "128add7d-3638-4c79-886c-908ea0c25c34"
+version = "0.1.4"
+
 [[deps.MicrosoftMPI_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
 git-tree-sha1 = "a8027af3d1743b3bfae34e54872359fdebb31422"
@@ -1361,6 +1398,12 @@ git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
 uuid = "189a3867-3050-52da-a836-e630ba90ab69"
 version = "1.2.2"
 
+[[deps.Referenceables]]
+deps = ["Adapt"]
+git-tree-sha1 = "e681d3bfa49cd46c3c161505caddf20f0e62aaa9"
+uuid = "42d2dcc6-99eb-4e98-b66c-637b7d73030e"
+version = "0.1.2"
+
 [[deps.RelocatableFolders]]
 deps = ["SHA", "Scratch"]
 git-tree-sha1 = "90bc7a7c96410424509e4263e277e43250c05691"
@@ -1497,6 +1540,12 @@ git-tree-sha1 = "ef28127915f4229c971eb43f3fc075dd3fe91880"
 uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
 version = "2.2.0"
 
+[[deps.SplittablesBase]]
+deps = ["Setfield", "Test"]
+git-tree-sha1 = "e08a62abc517eb79667d0a29dc08a3b589516bb5"
+uuid = "171d559e-b47b-412a-8079-5efa626c420e"
+version = "0.1.15"
+
 [[deps.Static]]
 deps = ["IfElse"]
 git-tree-sha1 = "7f5a513baec6f122401abfc8e9c074fdac54f6c1"
@@ -1575,7 +1624,7 @@ uuid = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
 version = "0.12.2"
 
 [[deps.TempestRemap_jll]]
-deps = ["Artifacts", "JLLWrappers", "Libdl", "NetCDF_jll", "OpenBLAS32_jll", "Pkg"]
+deps = ["Artifacts", "HDF5_jll", "JLLWrappers", "Libdl", "NetCDF_jll", "OpenBLAS32_jll", "Pkg"]
 git-tree-sha1 = "88c3818a492ad1a94b1aa440b01eab5d133209ff"
 uuid = "8573a8c5-1df0-515e-a024-abad257ee284"
 version = "2.1.6+1"
@@ -1602,6 +1651,12 @@ git-tree-sha1 = "c97f60dd4f2331e1a495527f80d242501d2f9865"
 uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5"
 version = "0.5.1"
 
+[[deps.ThreadsX]]
+deps = ["ArgCheck", "BangBang", "ConstructionBase", "InitialValues", "MicroCollections", "Referenceables", "Setfield", "SplittablesBase", "Transducers"]
+git-tree-sha1 = "34e6bcf36b9ed5d56489600cf9f3c16843fa2aa2"
+uuid = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
+version = "0.1.11"
+
 [[deps.TimerOutputs]]
 deps = ["ExprTools", "Printf"]
 git-tree-sha1 = "f548a9e9c490030e545f72074a41edfd0e5bcdd7"
@@ -1614,6 +1669,12 @@ git-tree-sha1 = "9a6ae7ed916312b41236fcef7e0af564ef934769"
 uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
 version = "0.9.13"
 
+[[deps.Transducers]]
+deps = ["Adapt", "ArgCheck", "BangBang", "Baselet", "CompositionsBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "Setfield", "SplittablesBase", "Tables"]
+git-tree-sha1 = "25358a5f2384c490e98abd565ed321ffae2cbb37"
+uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999"
+version = "0.4.76"
+
 [[deps.TriangularSolve]]
 deps = ["CloseOpenIntervals", "IfElse", "LayoutPointers", "LinearAlgebra", "LoopVectorization", "Polyester", "Static", "VectorizationBase"]
 git-tree-sha1 = "31eedbc0b6d07c08a700e26d31298ac27ef330eb"

examples/Project.toml

@@ -28,6 +28,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [compat]

examples/hybrid/plane/inertial_gravity_wave.jl

@@ -282,6 +282,11 @@ function linear_solution_cache(ᶜlocal_geometry, ᶠlocal_geometry)
     @time "ρfb_init_coefs!" IGWU.ρfb_init_coefs!(FT, ρfb_init_array_params)
     (; ρfb_init_array, ᶜρb_init_xz, unit_integral) = ρfb_init_array_params
     max_ikx, max_ikz = (size(ρfb_init_array) .- 1) .÷ 2
+
+    get_xz(lg) = (; x = lg.coordinates.x, z = lg.coordinates.z)
+    ᶠxz = get_xz.(ᶠlocal_geometry)
+    ᶜxz = get_xz.(ᶜlocal_geometry)
+
     ᶜp₀ = @. p₀(ᶜz)
     return (;
         # globals
@@ -305,6 +310,8 @@ function linear_solution_cache(ᶜlocal_geometry, ᶠlocal_geometry)
         ᶠx = ᶠlocal_geometry.coordinates.x,
         ᶜz,
         ᶠz,
+        ᶜxz,
+        ᶠxz,
 
         # background state
         ᶜp₀,

examples/hybrid/plane/intertial_gravity_wave_utils.jl

@@ -35,15 +35,83 @@ function ρfb_init_coefs!(::Type{FT}, params) where {FT}
 end
 
 function Bretherton_transforms!(lin_cache, t, ::Type{FT}) where {FT}
-    # Bretherton_transforms_partial_sums! is fastest because
-    # we can multithread across
-    #       `Iterators.product((-max_ikx):max_ikx, (-max_ikz):max_ikz)`
-    # and apply sums for center and face fields. Using mapreduce requires
-    # two calls and, as a result in ~20 slower.
-
-    Bretherton_transforms_original!(lin_cache, t, FT)
-    # Bretherton_transforms_partial_sums!(lin_cache, t, FT)
-    # Bretherton_transforms_threaded_mapreduce!(lin_cache, t, FT)
+    # Bretherton_transforms! is the most computationally
+    # expensive part of this example and was therefore
+    # optimized a bit.
+    # Bretherton_transforms_original!(lin_cache, t, FT)
+    Bretherton_transforms_threaded_mapreduce!(lin_cache, t, FT)
+end
+
+import ThreadsX
+function Bretherton_transforms_threaded_mapreduce!(
+    lin_cache,
+    t,
+    ::Type{FT},
+) where {FT}
+    # @info "Computing Bretherton_transforms! (threaded mapreduce)..."
+    (; ᶠwb) = lin_cache
+    (; ᶜx, ᶠx, ᶜz, ᶠz, ᶠxz, ᶜxz) = lin_cache
+    (; ᶜρb_init_xz, ρfb_init_array, unit_integral, x_max, z_max) = lin_cache
+    (; max_ikx, max_ikz, u₀) = lin_cache
+    combine(ᶜpb, ᶜρb, ᶜub, ᶜvb) = (; ᶜpb, ᶜρb, ᶜub, ᶜvb)
+    ᶜbretherton_fields =
+        combine.(lin_cache.ᶜpb, lin_cache.ᶜρb, lin_cache.ᶜub, lin_cache.ᶜvb)
+
+    # TODO: could we, and is it advantageous to, combine
+    #       this into a single mapreduce call?
+    ip = Iterators.product((-max_ikx):max_ikx, (-max_ikz):max_ikz)
+    bc_add = (a, b) -> a .+ b
+    ᶠwb .= ThreadsX.mapreduce(
+        bc_add,
+        ip;
+        init = zeros(FT, axes(ᶠwb)),
+    ) do (ikx, ikz)
+        (; pfb, ρfb, ufb, vfb, wfb) =
+            Bretherton_transform_coeffs(lin_cache, ikx, ikz, t, FT)
+
+        # Fourier coefficient of ᶜρb_init (for current kx and kz)
+        kx::FT = 2 * π / x_max * ikx
+        kz::FT = 2 * π / (2 * z_max) * ikz
+
+        # Fourier factors, shifted by u₀ * t along the x-axis
+        map(ᶠxz) do nt
+            real(wfb * exp(im * (kx * (nt.x - u₀ * t) + kz * nt.z)))
+        end
+    end
+
+    bc_add = (a, b) -> a .+ b
+    zeroᶜbretherton_fields =
+        zeros(eltype(ᶜbretherton_fields), axes(ᶜbretherton_fields))
+    ᶜbretherton_fields .= ThreadsX.mapreduce(
+        bc_add,
+        ip;
+        init = zeroᶜbretherton_fields,
+    ) do (ikx, ikz)
+        (; pfb, ρfb, ufb, vfb, wfb) =
+            Bretherton_transform_coeffs(lin_cache, ikx, ikz, t, FT)
+
+        # Fourier coefficient of ᶜρb_init (for current kx and kz)
+        kx::FT = 2 * π / x_max * ikx
+        kz::FT = 2 * π / (2 * z_max) * ikz
+
+        # Fourier factors, shifted by u₀ * t along the x-axis
+        map(ᶜxz) do nt
+            ᶜpb::FT = real(pfb * exp(im * (kx * (nt.x - u₀ * t) + kz * nt.z)))
+            ᶜρb::FT =
+                real(ρfb * exp(im * (kx * (nt.x - u₀ * t) + kz * nt.z)))
+            ᶜub::FT =
+                real(ufb * exp(im * (kx * (nt.x - u₀ * t) + kz * nt.z)))
+            ᶜvb::FT =
+                real(vfb * exp(im * (kx * (nt.x - u₀ * t) + kz * nt.z)))
+            (; ᶜpb, ᶜρb, ᶜub, ᶜvb)
+        end
+    end
+
+    lin_cache.ᶜpb .= ᶜbretherton_fields.ᶜpb
+    lin_cache.ᶜρb .= ᶜbretherton_fields.ᶜρb
+    lin_cache.ᶜub .= ᶜbretherton_fields.ᶜub
+    lin_cache.ᶜvb .= ᶜbretherton_fields.ᶜvb
+    return nothing
 end
 
 function Bretherton_transforms_original!(lin_cache, t, ::Type{FT}) where {FT}

src/Spaces/dss_transform.jl

@@ -68,6 +68,11 @@ end
     local_geometry::Geometry.LocalGeometry,
     weight,
 ) where {T, N} = arg * weight
+@inline dss_transform(
+    arg::Geometry.AxisTensor{T, N, <:Tuple{}},
+    local_geometry::Geometry.LocalGeometry,
+    weight,
+) where {T, N} = arg * weight
 @inline dss_transform(
     arg::Geometry.LocalVector,
     local_geometry::Geometry.LocalGeometry,

test/aqua.jl

@@ -20,7 +20,7 @@ using Aqua
     # then please lower the limit based on the new number of ambiguities.
     # We're trying to drive this number down to zero to reduce latency.
     @info "Number of method ambiguities: $(length(ambs))"
-    @test length(ambs) ≤ 16
+    @test length(ambs) ≤ 15
 
     # returns a vector of all unbound args
     # ua = Aqua.detect_unbound_args_recursively(ClimaCore)