5

When type inference falters (::Any in @code_warntype printout), my understanding is that function calls are dynamically dispatched. In other words, at run-time, the arguments' types are checked to find the specialization (MethodInstance) for the concrete argument types. Needing to do this at run-time instead of compile-time incurs performance costs.

(EDIT: originally, I said "multiple dispatch finds the fitting method" between the type-checking and specialization-finding, but I don't actually know if this part happens at runtime. It seems that it only needs to happen if no valid specialization exists and one needs to be compiled.)

In cases where only one argument's concrete type needs to be checked, is it possible to do a faster dynamic single dispatch instead, like in some sort of lookup table of specializations? I just can't find a way to access and call MethodInstances as if they were functions.

When it comes to altering dispatch or specialization, I thought of invoke and @nospecialize. invoke looks like it might skip right to a specified method, but checking multiple argument types and specialization must still happen. @nospecialize doesn't skip any part of the dispatch process, just results in different specializations.

EDIT: A minimal example with comments that hopefully describe what I'm talking about.

struct Foo end
struct Bar end

#   want to dispatch only on 1st argument
#          still want to specialize on 2nd argument
baz(::Foo, ::Integer) = 1
baz(::Foo, ::AbstractFloat) = 1.0
baz(::Bar, ::Integer) = 1im
baz(::Bar, ::AbstractFloat) = 1.0im

x = Any[Foo(), Bar(), Foo()]

# run test1(x, 1) or test1(x, 1.0)
function test1(x, second)
  #   first::Any in @code_warntype printout
  for first in x
    # first::Any requires dynamic dispatch of baz
    println(baz(first, second))
    # Is it possible to only dispatch -baz- on -first- given
    # the concrete types of the other arguments -second-?
  end
end
BatWannaBe
  • 4,330
  • 1
  • 14
  • 23
  • 4
    one way is to put the rest arguments into keyword arguments as those don't participate in MD, but maybe an example would be helpful – jling Jan 02 '22 at 21:02
  • I'll write an example to run but I won't be able to demonstrate the inner workings so much. – BatWannaBe Jan 03 '22 at 03:51
  • 1
    Probably not what you're looking for but you could use explicit if-branches to "dispatch" based on the first argument as a workaround. This will be fast. – carstenbauer Jan 03 '22 at 08:30
  • 3
    You might, perhaps, want to take a look at https://github.com/jlapeyre/ManualDispatch.jl and (maybe) https://github.com/YingboMa/Unityper.jl. – carstenbauer Jan 03 '22 at 08:37
  • @carstenbauer 1) Yes, a redundant if-elseif checking the concrete type of `first` in my example caused static dispatch in the branches. Have not checked if this scales to more types or hits something like a Union-splitting limit. 2) ManualDispatch.jl seems to simplify the if-elseif check and must check all types. Checking multiple types in a linear if-elseif doesn't seem to scale optimally, but the branchwise static dispatch could be worth it. 3) Unityper.jl is cool, but it combines types into 1, so it seems I can't write a multimethod, just different functions for each branch. – BatWannaBe Jan 03 '22 at 14:16
  • Static dispatch in a branch is equivalent to specifying a specialization, so ManualDispatch.jl is close to what I'm going for, and it has an additional benefit of improved type inference in the branch. I just don't know how well checking types sequentially and repetitively would scale, and I have to explicitly specify what types to check like it was a lookup table, whereas multiple dispatch works no matter how many methods are added. – BatWannaBe Jan 03 '22 at 19:50

1 Answers1

0

The easiest way to do what you ask is to simply not dispatch on the second argument (by not specifying a type assertion on the second variable specific enough to trigger dispatch), and instead specialize with an if statement within your function. For example:

struct Foo end
struct Bar end

# Note lack of type assertion on second variable. 
# We could also write `baz(::Foo, n::Number)` for same effect in this case, 
# but type annotations have no performance benefit in Julia if you're not 
# dispatching on them anyways.
function baz(::Foo, n) 
    if isa(n, Integer)
        1
    elseif isa(n, AbstractFloat)
        1.0
    else
        error("unsupported type")
    end
end

function baz(::Bar, n)
    if isa(n, Integer)
        1im
    elseif isa(n, AbstractFloat)
        1.0im
    else
        error("unsupported type")
    end
end

Now, this will do what you want

julia> x = Any[Foo(), Bar(), Foo()]
3-element Vector{Any}:
 Foo()
 Bar()
 Foo()

julia> test1(x, 1)
1
0 + 1im
1

julia> test1(x, 1.0)
1.0
0.0 + 1.0im
1.0

and since this effectively manually picks only two cases to specialize out of all the possible types to specialize on, I could imagine scenarios where this sort of technique has performance benefits (though, of course, it goes without saying in Julia that generally even better would be to find and eliminate the source of the type instability in the first place if at all possible).

However, it is critically important in the context of this question as written to point out that that even though we have eliminated dispatch on the second argument of the function, these baz functions may still have poor performance if the first argument (i.e., the one you are dispatching on) is type-unstable – as is the case in the question as written because of the use of an Array{Any}.

Instead, try to use an array with at least some type constraint. Ex:

julia> function test2(x, second)
           s = 1+1im
           for first in x
               s += baz(first, second)
           end
           s
       end
test2 (generic function with 1 method)

julia> using BenchmarkTools

julia> x = Any[Foo(), Bar(), Foo()];

julia> @benchmark test2($x, 1)
BenchmarkTools.Trial: 10000 samples with 998 evaluations.
 Range (min … max):  13.845 ns … 71.554 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     13.869 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   15.397 ns ±  3.821 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  █▅  ▃ ▄  ▄      ▄       ▄                                 ▃ ▁
  ██▇▆█▇██▄█▇▇▄▃▁▁██▁▃▃▁▁▃██▃▁▃▁▁▄▃▃▃▆▆▅▆▆▅▅▄▁▁▄▃▃▃▁▃▁▄▁▁▃▄▄█ █
  13.8 ns      Histogram: log(frequency) by time      30.2 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

julia> x = Union{Foo,Bar}[Foo(), Bar(), Foo()];

julia> @benchmark test2($x, 1)
BenchmarkTools.Trial: 10000 samples with 1000 evaluations.
 Range (min … max):  4.654 ns … 62.311 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     4.707 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   5.471 ns ±  1.714 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  █▂▂▃▄ ▃  ▄▁    ▄▂      ▅▁                               ▁▄ ▁
  ███████▁▁██▁▁▁▁██▁▁▁▁▁▁██▁▁▁▄▁▃▁▃▁▁▁▁▃▁▁▁▁▃▁▃▃▁▁▁▁▃▁▁▁▁▁██ █
  4.65 ns      Histogram: log(frequency) by time     10.2 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.
cbk
  • 4,225
  • 6
  • 27
  • 1) Combining methods wrt the 2nd argument reduces the number of methods, but not the number of specializations (`baz(::Foo, n)` still compiles `baz(::Foo,::Int)` and `baz(::Foo,::Float64)`). I don't know how to profile what part of dispatch is compiled, but in your approach, all concrete types must still be checked at run time to find the right specialization. `@nospecialize` would reduce specializations, but I do want the 2nd argument's type to be inferred within the method. 2) The type constraint uses (small-)Union-splitting optimization, which is useful but I was imagining many more types. – BatWannaBe Jan 07 '22 at 17:59
  • If you neither want to dispatch on the second argument _nor_ have the compiler specialize on the second argument, then just add a `@nospecialize` to the functions above. The way you worded your question though, it sounded like you wanted to avoid dispatch on the second argument, not specialization. – cbk Jan 07 '22 at 18:10
  • If you are asking for something very different than what I answered, it is doubtful to me that will have performance benefits, but you may need to word your question more clearly. – cbk Jan 07 '22 at 18:16
  • The wording is accurate, I want to avoid runtime dispatch on the second argument but still specialize on it. Perhaps if I word it this way: if I fix the 2nd argument type as `Int`, I could theoretically check only the 1st argument type to find the proper specialization. What I read is that specializations are always found by checking all the argument types, as if a Tuple of them is a key in some hash table; this implementation can't be altered by rewriting methods. Not sure because I can't track down thorough sources on the implementation-level details of dispatch at compile- and run-time. – BatWannaBe Jan 07 '22 at 18:44
  • One of the comments on my post referenced `ManualDispatch.jl`, and I've seen similar approaches in other threads. It doesn't alter how specializations are found, but type-checks of the poorly inferred arguments in if-else branches can allow specializations to be determined at compile-time instead. Of course, the downside is that if-else branches aren't as extendable as tables. – BatWannaBe Jan 07 '22 at 18:49
  • `I want to avoid runtime dispatch on the second argument but still specialize on it.` This is exactly what the answer above does. Perhaps we are using different meanings of "specialize". However, there still seems to be a bit of a misunderstanding inherent in the way you asked the question, in that there is no performance benefit to doing this if the second argument is not the type-unstable one; for the example vector you gave you might rather dispatch normally for the second argument and avoid dispatch (instead manually branching) on the first argument. – cbk Jan 07 '22 at 19:02
  • 1
    Let us [continue this discussion in chat](https://chat.stackoverflow.com/rooms/240843/discussion-between-batwannabe-and-cbk). – BatWannaBe Jan 07 '22 at 20:59