Merge pull request #12 from CliMA/dy/val_unrolled_reduce

Change val_unrolled_reduce to avoid compilation hanging in ClimaAtmos

Project.toml

@@ -1,7 +1,7 @@
 name = "UnrolledUtilities"
 uuid = "0fe1646c-419e-43be-ac14-22321958931b"
 authors = ["CliMA Contributors <[email protected]>"]
-version = "0.1.3"
+version = "0.1.4"
 
 [compat]
 julia = "1.10"

src/UnrolledUtilities.jl

@@ -22,12 +22,12 @@ export unrolled_any,
     StaticOneTo,
     StaticBitVector
 
+struct NoInit end # Analogue of Base._InitialValue for reduction/accumulation.
+
 include("unrollable_iterator_interface.jl")
 include("recursively_unrolled_functions.jl")
 include("generatively_unrolled_functions.jl")
 
-struct NoInit end # Analogue of Base._InitialValue for reduction/accumulation.
-
 @inline unrolled_any(f, itr) =
     (rec_unroll(itr) ? rec_unrolled_any : gen_unrolled_any)(f, itr)
 @inline unrolled_any(itr) = unrolled_any(identity, itr)
@@ -65,7 +65,7 @@ struct NoInit end # Analogue of Base._InitialValue for reduction/accumulation.
 # TODO: Figure out why unrolled_reduce(op, Val(N), init) compiles faster than
 # unrolled_reduce(op, StaticOneTo(N), init) for the non-orographic gravity wave
 # parametrization test in ClimaAtmos, to the point where the StaticOneTo version
-# appears to completely hang while the Val version compiles in under a second.
+# completely hangs while the Val version compiles in only a few seconds.
 @inline unrolled_reduce(op, val_N::Val, init) =
     val_unrolled_reduce(op, val_N, init)
 @inline unrolled_reduce(op, val_N::Val; init = NoInit()) =

src/generatively_unrolled_functions.jl

@@ -53,8 +53,19 @@ end
 @inline gen_unrolled_accumulate(op, itr, init, transform) =
     _gen_unrolled_accumulate(Val(length(itr)), op, itr, init, transform)
 
-# NOTE: The following is experimental and will likely be removed in the future.
-@inline @generated val_unrolled_reduce(op, ::Val{N}, init) where {N} =
-    foldl(init <: NoInit ? (1:N) : (:init, 1:N...)) do prev_op_expr, item_expr
+# TODO: The following is experimental and will likely be removed in the future.
+# For some reason, combining these two methods into one (or combining them with
+# the method for gen_unrolled_reduce defined above) causes compilation of the
+# non-orographic gravity wave parametrization test in ClimaAtmos to hang. Even
+# more bizarrely, using the assignment form of the first method definition below
+# (as opposed to the function syntax used here) causes compilation to hang as
+# well. This behavior has not yet been replicated in a minimal working example.
+@inline @generated function val_unrolled_reduce(op, ::Val{N}, init) where {N}
+    return foldl((:init, 1:N...)) do prev_op_expr, item_expr
         :(op($prev_op_expr, $item_expr))
-    end # Use foldl instead of reduce to guarantee left associativity.
+    end
+end
+@inline @generated val_unrolled_reduce(op, ::Val{N}, ::NoInit) where {N} =
+    foldl(1:N) do prev_op_expr, item_expr
+        :(op($prev_op_expr, $item_expr))
+    end