Julia: Function redefinition can improve performance/inlining/inference

[schmrlng]

I don't think this is a particularly new observation and is likely related to #28683, #28940, enduring broadcasting struggles for StaticArrays (https://github.com/JuliaArrays/StaticArrays.jl/issues/482, https://github.com/JuliaArrays/StaticArrays.jl/issues/609), and associated lore about "recursion limiting heuristics" (#29816, #29294), but most of those linked issues are closed now so I figured I'd resurface this old test case that still seems to work in 1.2.0-rc2.0:

  | | |_| | | | (_| |  |  Version 1.2.0-rc2.0 (2019-07-08)
 _/ |\__'_|_|_|\__'_|  |  Official https://julialang.org/ release
|__/                   |

(v1.2) pkg> st
    Status `~/.julia/environments/v1.2/Project.toml`
  [6e4b80f9] BenchmarkTools v0.4.2
  [90137ffa] StaticArrays v0.11.0

julia> using LinearAlgebra, Test, BenchmarkTools, StaticArrays

julia> perp(v) = SVector(v[2], -v[1])
perp (generic function with 1 method)

julia> f(v) = perp.(v) ./ norm.(v)
f (generic function with 1 method)

julia> x = rand(SVector{1,SVector{2,Float64}})
1-element SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1} with indices SOneTo(1):
 [0.5991125473138443, 0.13915128285280431]

julia> @inferred f(x)
ERROR: return type SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1} does not match inferred return type SArray{Tuple{1},_A,1,1} where _A
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] top-level scope at REPL[5]:1

julia> @benchmark f($x)
BenchmarkTools.Trial: 
  memory estimate:  96 bytes
  allocs estimate:  3
  --------------
  minimum time:     25.924 ns (0.00% GC)
  median time:      27.211 ns (0.00% GC)
  mean time:        36.222 ns (19.56% GC)
  maximum time:     36.263 μs (99.83% GC)
  --------------
  samples:          10000
  evals/sample:     993

julia> @code_warntype optimize=true f(x)
Variables
  #self#::Core.Compiler.Const(f, false)
  v::SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}

Body::SArray{Tuple{1},_A,1,1} where _A
1 ─ %1 = Core.tuple(v)::Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}
│   %2 = %new(Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(perp),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}, perp, %1, nothing)::Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(perp),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}
│   %3 = Core.tuple(v)::Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}
│   %4 = %new(Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(norm),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}, LinearAlgebra.norm, %3, nothing)::Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(norm),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}
│   %5 = Core.tuple(%2, %4)::Tuple{Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(perp),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}},Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(norm),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}}
│   %6 = %new(Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Tuple{SOneTo{1}},typeof(/),Tuple{Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(perp),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}},Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(norm),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}}}, /, %5, (SOneTo(1),))::Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Tuple{SOneTo{1}},typeof(/),Tuple{Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(perp),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}},Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(norm),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}}}
│   %7 = invoke Base.Broadcast.copy(%6::Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Tuple{SOneTo{1}},typeof(/),Tuple{Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(perp),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}},Base.Broadcast.Broadcasted{StaticArrays.StaticArrayStyle{1},Nothing,typeof(norm),Tuple{SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}}}}})::SArray{Tuple{1},_A,1,1} where _A
└──      return %7

julia> perp(v) = SVector(v[2], -v[1])
perp (generic function with 1 method)

julia> @inferred f(x)
1-element SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1} with indices SOneTo(1):
 [0.22624014036577766, -0.9740715573751618]

julia> @benchmark f($x)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     2.461 ns (0.00% GC)
  median time:      2.954 ns (0.00% GC)
  mean time:        3.125 ns (0.00% GC)
  maximum time:     17.793 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1000

julia> @code_warntype optimize=true f(x)
Variables
  #self#::Core.Compiler.Const(f, false)
  v::SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}

Body::SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}
1 ── %1  = Base.getfield(v, :data)::Tuple{SArray{Tuple{2},Float64,1,2}}
│    %2  = Base.getfield(%1, 1, false)::SArray{Tuple{2},Float64,1,2}
│    %3  = Base.getfield(v, :data)::Tuple{SArray{Tuple{2},Float64,1,2}}
│    %4  = Base.getfield(%3, 1, false)::SArray{Tuple{2},Float64,1,2}
│    %5  = Base.getfield(%4, :data)::Tuple{Float64,Float64}
│    %6  = Base.getfield(%5, 1, false)::Float64
│    %7  = Base.mul_float(%6, %6)::Float64
│    %8  = Base.getfield(%4, :data)::Tuple{Float64,Float64}
│    %9  = Base.getfield(%8, 2, false)::Float64
│    %10 = Base.mul_float(%9, %9)::Float64
│    %11 = Base.add_float(%7, %10)::Float64
│    %12 = Base.lt_float(%11, 0.0)::Bool
└───       goto #3 if not %12
2 ──       invoke Base.Math.throw_complex_domainerror(:sqrt::Symbol, %11::Float64)
└───       $(Expr(:unreachable))
3 ┄─ %16 = Base.Math.sqrt_llvm(%11)::Float64
└───       goto #4
4 ──       goto #5
5 ──       goto #6
6 ──       goto #7
7 ── %21 = Base.getfield(%2, :data)::Tuple{Float64,Float64}
│    %22 = Base.getfield(%21, 2, true)::Float64
│    %23 = Base.getfield(%2, :data)::Tuple{Float64,Float64}
│    %24 = Base.getfield(%23, 1, true)::Float64
│    %25 = Base.neg_float(%24)::Float64
└───       goto #8
8 ── %27 = Base.div_float(%22, %16)::Float64
│    %28 = Base.div_float(%25, %16)::Float64
│    %29 = StaticArrays.tuple(%27, %28)::Tuple{Float64,Float64}
│    %30 = %new(SArray{Tuple{2},Float64,1,2}, %29)::SArray{Tuple{2},Float64,1,2}
└───       goto #9
9 ── %32 = StaticArrays.tuple(%30)::Tuple{SArray{Tuple{2},Float64,1,2}}
│    %33 = %new(SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}, %32)::SArray{Tuple{1},SArray{Tuple{2},Float64,1,2},1,1}
└───       goto #10
10 ─       goto #11
11 ─       goto #12
12 ─       return %33

[tkf]

FWIW here is a code snippet to trigger a similar issue about inference. ~Like broadcasting, it's heavily relying on recursion.~ (edit: it was broadcasting as well) The inference succeeds in the second try with Julia 1.2.0-rc1.2, 1.2.0-rc2.0, and 1.3.0-DEV.533. In 1.0.3 and 1.1.1, it fails in the second try as well. Here is the record in Travis: https://travis-ci.com/tkf/NDReducibles.jl/builds/118867594

using Pkg
pkg"add https://github.com/tkf/NDReducibles.jl#inference-quirk"
pkgid = Base.PkgId(Base.UUID(0xe7e9fd8f232e412695ea53110d05d1fe), "NDReducibles")
path = Base.locate_package(pkgid)
testdir = joinpath(dirname(dirname(path)), "test")
Pkg.activate(testdir)
Pkg.instantiate()

using NDReducibles: AccessPattern, Index, plan

f1() = plan(AccessPattern.((
    ones(0) => (Index(:i),),
    ones(0, 0) => (Index(:i), Index(:j)),
))...)

using Test
try
    @inferred f1()
catch exception
    @info "Failed as expected" exception
end

@info "Redefining it..."
f2() = plan(AccessPattern.((
    ones(0) => (Index(:i),),
    ones(0, 0) => (Index(:i), Index(:j)),
))...)

@inferred f2()
@info "It worked!"

[tkf]

For my case, I managed to turn it into a smaller example:

g(p) = p[1] => identity.(p[2])
f() = g.((1 => (Val(2), Val(3)), 4 => (Val(5), Val(6))))

Example session:

   _       _ _(_)_     |  Documentation: https://docs.julialang.org
  (_)     | (_) (_)    |
   _ _   _| |_  __ _   |  Type "?" for help, "]?" for Pkg help.
  | | | | | | |/ _` |  |
  | | |_| | | | (_| |  |  Version 1.3.0-DEV.533 (2019-07-10)
 _/ |\__'_|_|_|\__'_|  |  Commit 073cbbb491 (2 days old master)
|__/                   |

julia> using Test

julia> g(p) = p[1] => identity.(p[2])
g (generic function with 1 method)

julia> f() = g.((1 => (Val(2), Val(3)), 4 => (Val(5), Val(6))))
f (generic function with 1 method)

julia> @inferred f()
ERROR: return type Tuple{Pair{Int64,Tuple{Val{2},Val{3}}},Pair{Int64,Tuple{Val{5},Val{6}}}} does not match inferred return type Tuple{Pair{Int64,_A} where _A,Pair{Int64,_A} where _A}
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] top-level scope at REPL[4]:1

julia> f() = g.((1 => (Val(2), Val(3)), 4 => (Val(5), Val(6))))
f (generic function with 1 method)

julia> @inferred f()
(1 => (Val{2}(), Val{3}()), 4 => (Val{5}(), Val{6}()))

[perrutquist]

In the last example, if I run g(0 => (Val(5), Val(6))); g(0 => (Val(2), Val(3))) before defining f then inference succeeds. Conversely, defining f twice without running it in between does not help inference.

[chakravala]

Another example of this issue was discussed on the discourse recently.

Are there any plans for resolving this issue?