Type stability with variable arguments

General Usage Performance
12

That’s not how type inference works in Julia. The key thing is what happens at compile-time, not at parse time (the parser only knows how things are spelled).

It’s totally fine to declare lexicographic2(a,b,c) with no types for the arguments — the compiler knows where the function is called, and (assuming the call site is type-stable/grounded/inferable) it compile a specialized version of lexicographic2 (and hence a specialized version of the returned function lt) for those argument types.

The problem with lexicographic2 is that the local variable v is type-unstable — it takes on 3 different types because (a,b,c) has 3 different types. If instead you write out the loop:

       function lexicographic3(a,b,c )
           function lt( x::Int, y::Int )
               if a[x] < a[y]
                   return true
               elseif a[x] > a[y]
                   return false
               elseif b[x] < b[y]
                   return true
               elseif b[x] > b[y]
                   return false
               elseif c[x] < c[y]
                   return true
               elseif c[x] > c[y]
                   return false
               else
                   return false
               end
           end
           return lt
       end

then the type-instability of v goes away and the performance is good:

julia> @btime sort!( $indices, lt=$(lexicographic3( x, y, z )));
  48.989 ms (0 allocations: 0 bytes)

(Note the lack of allocations — if you have any allocations here, that is indicative of a type-instability.)

The trick is to get the compiler to unroll the loop for you, for any number of arguments. I usually use map or ntuple for this, since it is specialized for tuples, and it is fine to construct a tuple result that is discarded (the compiler will eliminate the tuple allocation). I’d rather use foreach, but that is currently not specialized for tuples (julia#31901):

function lexicographic4(args...) # no type declarations needed here: the compiler will figure it out
    function lt(x, y) # similarly, no type-declarations needed here either
        map(args) do v  # the compiler will completely unroll a map of a tuple
            if v[x] < v[y]
                return true
            elseif v[x] > v[y]
                return false
            end
        end # map result (a tuple) is discarded … compiler will elide it
        return false
    end
    return lt
end

which gives:

julia> @btime sort!( $indices, lt=$(lexicographic4( x, y, z )));
  44.779 ms (0 allocations: 0 bytes)

as desired.

6 Likes