Speed up IGW example with threaded mapreduce

Author: charleskawczynski

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}