Julia: Splitting methods of a function across two modules prevents some optimizations

[mateuszbaran]

I'm experiencing a problem when splitting methods of a function between two modules. I have the following two codes:

using BenchmarkTools
number_eltype(::AbstractArray{T}) where T = T

struct ProductRepr{TM<:Tuple}
    parts::TM
end

ProductRepr(points...) = ProductRepr{typeof(points)}(points)

function number_eltype(x::ProductRepr)
    return typeof(reduce(+, one(number_eltype(eti)) for eti in x.parts))
end

b = ProductRepr([1.0, 2.0], [1.0, 3.0])
@benchmark number_eltype($b)

and

using BenchmarkTools

module A
number_eltype(::AbstractArray{T}) where T = T
end

module B
struct ProductRepr{TM<:Tuple}
    parts::TM
end

ProductRepr(points...) = ProductRepr{typeof(points)}(points)

function Main.A.number_eltype(x::ProductRepr)
    return typeof(reduce(+, one(Main.A.number_eltype(eti)) for eti in x.parts))
end
end

b = B.ProductRepr([1.0, 2.0], [1.0, 3.0])
@benchmark A.number_eltype($b)

As far as I can see the only significant difference is that in the second example number_eltype is split across two modules. Now the fun part is benchmarking.

First, 1.4.1:

julia> versioninfo()
Julia Version 1.4.1
Commit 381693d3df* (2020-04-14 17:20 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: Intel(R) Core(TM) i7-4800MQ CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-8.0.1 (ORCJIT, haswell)

First code:

julia> @benchmark number_eltype($b)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     2.250 ns (0.00% GC)
  median time:      2.262 ns (0.00% GC)
  mean time:        2.261 ns (0.00% GC)
  maximum time:     6.825 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

Second code:

julia> @benchmark A.number_eltype($b)
BenchmarkTools.Trial: 
  memory estimate:  32 bytes
  allocs estimate:  2
  --------------
  minimum time:     551.080 ns (0.00% GC)
  median time:      569.809 ns (0.00% GC)
  mean time:        576.431 ns (0.07% GC)
  maximum time:     4.701 μs (87.24% GC)
  --------------
  samples:          10000
  evals/sample:     188

Julia 1.5 master gives very similar results:

julia> versioninfo()
Julia Version 1.5.0-DEV.814
Commit f1d10e71ab (2020-05-04 11:41 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: Intel(R) Core(TM) i7-4800MQ CPU @ 2.70GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-9.0.1 (ORCJIT, haswell)

[JeffBezanson]

I believe this is due not to where the methods are located, but rather to the Main.A reference, which is not constant. Main isn't a specific package, but just "whatever is in the interactive environment right now". We could possibly change that, but in the meantime it should be fixed if module B does import ..A instead.

[yuyichao]

This is because Main isn't a constant binding. Making it a constant binding gives good result.

using BenchmarkTools

module A
number_eltype(::AbstractArray{T}) where T = T
end

module B
const M = Main

struct ProductRepr{TM<:Tuple}
    parts::TM
end

ProductRepr(points...) = ProductRepr{typeof(points)}(points)

function M.A.number_eltype(x::ProductRepr)
    return typeof(reduce(+, one(M.A.number_eltype(eti)) for eti in x.parts))
end
end

b = B.ProductRepr([1.0, 2.0], [1.0, 3.0])
@benchmark A.number_eltype($b)

Ooops, Jeff commented while I was doing the testing...

[mateuszbaran]

Yes, thanks, I didn't know that. I was hoping it was related to the actual issue I had (and couldn't reduce to anything else). Just in case:

using BenchmarkTools, Manifolds

b = ProductRepr([1.0, 2.0], [1.0, 3.0])
@benchmark number_eltype($b)

actually should execute the same code as posted above but I get

julia> @benchmark number_eltype($b)
BenchmarkTools.Trial: 
  memory estimate:  176 bytes
  allocs estimate:  6
  --------------
  minimum time:     102.802 ns (0.00% GC)
  median time:      105.560 ns (0.00% GC)
  mean time:        114.749 ns (5.96% GC)
  maximum time:     2.200 μs (94.17% GC)
  --------------
  samples:          10000
  evals/sample:     940

on master and

julia> @benchmark number_eltype($b)
BenchmarkTools.Trial: 
  memory estimate:  32 bytes
  allocs estimate:  2
  --------------
  minimum time:     63.204 ns (0.00% GC)
  median time:      63.605 ns (0.00% GC)
  mean time:        66.557 ns (0.75% GC)
  maximum time:     917.006 ns (92.30% GC)
  --------------
  samples:          10000
  evals/sample:     980

on 1.4.1