Hi, I want to perform dispatch based on functions of the same name within a module. There’s two ways to go about it: having a function f(m::Module,...) or defining a trait type and saying f(::Type{Mod{M}}, ...) where M.
module A
foo(bar) = bar + bar
end
module B
foo(bar) = bar * bar
end
foo_twice(M::Module, bar) = M.foo(M.foo(bar))
or
struct Mod{M} end
foo_twice(::Type{Mod{M}}, bar) where M = M.foo(M.foo(bar))
I found that the second is much much faster than the first. But I’m wondering why. I know that modules could be considered mutable in a way, and that would mean that M.foo falls under mutable global lookup somehow? Still, I can’t shake the feeling that these implementations should be equivalent. Am I missing something?
Cheers, Johan
Welcome!
I have never attempted something like this before, and there’s almost certainly a better way to do what you’re doing.
But the issue is probably that Module is a type (although I didn’t even know that until now). When compiling this function, Julia does not specialize on the instance of the Module but rather on the type Module itself. So the exact same compiled code runs whether your module is A, B, Base, LinearAlgebra, etc.
Because the function does not specialize on the specific module you pass, M.foo must be a run-time lookup. This runs much slower than doing this work at compile time.
Your pattern of Type{Mod{M}} forces the instance of the module into the type domain, in which case it is specialized. Although I would have simply used ::Val{M} rather than Type{Mod{M}} to do that.
Perhaps you can provide a little more info on what you’re trying to accomplish? People here can probably suggest a better way to do it. My first impulse is that you should have a single function with multiple methods or pass the function directly, rather than pass modules and do lookups within them.
1 Like
Thanks, this is more or less what I suspected. I’m working on a larger post on exactly why I’m doing it this way. Short answer: it’s not one function I’m worried about, but a group of related structs and functions that I don’t all want to give trait arguments. Grouping them into a module makes my code more readable.
The funny thing is, I can define
@inline foo_twice(M::Module, bar) = foo_twice(Mod{M}, bar)
regaining a simpler interface with the performance of the trait implementation (as far as I can tell).
Regarding your @inline workaround. This will work sometimes, because inlining makes constant propagation more potent. But it’s only a hint and may not work in all cases. If the module is not a known constant at the call site, you can still end up in a similar situation to the original (except with the dynamic dispatch happening within the ::Module method rather than the ::Type{Mod} method). The function barrier this introduces may result in different performance characteristics (often an improvement).
A more idiomatic pattern similar to what you’re doing might be this (I use slightly different patterns for A and B, you can choose which you prefer):
module ABCore
export foo
function foo end # function stub we'll add to later
end
module A
using ..ABCore # may need to fiddle with the number of `.` depending on your context
export AA
struct AAType end # singleton type
AA = AAType() # construct an instance
ABCore.foo(::AAType, bar) = bar + bar # add a method
end
module B
using ..ABCore # may need to fiddle with the number of `.` depending on your context
export BB
struct BB end # singleton type
# here I use a slightly different pattern where BB is a type, which will change how this is called
ABCore.foo(::BB, bar) = bar * bar # add a method
end
julia> using .ABCore, .A, .B
julia> foo(AA, 3) # AA is an instance, so call like this
6
julia> foo(BB(), 3) # BB is a type, so use () to construct an instance
9
julia> methods(foo) # all methods of the same function now
# 2 methods for generic function "foo" from Main.ABCore:
[1] foo(::BB, bar)
@ Main.B REPL[3]:9
[2] foo(::Main.A.AAType, bar)
@ Main.A REPL[2]:10
In this pattern, rather than passing around Modules, you use the singleton type provided by each module to control which method of foo is called. There is only a single foo function, but different modules create different methods to implement it.