Replace UnrolledFunctions with UnrolledUtilities (#1713)

charleskawczynski · dennisYatunin · web-flow · commit 7cc1c9358510 · 2025-02-26T05:47:05.000-08:00
Co-authored-by: Dennis Yatunin &lt;dyatun@gmail.com&gt;
diff --git a/.buildkite/Manifest.toml b/.buildkite/Manifest.toml
@@ -327,10 +327,10 @@ weakdeps = ["CUDA", "MPI"]
     ClimaCommsMPIExt = "MPI"
 
 [[deps.ClimaCore]]
-deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "Unrolled"]
+deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
 path = ".."
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.14.24"
+version = "0.14.26"
 weakdeps = ["CUDA", "Krylov"]
 
     [deps.ClimaCore.extensions]
@@ -344,7 +344,7 @@ uuid = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 version = "0.2.11"
 
 [[deps.ClimaCoreTempestRemap]]
-deps = ["ClimaComms", "ClimaCore", "CommonDataModel", "Dates", "LinearAlgebra", "NCDatasets", "PkgVersion", "TempestRemap_jll"]
+deps = ["ClimaComms", "ClimaCore", "CommonDataModel", "Dates", "DiskArrays", "LinearAlgebra", "NCDatasets", "PkgVersion", "TempestRemap_jll"]
 path = "../lib/ClimaCoreTempestRemap"
 uuid = "d934ef94-cdd4-4710-83d6-720549644b70"
 version = "0.3.18"
@@ -671,9 +671,9 @@ version = "0.6.40"
 
 [[deps.DiskArrays]]
 deps = ["LRUCache", "Mmap", "OffsetArrays"]
-git-tree-sha1 = "64650943240652ebedc6c43d03cccda247b327a3"
+git-tree-sha1 = "4687e77a603fcd86738a92758086717cd06cdaae"
 uuid = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
-version = "0.4.9"
+version = "0.4.8"
 
 [[deps.Distributed]]
 deps = ["Random", "Serialization", "Sockets"]
@@ -2281,11 +2281,14 @@ git-tree-sha1 = "975c354fcd5f7e1ddcc1f1a23e6e091d99e99bc8"
 uuid = "45397f5d-5981-4c77-b2b3-fc36d6e9b728"
 version = "1.6.4"
 
-[[deps.Unrolled]]
-deps = ["MacroTools"]
-git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b"
-uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
-version = "0.1.5"
+[[deps.UnrolledUtilities]]
+git-tree-sha1 = "5caf11dfadeee25daafa7caabb3f252a977ffe72"
+uuid = "0fe1646c-419e-43be-ac14-22321958931b"
+version = "0.1.6"
+weakdeps = ["StaticArrays"]
+
+    [deps.UnrolledUtilities.extensions]
+    UnrolledUtilitiesStaticArraysExt = "StaticArrays"
 
 [[deps.UnsafeAtomics]]
 git-tree-sha1 = "b13c4edda90890e5b04ba24e20a310fbe6f249ff"
diff --git a/Project.toml b/Project.toml
@@ -26,7 +26,7 @@ RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-Unrolled = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
+UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -77,7 +77,7 @@ Statistics = "1"
 StatsBase = "0.34"
 TerminalLoggers = "0.1"
 Test = "1"
-Unrolled = "0.1.5"
+UnrolledUtilities = "0.1.6"
 julia = "1.10"
 
 [extras]
diff --git a/benchmarks/bickleyjet/Manifest.toml b/benchmarks/bickleyjet/Manifest.toml
@@ -232,10 +232,10 @@ version = "0.6.6"
     MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 
 [[deps.ClimaCore]]
-deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "Unrolled"]
+deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
 path = "../.."
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.14.24"
+version = "0.14.26"
 
     [deps.ClimaCore.extensions]
     ClimaCoreCUDAExt = "CUDA"
@@ -1490,11 +1490,14 @@ git-tree-sha1 = "975c354fcd5f7e1ddcc1f1a23e6e091d99e99bc8"
 uuid = "45397f5d-5981-4c77-b2b3-fc36d6e9b728"
 version = "1.6.4"
 
-[[deps.Unrolled]]
-deps = ["MacroTools"]
-git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b"
-uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
-version = "0.1.5"
+[[deps.UnrolledUtilities]]
+git-tree-sha1 = "5caf11dfadeee25daafa7caabb3f252a977ffe72"
+uuid = "0fe1646c-419e-43be-ac14-22321958931b"
+version = "0.1.6"
+weakdeps = ["StaticArrays"]
+
+    [deps.UnrolledUtilities.extensions]
+    UnrolledUtilitiesStaticArraysExt = "StaticArrays"
 
 [[deps.UnsafeAtomics]]
 git-tree-sha1 = "b13c4edda90890e5b04ba24e20a310fbe6f249ff"
diff --git a/docs/src/api.md b/docs/src/api.md
@@ -9,7 +9,6 @@ CurrentModule = ClimaCore
 ```@docs
 Utilities.PlusHalf
 Utilities.half
-Utilities.UnrolledFunctions
 ```
 
 ### Utilities.Cache
diff --git a/ext/cuda/matrix_fields_multiple_field_solve.jl b/ext/cuda/matrix_fields_multiple_field_solve.jl
@@ -6,7 +6,6 @@ import ClimaCore.MatrixFields
 import ClimaCore.MatrixFields: _single_field_solve!
 import ClimaCore.MatrixFields: multiple_field_solve!
 import ClimaCore.MatrixFields: is_CuArray_type
-import ClimaCore.Utilities.UnrolledFunctions: unrolled_map
 
 is_CuArray_type(::Type{T}) where {T <: CUDA.CuArray} = true
 
diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -23,11 +23,12 @@ import ..Grids: ColumnIndex, local_geometry_type
 import ..Spaces: Spaces, AbstractSpace, AbstractPointSpace, cuda_synchronize
 import ..Spaces: nlevels, ncolumns
 import ..Geometry: Geometry, Cartesian12Vector
-import ..Utilities: PlusHalf, half, UnrolledFunctions
+import ..Utilities: PlusHalf, half
 
 using ..RecursiveApply
 using ClimaComms
 import Adapt
+import UnrolledUtilities: unrolled_map
 
 import StaticArrays, LinearAlgebra, Statistics, InteractiveUtils
 
diff --git a/src/Fields/fieldvector.jl b/src/Fields/fieldvector.jl
@@ -429,7 +429,7 @@ LinearAlgebra.ldiv!(A::LinearAlgebra.LU, x::FieldVector) =
     x .= LinearAlgebra.ldiv!(A, Vector(x))
 
 function LinearAlgebra.norm_sqr(x::FieldVector)
-    value_norm_sqrs = UnrolledFunctions.unrolled_map(_values(x)) do value
+    value_norm_sqrs = unrolled_map(_values(x)) do value
         LinearAlgebra.norm_sqr(backing_array(value))
     end
     return sum(value_norm_sqrs; init = zero(eltype(x)))
@@ -440,9 +440,8 @@ end
 
 import ClimaComms
 
-ClimaComms.array_type(x::FieldVector) = promote_type(
-    UnrolledFunctions.unrolled_map(ClimaComms.array_type, _values(x))...,
-)
+ClimaComms.array_type(x::FieldVector) =
+    promote_type(unrolled_map(ClimaComms.array_type, _values(x))...)
 
 function __rprint_diff(
     io::IO,
diff --git a/src/MatrixFields/MatrixFields.jl b/src/MatrixFields/MatrixFields.jl
@@ -52,6 +52,7 @@ import KrylovKit
 import ClimaComms
 import NVTX
 import Adapt
+using UnrolledUtilities
 
 import ..Utilities: PlusHalf, half
 import ..RecursiveApply:
@@ -65,8 +66,6 @@ import ..Spaces
 import ..Spaces: local_geometry_type
 import ..Fields
 import ..Operators
-
-using ..Utilities.UnrolledFunctions
 using ..Geometry:
     rmul_with_projection,
     mul_with_projection,
diff --git a/src/MatrixFields/field_name.jl b/src/MatrixFields/field_name.jl
@@ -160,10 +160,11 @@ get_subtree_at_name(name, tree) =
     if name == tree.name
         tree
     else
-        subtree = unrolled_findonly(tree.subtrees) do subtree
+        subtrees_at_name = unrolled_filter(tree.subtrees) do subtree
             is_valid_name(name, subtree)
         end
-        get_subtree_at_name(name, subtree)
+        @assert length(subtrees_at_name) == 1
+        get_subtree_at_name(name, subtrees_at_name[1])
     end
 
 ################################################################################
diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl
@@ -137,7 +137,7 @@ Base.:(==)(dict1::FieldNameDict, dict2::FieldNameDict) =
 function Base.getindex(dict::FieldNameDict, key)
     key in keys(dict) || throw(KeyError(key))
     key′, entry′ =
-        unrolled_findonly(pair -> is_child_value(key, pair[1]), pairs(dict))
+        unrolled_filter(pair -> is_child_value(key, pair[1]), pairs(dict))[1]
     internal_key = get_internal_key(key, key′)
     return get_internal_entry(entry′, internal_key, KeyError(key))
 end
diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl
@@ -1,4 +1,5 @@
-import ..Utilities: PlusHalf, half, UnrolledFunctions
+import ..Utilities: PlusHalf, half
+import UnrolledUtilities: unrolled_map
 
 const AllFiniteDifferenceSpace =
     Union{Spaces.FiniteDifferenceSpace, Spaces.ExtrudedFiniteDifferenceSpace}
@@ -3324,10 +3325,7 @@ Adapt.adapt_structure(to, op::FiniteDifferenceOperator) =
     unionall_type(typeof(op))(; adapt_bcs(to, bcs)...)
 
 @inline adapt_bcs(to, bcs) = NamedTuple{keys(bcs)}(
-    UnrolledFunctions.unrolled_map(
-        bc -> Adapt.adapt_structure(to, bc),
-        values(bcs),
-    ),
+    unrolled_map(bc -> Adapt.adapt_structure(to, bc), values(bcs)),
 )
 
 """
diff --git a/src/Utilities/Utilities.jl b/src/Utilities/Utilities.jl
@@ -1,7 +1,6 @@
 module Utilities
 
 include("plushalf.jl")
-include("unrolled_functions.jl")
 include("cache.jl")
 
 """
diff --git a/src/Utilities/unrolled_functions.jl b/src/Utilities/unrolled_functions.jl
