`haskey(::@Kwargs{item::Bool}, :item)` does not constant-fold

[oscardssmith]

The following (fairly simplified example) loses

using OrdinaryDiffEqRosenbrock, OrdinaryDiffEqCore, ForwardDiff, ADTypes

code_typed(OrdinaryDiffEqCore.:(var"#remake#8"), (Base.Pairs{Symbol, ADTypes.AutoForwardDiff{1, ForwardDiff.Tag{Nothing, Float64}}, Tuple{Symbol}, @NamedTuple{autodiff::ADTypes.AutoForwardDiff{1, ForwardDiff.Tag{Nothing, Float64}}}}, typeof(SciMLBase.remake), OrdinaryDiffEqRosenbrock.Rodas5P{1, Nothing, Nothing, Nothing, Nothing, true, nothing, Nothing, Nothing}))

returns a poorly inferred Rodas5P because the OrdinaryDiffEqCore.var"#remake#8 (in src/algirithms.jl:54) is getting called with badly inferred kwargs even though we specifically are calling it with a very well inferred namedTuple.

@serenity4 I was told by @topolarity that you might be interested in this.

[KristofferC]

It would be helpful to extract this into a stand alone MWE.

[topolarity]

# setup
julia> solver = OrdinaryDiffEqRosenbrock.Rodas5P{1, Nothing, Nothing, Nothing, Nothing, true, nothing, Nothing, Nothing}(nothing, nothing, nothing, nothing, nothing);
julia> autodiff = (ADTypes.AutoForwardDiff{1, ForwardDiff.Tag{Nothing, Float64}}(ForwardDiff.Tag{Nothing, Float64}()));

One thing I don't understand is why this code_typed is good:

julia> @code_typed optimize=false OrdinaryDiffEqCore.remake(solver; autodiff)
CodeInfo(
...
) => Rodas5P{1, AutoForwardDiff{1, ForwardDiff.Tag{Nothing, Float64}}, Nothing, Nothing, Val{:forward}(), true, nothing, Nothing, Nothing}

but a manual code_typed on the internal call is bad:

julia> @code_typed OrdinaryDiffEqCore.var"#remake#6"(Base.pairs((; autodiff)), OrdinaryDiffEqCore.remake, solver)
CodeInfo(
...
) => Rodas5P

As a side note, some solvers end up w/ better effects (e.g. solver = OrdinaryDiffEqRosenbrock.Rodas5P()), which causes everything to be concrete-eval'd. The divergence here seems to be in the case that this function doesn't have good effects, though.

[topolarity]

Possible MWE:

julia> remake(; kwargs...) = haskey(kwargs, :autodiff) ? Val{1}() : Val{2}()

# outer function infers well
julia> @code_typed remake(; autodiff=true)
CodeInfo(
1 ─     return Val{1}()
) => Val{1}

# inner function does not?
julia> @code_typed var"#remake#1"(Base.pairs((; autodiff=true)), remake)
CodeInfo(
1 ─ %1 =   builtin Base.Iterators.getfield(kwargs, :itr)::Tuple{Symbol}
│   %2 =    invoke Base.sym_in(:autodiff::Symbol, %1::Tuple{Symbol})::Bool
└──      goto #3 if not %2
2 ─      return Val{1}()
3 ─      return Val{2}()
) => Union{Val{1}, Val{2}}

The concrete issue that @oscardssmith ran into is that we can end up invoke-ing the inner function, and it will infer much worse than if we had invoke'd the outer one.

[topolarity]

A major part of the issue is that Base.Pairs{Symbol, Bool, Tuple{Symbol}, @NamedTuple{autodiff::Bool}} is not sufficient type information to resolve getfield(pairs, :itr).

This isn't a problem for the wrapper function because inference sees the Base.Pairs get constructed.

The Base.pairs(kwargs) call in the outer function lets inference know the field values: Core.PartialStruct(Base.Pairs{...}, ..., Any[@NamedTuple{autodiff::Bool}, Core.Const((:autodiff,))]), enabling it to resolve the getfield(kwargs, :itr) and eventually the Base.sym_in call.

[topolarity]

It would be nice to define something like:

keys(v::Pairs) = getfield(v, :data) isa NamedTuple ? keys(getfield(v, :data)) : getfield(v, :itr)
haskey(v::Pairs, key) = key in keys(v)

which does resolve the issue. But this is problematic, since it changes the behavior of Base.Pairs when itr does not match the keys of the NamedTuple.

@serenity4 Any ideas how we might improve this?

[serenity4]

Thanks for getting that MWE @topolarity, that's very helpful.

It seems to me that the only non-problematic way to address this issue is to get field values into the type parameter I, something like Val{(:autodiff,)} instead of Tuple{Symbol}. On the other hand, that doesn't seem really workable in the presence of unknown parameters, because we then introduce a type instability, which is way worse (although I'm not sure how often we'd encounter that situation).

But why call these inner lowering-produced functions, over the outer one?

[topolarity]

But why call these inner lowering-produced functions, over the outer one?

The trouble is that inference can end up making the same choice to call an inner function, often because the inlining heuristics prevent it from being inlined into the outer one.

The concrete example that @oscardssmith hit was that he had a --trim failure due to haskey(kwargs, :autodiff) not const-folding

[serenity4]

The trouble is that inference can end up making the same choice to call an inner function, often because the inlining heuristics prevent it from being inlined into the outer one.

Isn't that an inference issue, rather than an inlining issue? Or is it that if it doesn't get inlined, we can't perform constprop inside of the inner function (because we reuse the generically-inferred code?

In any case, at the end of the day, --trim is the authority here, so that's something to fix, I agree.

[topolarity]

Yeah, --trim is especially an issue because the verifier will inspect the inner function call w/o any const-prop information (it explores the call as if it saw the full invoke worth of arguments)

If it weren't for that, semi-concrete eval (a.k.a. inter-procedural const-prop) would be able to resolve this, if inference decides to run it, although I think in the OP case it was not IIRC. Tuning those heuristics can be a bit fragile.

[topolarity]

But this is problematic, since it changes the behavior of Base.Pairs when itr does not match the keys of the NamedTuple.

Oops! Looks like we already do this in Base.merge:

julia/base/namedtuple.jl

Line 346 in 07716cb

merge(a::NamedTuple, b::Iterators.Pairs{<:Any,<:Any,<:Any,<:NamedTuple}) = merge(a, getfield(b, :data))