Type stability when arguments are functions

Sorry if this has been asked before. I have functions that take vectors of functions as arguments, similarly to the following:

function g2(fs::Vector{Function}, x::Float64)::Float64
       a::Float64 = 0.0
       for f in fs
       a += f(x)
       end
       a
end

Let’s say we define

julia> f1(x::Float64)::Float64 = sin(x);
julia> f2(x::Float64)::Float64 = cos(x);
julia> f3(x::Float64)::Float64 = tan(x);

Do the Any in the following @code_warntype indicate type instability? If so, is there a way to make g2 type stable?

julia> @code_warntype g2([f1, f2, f3], 3.14)
Variables:
  #self#::#g2
  fs::Array{Function,1}
  x::Float64
  a::Float64
  #temp#::Int64
  f::F

Body:
  begin 
      SSAValue(0) = Main.Float64
      SSAValue(1) = 0.0
      a::Float64 = SSAValue(1) # line 3:
      #temp#::Int64 = $(QuoteNode(1))
      6: 
      unless (Base.box)(Base.Bool,(Base.not_int)((#temp#::Int64 === (Base.box)(Int64,(Base.add_int)((Base.arraylen)(fs::Array{Function,1})::Int64,1)))::Bool)) goto 17
      SSAValue(5) = (Base.arrayref)(fs::Array{Function,1},#temp#::Int64)::F
      SSAValue(6) = (Base.box)(Int64,(Base.add_int)(#temp#::Int64,1))
      f::F = SSAValue(5)
      #temp#::Int64 = SSAValue(6) # line 4:
      SSAValue(4) = (a::Float64 + (f::F)(x::Float64)::Any)::Any
      a::Float64 = (Core.typeassert)((Base.convert)(Main.Float64,SSAValue(4))::Any,Main.Float64)::Float64
      15: 
      goto 6
      17:  # line 6:
      return a::Float64
  end::Float64

Is there a way to indicate that the input argument fs should be a vector of Functions each with a certain return type?

Thanks!

FunctionWrappers.jl could perhaps be of interest. Note that your function is type stable (because of the type assert) but the result of f is not.

4 Likes

You could replace fs::Vector{Function} with fs::Tuple{Vararg{Function} and pass a tuple of functions rather than an array. That will allow the compiler to know the exact types of the functions. Unfortunately, the return types do not seem to be inferred in that case: the fact that there is a loop over the functions does not appear to be handled (yet).

Try “lispy tuple recursion,”

@inline g2(fs::Tuple{Function, Vararg{Function}}, x::Float64) = _g2(0.0, fs, x)
# _gs is a "private" method that processes the first call, then discards that function and recursively calls itself 
@inline _g2(a, fs::Tuple{Function, Vararg{Function}}, x::Float64) = _g2(a + fs[1](x), Base.tail(fs), x)
# But we have to terminate the recursion. This method is called when we've "used up" all the functions
_g2(a, ::Tuple{}, x::Float64) = a

@code_warntype g2((sin, cos, tan), 3.14)

You’ll see this is type-stable even without any of the type-assertions/declarations.

6 Likes

How does compile time scale with the number of functions and is there a threshold where this does not work? Would try it myself but I don’t have access to a computer for a while.

O(N^2)

is there a threshold where this does not work

Set by MAX_TUPLETYPE_LEN. Currently 15, but Jameson has shown that recent improvements to inference.jl may allow that to be changed to >1000.

3 Likes

This seems like a pretty common pattern in high performance julia. I wonder if there’s a way to get it into a macro.

I’ve got an unregistered package Unrolled.jl with generated functions that achieve the same thing, but without the quadratic compile-time, and for arbitrary-long tuples:

julia> using Unrolled, BenchmarkTools

julia> g2(fs::Tuple, x::Float64) =
           unrolled_reduce(+, 0.0, unrolled_map(f->f(x), fs))
g2 (generic function with 1 method)

julia> @benchmark g2((sin, cos, tan), 10.0)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     58.724 ns (0.00% GC)
  median time:      58.822 ns (0.00% GC)
  mean time:        63.931 ns (0.00% GC)
  maximum time:     280.161 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     984

julia> g3(x) = sin(x) + cos(x) + tan(x)
g3 (generic function with 1 method)

julia> @benchmark g3(10.0)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     54.456 ns (0.00% GC)
  median time:      54.538 ns (0.00% GC)
  mean time:        62.114 ns (0.00% GC)
  maximum time:     269.921 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     985

The definitions are straight-forward, as far as these things go:

@generated function unrolled_map(f, seq::Tuple) 
    :(tuple($((:(f(seq[$i])) for i in 1:type_length(seq))...)))
end

@generated function unrolled_reduce(f, v0, seq) 
    niter = type_length(seq)
    expand(i) = i == 0 ? :v0 : :(f(seq[$i], $(expand(i-1))))
    return expand(niter)
end
6 Likes