Merge pull request #23 from CliMA/dy/ntuple_bugfix

Fix recursion instability in unrolled_product

Author: dennisYatunin

Project.toml

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

src/UnrolledUtilities.jl

@@ -399,8 +399,10 @@ include("generatively_unrolled_functions.jl")
 @inline unrolled_product(itrs...) =
     ntuple(Val(unrolled_prod(length, itrs))) do n
         @inline
+        Base.@assume_effects :foldable
         items = ntuple(Val(length(itrs))) do itr_index
             @inline
+            Base.@assume_effects :foldable
             cur_length = length(itrs[itr_index])
             prev_length = unrolled_prod(length, itrs[1:(itr_index - 1)])
             item_index = (n - 1) ÷ prev_length % cur_length + 1

test/test_and_analyze.jl

@@ -289,11 +289,16 @@ macro test_unrolled(
             isempty(JET.get_reports(@report_opt reference_func($(args...))))
         $(esc(skip_type_stability_test)) || @test_opt unrolled_func($(args...))
 
+        unrolled_code = code_instance(unrolled_func, $(args...))
+        reference_code = code_instance(reference_func, $(args...))
+
         # Test for constant propagation.
         is_unrolled_const =
-            isdefined(code_instance(unrolled_func, $(args...)), :rettype_const)
+            isdefined(unrolled_code, :rettype_const) &&
+            isbits(unrolled_code.rettype_const)
         is_reference_const =
-            isdefined(code_instance(reference_func, $(args...)), :rettype_const)
+            isdefined(reference_code, :rettype_const) &&
+            isbits(reference_code.rettype_const)
 
         buffer = IOBuffer()
         args_type = Tuple{map(typeof, ($(args...),))...}
@@ -761,12 +766,15 @@ for itr in (
             str,
         )
 
+        # TODO: Testing with coverage triggers allocations for unrolled_product!
         if length(itr) <= 32
             @test_unrolled(
                 (itr,),
                 unrolled_product(itr, itr),
                 Tuple(Iterators.product(itr, itr)),
                 str,
+                "fast_mode" in ARGS,
+                "fast_mode" in ARGS,
             )
         end # This can take several minutes to compile when the length is 128.
         if length(itr) <= 8
@@ -775,6 +783,8 @@ for itr in (
                 unrolled_product(itr, itr, itr),
                 Tuple(Iterators.product(itr, itr, itr)),
                 str,
+                "fast_mode" in ARGS,
+                "fast_mode" in ARGS,
             )
         end # This can take several minutes to compile when the length is 32.