Replace Unrolled with UnrolledUtilities

Author: charleskawczynski

Project.toml

@@ -27,7 +27,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"
@@ -45,13 +45,13 @@ AssociatedLegendrePolynomials = "1"
 BandedMatrices = "0.17, 1"
 BenchmarkTools = "1"
 BlockArrays = "0.16, 1"
+CUDA = "5"
 ClimaComms = "0.6"
 Combinatorics = "1"
 CountFlops = "0.1"
 CubedSphere = "0.2, 0.3"
-CUDA = "5"
-Dates = "1"
 DataStructures = "0.18"
+Dates = "1"
 DocStringExtensions = "0.8, 0.9"
 FastBroadcast = "0.3"
 ForwardDiff = "0.10"
@@ -63,8 +63,8 @@ IntervalSets = "0.5, 0.6, 0.7"
 JET = "0.9"
 Krylov = "0.9"
 KrylovKit = "0.6, 0.7, 0.8"
-LinearAlgebra = "1"
 LazyBroadcast = "0.1"
+LinearAlgebra = "1"
 Logging = "1"
 MPI = "0.20"
 MultiBroadcastFusion = "0.3"
@@ -82,7 +82,6 @@ Statistics = "1"
 StatsBase = "0.34"
 TerminalLoggers = "0.1"
 Test = "1"
-Unrolled = "0.1"
 julia = "1.9"
 
 [extras]
@@ -94,8 +93,8 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 CountFlops = "1db9610d-79e1-487a-8d40-77f3295c7593"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
-Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
@@ -109,4 +108,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Aqua", "ArgParse", "AssociatedLegendrePolynomials", "BenchmarkTools", "Combinatorics", "CountFlops", "Dates", "FastBroadcast", "Krylov", "JET", "LazyBroadcast", "Logging", "MPI", "OrderedCollections", "PrettyTables", "Random", "SafeTestsets", "StatsBase", "TerminalLoggers", "Test"]
-

docs/src/api.md

@@ -9,7 +9,6 @@ CurrentModule = ClimaCore
 ```@docs
 Utilities.PlusHalf
 Utilities.half
-Utilities.UnrolledFunctions
 ```
 
 ### Utilities.Cache

ext/cuda/matrix_fields_multiple_field_solve.jl

@@ -6,7 +6,7 @@ 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
+import UnrolledUtilities: unrolled_map
 
 is_CuArray_type(::Type{T}) where {T <: CUDA.CuArray} = true
 

src/Fields/Fields.jl

@@ -17,7 +17,8 @@ import ..Quadratures
 import ..Grids: ColumnIndex, local_geometry_type
 import ..Spaces: Spaces, AbstractSpace, AbstractPointSpace, cuda_synchronize
 import ..Geometry: Geometry, Cartesian12Vector
-import ..Utilities: PlusHalf, half, UnrolledFunctions
+import ..Utilities: PlusHalf, half
+import UnrolledUtilities
 
 using ..RecursiveApply
 using ClimaComms

src/Fields/fieldvector.jl

@@ -352,7 +352,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 = UnrolledUtilities.unrolled_map(_values(x)) do value
         LinearAlgebra.norm_sqr(backing_array(value))
     end
     return sum(value_norm_sqrs; init = zero(eltype(x)))
@@ -364,7 +364,7 @@ end
 import ClimaComms
 
 ClimaComms.array_type(x::FieldVector) = promote_type(
-    UnrolledFunctions.unrolled_map(ClimaComms.array_type, _values(x))...,
+    UnrolledUtilities.unrolled_map(ClimaComms.array_type, _values(x))...,
 )
 
 function __rprint_diff(

src/MatrixFields/MatrixFields.jl

@@ -66,7 +66,20 @@ import ..Spaces: local_geometry_type
 import ..Fields
 import ..Operators
 
-using ..Utilities.UnrolledFunctions
+using UnrolledUtilities:
+    unrolled_map,
+    unrolled_take,
+    unrolled_drop,
+    unrolled_any,
+    unrolled_foreach,
+    unrolled_filter,
+    unrolled_flatmap,
+    unrolled_all,
+    unrolled_in,
+    unrolled_product,
+    unrolled_unique,
+    unrolled_split
+
 using ..Geometry:
     rmul_with_projection,
     mul_with_projection,

src/MatrixFields/field_name.jl

@@ -144,15 +144,18 @@ function child_names(name, tree)
     subtree isa FieldNameTreeNode || error("$name does not have child names")
     return unrolled_map(subsubtree -> subsubtree.name, subtree.subtrees)
 end
-get_subtree_at_name(name, tree) =
+function get_subtree_at_name(name, tree)
     if name == tree.name
         tree
     else
-        subtree = unrolled_findonly(tree.subtrees) do subtree
+        filtered_values = unrolled_filter(tree.subtrees) do subtree
             is_valid_name(name, subtree)
         end
+        @assert length(filtered_values) == 1
+        subtree = filtered_values[1]
         get_subtree_at_name(name, subtree)
     end
+end
 
 ################################################################################
 

src/MatrixFields/field_name_dict.jl

@@ -120,8 +120,11 @@ 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))
+    filtered_values = unrolled_filter(pairs(dict)) do pair
+        is_child_value(key, pair[1])
+    end
+    @assert length(filtered_values) == 1
+    key′, entry′ = filtered_values[1]
     return get_internal_entry(entry′, get_internal_key(key, key′))
 end
 

src/Operators/finitedifference.jl

@@ -1,4 +1,5 @@
-import ..Utilities: PlusHalf, half, UnrolledFunctions
+import ..Utilities: PlusHalf, half
+import UnrolledUtilities
 
 const AllFiniteDifferenceSpace =
     Union{Spaces.FiniteDifferenceSpace, Spaces.ExtrudedFiniteDifferenceSpace}
@@ -2624,7 +2625,7 @@ end
 function Adapt.adapt_structure(to, op::FiniteDifferenceOperator)
     if hasfield(typeof(op), :bcs)
         bcs_adapted = NamedTuple{keys(op.bcs)}(
-            UnrolledFunctions.unrolled_map(
+            UnrolledUtilities.unrolled_map(
                 bc -> Adapt.adapt_structure(to, bc),
                 values(op.bcs),
             ),

src/Utilities/Utilities.jl

@@ -1,7 +1,6 @@
 module Utilities
 
 include("plushalf.jl")
-include("unrolled_functions.jl")
 include("cache.jl")
 
 """