Advice for Higher Order Functions on Tuples

I have a package where I’m operating on lists of types, and I need everything to compile away. I have succeeded in doing so by using these two functions:

using Tricks

"""
Utility function which generates a list of expressions in the form 
`:($f(tuples[1][i], tuples[2][i]...))` for each index of the input tuples. 
Tuples must be the same length.
"""
function make_f_calls(f, tuples)
    tuple_lengths = fieldcount(tuples[1])
    @assert all(fieldcount(t) == tuple_lengths for t in tuples) "all tuples must be same length"
    f_calls = [:(f($(
                   (:(tuples[$j][$i]) for j in eachindex(tuples))...)
               )
               )
               for i in 1:tuple_lengths]
    return f_calls
end

@generated function tuple_all(f, tuples...)
    f_calls = make_f_calls(f, tuples)
    with_returns = [:($f_call || return false) for f_call in f_calls]
    quote
        $(with_returns...)
        return true
    end
end

# then in my code
function is_all_numbers(::Type{Ts}) where Ts
    types = fieldtypes(Ts)
    f = T -> T <: Number
    return tuple_all(f, types)
end

Questions:

  1. Is this silly and there is a simpler way to let the compiler constant propagate stuff like this?
  2. If not, is there already a package that does this kind of unrolling I should be using?
  3. If not 2, is there an existing package where it would be helpful to have this functionality added? I have one for tuple_all, tuple_map, and tuple_any. To me, this solution is in the weird space where it is too simple to be a package but complicated enough that it doesn’t make sense to have it reimplemented a bunch of places.

For this specific case, this compiles away on 1.10. Whether this can fully compile away may depend on the complexity of f

julia> function my_all_num(::Type{TS}) where {TS}
           tup = fieldtypes(TS)
           f = T -> T <: Number
           isnums = f.(tup)
           all(isnums)
       end
my_all_num (generic function with 1 method)

julia> my_all_num(T_tup)
true

julia> @code_llvm my_all_num(T_tup)
;  @ REPL[30]:1 within `my_all_num`
define i8 @julia_my_all_num_528() #0 {
top:
  ret i8 1
}
1 Like

to me, it seems like if Base.all (or any or map) doesn’t infer well enough for tuples, improving the implementation there would be the place to do it

1 Like

Well, I think the issue there is that, for my case, I know that not only are my Tuples type stable, their content is static too. I’m unrolling any length tuple, so I pay tons of extra compilation (worth it in my case). Unrolling a tuple where the content isn’t stable probably isn’t going to balance out to worth it.

For example, it would be a little silly to do tuple_all(x->x>3, some_tuple_of_ints), since we don’t know the value of the Int at compile time.

It’s hard for me to think of a heuristic for a Base implementation where you won’t harm the complication performance of most cases.

edit: some code to underline the point:

julia> function test_tuple_all(t::Type{T}) where T
           ts = static_fieldtypes(t)
           return tuple_all(ts) do T
               T isa Number
           end
       end
test_tuple_all (generic function with 2 methods)

julia> function test_tuple_all2(t::Type{T}) where T
           ts = static_fieldtypes(t)
           return all(ts) do T
               T isa Number
           end
       end
test_tuple_all2 (generic function with 2 methods)

julia> function test_tuple_all3(t)
           return tuple_all(t) do x
               x < 5
           end
       end
test_tuple_all3 (generic function with 1 method)

julia> function test_tuple_all4(t)
           return all(t) do x
               x < 5
           end
       end
test_tuple_all4 (generic function with 1 method)

# Tuple all outperforms base all when we are working on types
julia> @b Tuple{Int, Float64} test_tuple_all(_)
130.337 ns (0.04 allocs: 1.258 bytes)

julia> @b Tuple{Int, Float64} test_tuple_all2(_)
1.006 μs (4.24 allocs: 134.588 bytes)

# Ties Base, but with worse compilation behaviors in the normal case
julia> @b (1, 2) test_tuple_all3(_)
2.065 ns

julia> @b (1, 2) test_tuple_all4(_)
2.058 ns

I don’t understand why are you doing all this. I think a completely naive implementation of is_all_numbers should work fine (as long as length(types) is below the damned hardcoded cutoff chosen in Base). Can you give an example input where a non-naive implementation would be necessary?

What’s this?

You’re not using T in the body?

Probably you want <: instead of isa?

AFAIK this kind of invocation prevents type inference of the argument, i.e., Tuple{Int, Float64} is statically known to be merely DataType. I don’t think that’s your goal when benchmarking?

1 Like

This is likely used to force Julia to specialize the method on the type (see Be aware when Julia avoids specializing). I didn’t test this specific instance so I am unsure whether it is necessary here or not.

To answer this: Imo, if your application absolutely crucially depends on something being compiled away completely or a certain way, then using a @generated function is totally justified. It also protects you if at some in the future compiler development changes some rules that could break your code.
Generally, I am with @nsajko though and would start with simple implementation since Julia can be quite clever with compiling away stuff. I’d recommend keeping this in the back of your mind and focus on the functionality first. If it becomes apparent that you need to control the compilation more closely, you can always go back and optimize that with using @generated functions.

There is also a Base.@constprop :aggressive to encourage the compiler to do constant propagation which might help if you decide to forgo the @generated approach.

2 Likes