Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Minor improvement to inference of ntuple #54544

Merged
merged 1 commit into from
May 23, 2024

Conversation

martinholters
Copy link
Member

In the non-@generated branch, where N may be unknown, assert its type as Int so that the type of the created range can be inferred.

julia> Core.Compiler.return_type(ntuple, Tuple{typeof(identity), Val})
Tuple                # master
Tuple{Vararg{Int64}} # PR

Unfortunately, if the function argument could not be inferred concretely, inference is still sub-optimal:

julia> code_typed(Tuple{Any, Int}) do x, n
           ntuple(i -> (x; i), Val(n))
       end
1-element Vector{Any}:
 CodeInfo(
1 ── %1  = Core.typeof(x)::DataType%2  = Core.apply_type(Main.:(var"#26#28"), %1)::Type{var"#26#28"{_A}} where _A
│    %3  = %new(%2, x)::var"#26#28"%4  = Core.apply_type(Base.Val, n)::Type{Val{_A}} where _A
│    %5  = %new(%4)::Val%6  = (isa)(%5, Val{3})::Bool
└───       goto #3 if not %6
2 ──       goto #10
3 ── %9  = (isa)(%5, Val{2})::Bool
└───       goto #5 if not %9
4 ──       goto #10
5 ── %12 = (isa)(%5, Val{1})::Bool
└───       goto #7 if not %12
6 ──       goto #10
7 ── %15 = (isa)(%5, Val{0})::Bool
└───       goto #9 if not %15
8 ──       goto #10
9 ── %18 = Main.ntuple(%3, %5)::Tuple
└───       goto #10
10%20 = φ (#2 => (1, 2, 3), #4 => (1, 2), #6 => (1,), #8 => (), #9 => %18)::Tuple
└───       return %20
) => Tuple

Here, the inferred type of the closure is abstract as the type of the captured variable is unknown, but its return type can still be inferred, as witnessed for the small-N special cases. But for the general-N case, the return type is lost for some reason.

A further improvement to inference precision would be to replace Tuple(f(i) for i = 1:(N::Int)) with ((f(i) for i = 1:(N::Int))...,), but for statically unknown N, that has worse run-time performance characteristics AFAICT. I therefore stayed with the safer, minimally-invasive change here.

In the non-`@generated` branch, where `N` may be unknown, assert its
type as `Int` so that the type of the created range can be inferred.
@Seelengrab
Copy link
Contributor

I'm surprised this is necessary, since there's the N::Int assertion right above too, theoretically covering both @generated and non-@generated 🤔

@martinholters
Copy link
Member Author

I also thought we could propagate these asserts downstream, but we can't, be it for ordinary arguments or type parameters:

julia> function f(m, ::Val{N}) where N
           m::Int
           N::Int
           (m, N)
       end
f (generic function with 1 method)

julia> code_typed(f, Tuple{Any, Val})
1-element Vector{Any}:
 CodeInfo(
1 ─      Core.typeassert(m, Main.Int)::Int64%2 = $(Expr(:static_parameter, 1))::Any
│        Core.typeassert(%2, Main.Int)::Int64%4 = $(Expr(:static_parameter, 1))::Any%5 = Core.tuple(m, %4)::Tuple{Any, Any}
└──      return %5
) => Tuple{Any, Any}

@Seelengrab
Copy link
Contributor

Hmmm that's really weird

f(m, ::Val{N}) where N @ Main REPL[4]:1
Variables
  #self#::Core.Const(f)
  m::Any
  _::Val

%0 = invoke f(::Any,::Val)::Tuple{Any, Any}
    @ REPL[4]:2 within `unknown scope`
1%1 = Main.Int::Core.Const(Int64)
│        Core.typeassert(m, %1)::Int64
│   @ REPL[4]:3 within `unknown scope`%3 = $(Expr(:static_parameter, 1))::Any%4 = Main.Int::Core.Const(Int64)
│        Core.typeassert(%3, %4)::Int64
│   @ REPL[4]:4 within `unknown scope`%6 = $(Expr(:static_parameter, 1))::Any%7 = Core.tuple(m, %6)::Tuple{Any, Any}
└──      return %7

According to Cthulhu, the type is asserted on the original SSA value, but the values for the tuple are then taken from the unasserted SSA value. That sounds like a more general bug to me 🤔 No idea what's going on with m, which is reused but the assert is apparently unused?

@martinholters
Copy link
Member Author

The usual solution is to do m = m::Int instead of just m::Int, but that doesn't work with static parameters (like here) as they cannot be assigned to.

@KristofferC
Copy link
Member

Ref #38274 for the N::Int not "taking".

@martinholters
Copy link
Member Author

Thanks for digging that up. There is no reason to believe that will be fixed soon, so I don't see any reason we shouldn't merge this: For known N , it will be optimized out. Otherwise, it will safe a dynamic dispatch (and improve inference precision) at the cost of a type-check.

@martinholters martinholters merged commit 95cdf9b into master May 23, 2024
7 checks passed
@martinholters martinholters deleted the mh/improve-ntuple-inference branch May 23, 2024 11:51
@nsajko nsajko added the compiler:inference Type inference label Jul 5, 2024
lazarusA pushed a commit to lazarusA/julia that referenced this pull request Jul 12, 2024
In the non-`@generated` branch, where `N` may be unknown, assert its
type as `Int` so that the type of the created range can be inferred.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
compiler:inference Type inference
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants