Replace recursive unrolling with manual unrolling

Author: dennisYatunin

docs/src/developer_guide.md

@@ -5,7 +5,7 @@ CurrentModule = UnrolledUtilities
 ## How to Unroll
 
 There are two general ways to implement loop unrolling in Julia—recursively
-splatting iterator contents and manually generating unrolled expressions. For
+splatting iterator contents and manually constructing unrolled expressions. For
 example, a recursively unrolled version of the `foreach` function is
 
 ```julia
@@ -14,23 +14,42 @@ _unrolled_foreach(f) = nothing
 _unrolled_foreach(f, item, items...) = (f(item); _unrolled_foreach(f, items...))
 ```
 
-In contrast, a generatively unrolled implementation of this function looks like
+In contrast, a manually unrolled implementation of this function looks like
 
 ```julia
 unrolled_foreach(f, itr) = _unrolled_foreach(Val(length(itr)), f, itr)
 @generated _unrolled_foreach(::Val{N}, f, itr) where {N} =
     Expr(:block, (:(f(generic_getindex(itr, $n))) for n in 1:N)..., nothing)
 ```
 
-To switch between recursive and generative unrolling, this package defines the
-following function:
+Julia's compiler can only pass up to 32 values through function arguments
+without allocating heap memory, so recursive unrolling is not type-stable for
+iterators with lengths greater than 32. However, automatically generating
+functions often requires more time and memory resources during compilation than
+writing hard-coded functions. Recursive inlining adds overhead to compilation
+as well, but this is typically smaller than the overhead of generated functions
+for short iterators. To avoid sacrificing latency by using generated functions,
+several hard-coded methods can be added to the manually unrolled implementation:
 
-```@docs
-rec_unroll
+```julia
+_unrolled_foreach(::Val{0}, f, itr) = nothing
+_unrolled_foreach(::Val{1}, f, itr) = (f(generic_getindex(itr, 1)); nothing)
+_unrolled_foreach(::Val{2}, f, itr) =
+    (f(generic_getindex(itr, 1)); f(generic_getindex(itr, 2)); nothing)
+_unrolled_foreach(::Val{3}, f, itr) =
+    (f(generic_getindex(itr, 1)); f(generic_getindex(itr, 2)); f(generic_getindex(itr, 3)); nothing)
 ```
 
-The default choice for `rec_unroll` is motivated by the benchmarks for
-[Generative vs. Recursive Unrolling](@ref).
+With this modification, manual unrolling does not exceed the compilation
+requirements of recursive unrolling across a wide range of use cases. Since it
+also avoids type instabilities for arbitrarily large iterators, a combination
+of hard-coded and generated functions with manual unrolling serves as the basis
+of all unrolled functions defined in this package. Similarly, the
+[`ntuple(f, ::Val{N})`](https://github.com/JuliaLang/julia/blob/v1.11.0/base/ntuple.jl)
+function in `Base` uses this strategy to implement loop unrolling.
+
+For benchmarks that compare these two implementations, see
+[Manual vs. Recursive Unrolling](@ref).
 
 ## Interface API
 
@@ -54,10 +73,8 @@ StaticSequence
 
 To unroll over a statically sized iterator of some user-defined type `T`, follow
 these steps:
-- To enable recursive unrolling, add a method for `iterate(::T, [state])`
-- To enable generative unrolling, add a method for `getindex(::T, n)` (or for
-  `generic_getindex(::T, n)` if `getindex` should not be defined for iterators
-  of type `T`)
+- Add a method for `getindex(::T, n)`, or for `generic_getindex(::T, n)` if
+  `getindex` should not be defined for iterators of type `T`
 - If every unrolled function that needs to construct an iterator when given an
   iterator of type `T` can return a `Tuple` instead, stop here
 - Otherwise, to return a non-`Tuple` iterator whenever it is efficient to do so,

src/UnrolledUtilities.jl

@@ -53,6 +53,14 @@ include("unrollable_iterator_interface.jl")
 # Analogue of the non-public Base._InitialValue for reduction and accumulation.
 struct NoInit end
 
+@inline empty_reduction_value(::NoInit) =
+    error("unrolled_reduce requires an init value for empty iterators")
+@inline empty_reduction_value(init) = init
+
+@inline first_reduction_value(op, itr, ::NoInit) = generic_getindex(itr, 1)
+@inline first_reduction_value(op, itr, init) =
+    op(init, generic_getindex(itr, 1))
+
 # Analogue of ∘, but with only one function argument and guaranteed inlining.
 # Base's ∘ leads to type instabilities in unit tests on Julia 1.10 and 1.11.
 @inline ⋅(f1::F1, f2::F2) where {F1, F2} = x -> (@inline f1(f2(x)))
@@ -134,17 +142,17 @@ include("StaticBitVector.jl")
     unrolled_drop_into(inferred_output_type(itr), itr, val_N)
 
 ##
-## Functions unrolled using either recursion or generated expressions
+## Functions unrolled manually (hard-coded for N <= 3, auto-generated for N > 3)
 ##
 
-include("recursively_unrolled_functions.jl")
+include("hard_coded_unrolled_functions.jl")
 include("generatively_unrolled_functions.jl")
 
 # The unrolled_map function could also be implemented in terms of ntuple, but
 # then it would be subject to the same recursion limit as ntuple. On Julia 1.10,
 # this leads to type instabilities in several unit tests for nested iterators.
 @inline unrolled_map_into_tuple(f::F, itr) where {F} =
-    (rec_unroll(itr) ? rec_unrolled_map : gen_unrolled_map)(f, itr)
+    _unrolled_map(Val(length(itr)), f, itr)
 @inline unrolled_map_into(output_type, f::F, itr) where {F} =
     constructor_from_tuple(output_type)(unrolled_map_into_tuple(f, itr))
 @inline unrolled_map(f::F, itr) where {F} =
@@ -154,32 +162,26 @@ include("generatively_unrolled_functions.jl")
 
 @inline unrolled_any(itr) = unrolled_any(identity, itr)
 @inline unrolled_any(f::F, itr) where {F} =
-    (rec_unroll(itr) ? rec_unrolled_any : gen_unrolled_any)(f, itr)
+    _unrolled_any(Val(length(itr)), f, itr)
 
 @inline unrolled_all(itr) = unrolled_all(identity, itr)
 @inline unrolled_all(f::F, itr) where {F} =
-    (rec_unroll(itr) ? rec_unrolled_all : gen_unrolled_all)(f, itr)
+    _unrolled_all(Val(length(itr)), f, itr)
 
 @inline unrolled_foreach(f::F, itr) where {F} =
-    (rec_unroll(itr) ? rec_unrolled_foreach : gen_unrolled_foreach)(f, itr)
+    _unrolled_foreach(Val(length(itr)), f, itr)
 @inline unrolled_foreach(f, itrs...) = unrolled_foreach(splat(f), zip(itrs...))
 
 @inline unrolled_reduce(op::O, itr, init) where {O} =
-    isempty(itr) && init isa NoInit ?
-    error("unrolled_reduce requires an init value for empty iterators") :
-    (rec_unroll(itr) ? rec_unrolled_reduce : gen_unrolled_reduce)(op, itr, init)
+    _unrolled_reduce(Val(length(itr)), op, itr, init)
 @inline unrolled_reduce(op::O, itr; init = NoInit()) where {O} =
     unrolled_reduce(op, itr, init)
 
 @inline unrolled_mapreduce(f::F, op::O, itrs...; init = NoInit()) where {F, O} =
     unrolled_reduce(op, unrolled_map(f, itrs...), init)
 
 @inline unrolled_accumulate_into_tuple(op::O, itr, init) where {O} =
-    (rec_unroll(itr) ? rec_unrolled_accumulate : gen_unrolled_accumulate)(
-        op,
-        itr,
-        init,
-    )
+    _unrolled_accumulate(Val(length(itr)), op, itr, init)
 @inline unrolled_accumulate_into(output_type, op::O, itr, init) where {O} =
     constructor_from_tuple(output_type)(
         unrolled_accumulate_into_tuple(op, itr, init),
@@ -209,13 +211,7 @@ include("generatively_unrolled_functions.jl")
     itr,
     itrs...,
 ) where {F, I, E} =
-    (rec_unroll(itr) ? rec_unrolled_ifelse : gen_unrolled_ifelse)(
-        f,
-        get_if,
-        get_else,
-        itr,
-        itrs...,
-    )
+    _unrolled_ifelse(Val(length(itr)), f, get_if, get_else, itr, itrs...)
 
 ##
 ## Unrolled functions without any analogues in Base

src/generatively_unrolled_functions.jl

@@ -1,54 +1,45 @@
-@generated _gen_unrolled_map(::Val{N}, f, itr) where {N} = quote
+@generated _unrolled_map(::Val{N}, f, itr) where {N} = quote
     @inline
     return Base.Cartesian.@ntuple $N n -> f(generic_getindex(itr, n))
 end
-@inline gen_unrolled_map(f, itr) = _gen_unrolled_map(Val(length(itr)), f, itr)
 
-@generated _gen_unrolled_any(::Val{N}, f, itr) where {N} = quote
+@generated _unrolled_any(::Val{N}, f, itr) where {N} = quote
     @inline
     return Base.Cartesian.@nany $N n -> f(generic_getindex(itr, n))
 end
-@inline gen_unrolled_any(f, itr) = _gen_unrolled_any(Val(length(itr)), f, itr)
 
-@generated _gen_unrolled_all(::Val{N}, f, itr) where {N} = quote
+@generated _unrolled_all(::Val{N}, f, itr) where {N} = quote
     @inline
     return Base.Cartesian.@nall $N n -> f(generic_getindex(itr, n))
 end
-@inline gen_unrolled_all(f, itr) = _gen_unrolled_all(Val(length(itr)), f, itr)
 
-@generated _gen_unrolled_foreach(::Val{N}, f, itr) where {N} = quote
+@generated _unrolled_foreach(::Val{N}, f, itr) where {N} = quote
     @inline
     Base.Cartesian.@nexprs $N n -> f(generic_getindex(itr, n))
     return nothing
 end
-@inline gen_unrolled_foreach(f, itr) =
-    _gen_unrolled_foreach(Val(length(itr)), f, itr)
 
-@generated _gen_unrolled_reduce(::Val{N}, op, itr, init) where {N} = quote
+@generated _unrolled_reduce(::Val{N}, op, itr, init) where {N} = quote
     @inline
-    $N == 0 && return init
-    first_itr_item = generic_getindex(itr, 1)
-    value_1 = init isa NoInit ? first_itr_item : op(init, first_itr_item)
+    # $N == 0 && return empty_reduction_value(init)
+    # The N == 0 case is handled separately in a hard-coded method.
+    value_1 = first_reduction_value(op, itr, init)
     Base.Cartesian.@nexprs $(N - 1) n ->
         (value_{n + 1} = op(value_n, generic_getindex(itr, n + 1)))
     return $(Symbol(:value_, N))
 end
-@inline gen_unrolled_reduce(op, itr, init) =
-    _gen_unrolled_reduce(Val(length(itr)), op, itr, init)
 
-@generated _gen_unrolled_accumulate(::Val{N}, op, itr, init) where {N} = quote
+@generated _unrolled_accumulate(::Val{N}, op, itr, init) where {N} = quote
     @inline
-    $N == 0 && return ()
-    first_itr_item = generic_getindex(itr, 1)
-    value_1 = init isa NoInit ? first_itr_item : op(init, first_itr_item)
+    # $N == 0 && return ()
+    # The N == 0 case is handled separately in a hard-coded method.
+    value_1 = first_reduction_value(op, itr, init)
     Base.Cartesian.@nexprs $(N - 1) n ->
         (value_{n + 1} = op(value_n, generic_getindex(itr, n + 1)))
     return Base.Cartesian.@ntuple $N n -> value_n
 end
-@inline gen_unrolled_accumulate(op, itr, init) =
-    _gen_unrolled_accumulate(Val(length(itr)), op, itr, init)
 
-@generated _gen_unrolled_ifelse(::Val{N}, f, get_if, get_else, itr) where {N} =
+@generated _unrolled_ifelse(::Val{N}, f, get_if, get_else, itr) where {N} =
     quote
         @inline
         Base.Cartesian.@nexprs $N n -> begin
@@ -57,10 +48,8 @@ end
         end
         return get_else()
     end
-@inline gen_unrolled_ifelse(f, get_if, get_else, itr) =
-    _gen_unrolled_ifelse(Val(length(itr)), f, get_if, get_else, itr)
 
-@generated _gen_unrolled_ifelse2(
+@generated _unrolled_ifelse(
     ::Val{N},
     f,
     get_if,
@@ -74,5 +63,3 @@ end
     end
     return get_else()
 end
-@inline gen_unrolled_ifelse(f, get_if, get_else, itr1, itr2) =
-    _gen_unrolled_ifelse2(Val(length(itr1)), f, get_if, get_else, itr1, itr2)

src/hard_coded_unrolled_functions.jl

@@ -0,0 +1,111 @@
+@inline _unrolled_map(::Val{0}, f, itr) = ()
+@inline _unrolled_map(::Val{1}, f, itr) = (f(generic_getindex(itr, 1)),)
+@inline _unrolled_map(::Val{2}, f, itr) =
+    (f(generic_getindex(itr, 1)), f(generic_getindex(itr, 2)))
+@inline _unrolled_map(::Val{3}, f, itr) = (
+    f(generic_getindex(itr, 1)),
+    f(generic_getindex(itr, 2)),
+    f(generic_getindex(itr, 3)),
+)
+
+@inline _unrolled_any(::Val{0}, f, itr) = false
+@inline _unrolled_any(::Val{1}, f, itr) = f(generic_getindex(itr, 1))
+@inline _unrolled_any(::Val{2}, f, itr) =
+    f(generic_getindex(itr, 1)) || f(generic_getindex(itr, 2))
+@inline _unrolled_any(::Val{3}, f, itr) =
+    f(generic_getindex(itr, 1)) ||
+    f(generic_getindex(itr, 2)) ||
+    f(generic_getindex(itr, 3))
+
+@inline _unrolled_all(::Val{0}, f, itr) = true
+@inline _unrolled_all(::Val{1}, f, itr) = f(generic_getindex(itr, 1))
+@inline _unrolled_all(::Val{2}, f, itr) =
+    f(generic_getindex(itr, 1)) && f(generic_getindex(itr, 2))
+@inline _unrolled_all(::Val{3}, f, itr) =
+    f(generic_getindex(itr, 1)) &&
+    f(generic_getindex(itr, 2)) &&
+    f(generic_getindex(itr, 3))
+
+@inline _unrolled_foreach(::Val{0}, f, itr) = nothing
+@inline function _unrolled_foreach(::Val{1}, f, itr)
+    f(generic_getindex(itr, 1))
+    return nothing
+end
+@inline function _unrolled_foreach(::Val{2}, f, itr)
+    f(generic_getindex(itr, 1))
+    f(generic_getindex(itr, 2))
+    return nothing
+end
+@inline function _unrolled_foreach(::Val{3}, f, itr)
+    f(generic_getindex(itr, 1))
+    f(generic_getindex(itr, 2))
+    f(generic_getindex(itr, 3))
+    return nothing
+end
+
+@inline _unrolled_reduce(::Val{0}, op, itr, init) = empty_reduction_value(init)
+@inline _unrolled_reduce(::Val{1}, op, itr, init) =
+    first_reduction_value(op, itr, init)
+@inline _unrolled_reduce(::Val{2}, op, itr, init) =
+    op(first_reduction_value(op, itr, init), generic_getindex(itr, 2))
+@inline _unrolled_reduce(::Val{3}, op, itr, init) = op(
+    op(first_reduction_value(op, itr, init), generic_getindex(itr, 2)),
+    generic_getindex(itr, 3),
+)
+
+@inline _unrolled_accumulate(::Val{0}, op, itr, init) = ()
+@inline function _unrolled_accumulate(::Val{1}, op, itr, init)
+    value_1 = first_reduction_value(op, itr, init)
+    return (value_1,)
+end
+@inline function _unrolled_accumulate(::Val{2}, op, itr, init)
+    value_1 = first_reduction_value(op, itr, init)
+    value_2 = op(value_1, generic_getindex(itr, 2))
+    return (value_1, value_2)
+end
+@inline function _unrolled_accumulate(::Val{3}, op, itr, init)
+    value_1 = first_reduction_value(op, itr, init)
+    value_2 = op(value_1, generic_getindex(itr, 2))
+    value_3 = op(value_2, generic_getindex(itr, 3))
+    return (value_1, value_2, value_3)
+end
+
+@inline _unrolled_ifelse(::Val{0}, f, get_if, get_else, itr) = get_else()
+@inline function _unrolled_ifelse(::Val{1}, f, get_if, get_else, itr)
+    item_1 = generic_getindex(itr, 1)
+    f(item_1) && return get_if(item_1)
+    return get_else()
+end
+@inline function _unrolled_ifelse(::Val{2}, f, get_if, get_else, itr)
+    item_1 = generic_getindex(itr, 1)
+    f(item_1) && return get_if(item_1)
+    item_2 = generic_getindex(itr, 2)
+    f(item_2) && return get_if(item_2)
+    return get_else()
+end
+@inline function _unrolled_ifelse(::Val{3}, f, get_if, get_else, itr)
+    item_1 = generic_getindex(itr, 1)
+    f(item_1) && return get_if(item_1)
+    item_2 = generic_getindex(itr, 2)
+    f(item_2) && return get_if(item_2)
+    item_3 = generic_getindex(itr, 3)
+    f(item_3) && return get_if(item_3)
+    return get_else()
+end
+
+@inline _unrolled_ifelse(::Val{0}, f, get_if, get_else, itr1, itr2) = get_else()
+@inline function _unrolled_ifelse(::Val{1}, f, get_if, get_else, itr1, itr2)
+    f(generic_getindex(itr1, 1)) && return get_if(generic_getindex(itr2, 1))
+    return get_else()
+end
+@inline function _unrolled_ifelse(::Val{2}, f, get_if, get_else, itr1, itr2)
+    f(generic_getindex(itr1, 1)) && return get_if(generic_getindex(itr2, 1))
+    f(generic_getindex(itr1, 2)) && return get_if(generic_getindex(itr2, 2))
+    return get_else()
+end
+@inline function _unrolled_ifelse(::Val{3}, f, get_if, get_else, itr1, itr2)
+    f(generic_getindex(itr1, 1)) && return get_if(generic_getindex(itr2, 1))
+    f(generic_getindex(itr1, 2)) && return get_if(generic_getindex(itr2, 2))
+    f(generic_getindex(itr1, 3)) && return get_if(generic_getindex(itr2, 3))
+    return get_else()
+end

src/unrollable_iterator_interface.jl

@@ -1,14 +1,3 @@
-"""
-    rec_unroll(itr)
-
-Whether to use recursive loop unrolling instead of generative loop unrolling for
-the iterator `itr`. Recursive unrolling can lead to suboptimal LLVM code for
-iterators of more than 32 items, but it is typically faster than generative
-unrolling for short iterators. By default, recursive unrolling is used for
-iterators up to length 2, and generative unrolling is used for longer iterators.
-"""
-@inline rec_unroll(itr) = length(itr) <= 2
-
 """
     generic_getindex(itr, n)
 

src/recursively_unrolled_functions.jl test/recursively_unrolled_functions.jlsrc/recursively_unrolled_functions.jl renamed to test/recursively_unrolled_functions.jl

@@ -1,5 +1,4 @@
-# TODO: Replace all of these with manually unrolled functions, which should be
-# faster to compile. That is the pattern used in Base for ntuple, map, etc.
+using UnrolledUtilities: NoInit, generic_getindex, unrolled_drop
 
 @inline _rec_unrolled_map(f) = ()
 @inline _rec_unrolled_map(f, item, items...) =
@@ -53,11 +52,4 @@
     f(item1) ? get_if(item2) :
     _rec_unrolled_ifelse2(f, get_if, get_else, items...)
 @inline rec_unrolled_ifelse(f, get_if, get_else, itr1, itr2) =
-    _rec_unrolled_ifelse2(f, get_if, get_else, _unrolled_zip(itr1, itr2)...)
-# Using zip here triggers the recursion limit for one unit test on Julia 1.10.
-
-@inline _unrolled_zip(itr1, itr2) =
-    ntuple(Val(min(length(itr1), length(itr2)))) do n
-        @inline
-        (generic_getindex(itr1, n), generic_getindex(itr2, n))
-    end
+    _rec_unrolled_ifelse2(f, get_if, get_else, zip(itr1, itr2)...)