-
-
Notifications
You must be signed in to change notification settings - Fork 5.7k
Closed
Description
Currently Base.@locals returns a Dict, but this is unfortunate because it's slow, allocating, and needs to box all of its contents. If Base.@locals carried no runtime cost to use, then it could be used to solve #34168.
It'd be really great if @locals (or if we don't want to change it a new macro @staticlocals) could be a NamedTuple which we should be able to construct in a type stable manner just from syntax alone.
To be concrete, if I write
function f(x)
local y
Base.@locals
endthis lowers to
CodeInfo(
1 ─ Core.NewvarNode(:(y))
│ %2 = Core.apply_type(Base.Dict, Core.Symbol, Core.Any)
│ %3 = (%2)()
│ %4 = $(Expr(:isdefined, :(y)))
└── goto #3 if not %4
2 ─ Base.setindex!(%3, y, :y)
3 ┄ %7 = $(Expr(:isdefined, :(x)))
└── goto #5 if not %7
4 ─ Base.setindex!(%3, x, :x)
5 ┄ return %3
)But I think that we could emit code that looks more like this:
function Base.setindex(nt::NamedTuple{names}, x, ::Val{n}) where {names, n}
NamedTuple{(names..., n)}((values(nt)..., x))
end
function f(x)
local y
locals = NamedTuple()
if @isdefined x
locals = Base.setindex(locals, x, Val{:x}())
end
if @isdefined y
locals = Base.setindex(locals, y, Val{:y}())
end
locals
endwhich the optimizer has no problem with:
julia> @code_typed f(1)
CodeInfo(
1 ─ goto #3 if not true
2 ─ %2 = %new(NamedTuple{(:x,), Tuple{Int64}}, x)::NamedTuple{(:x,), Tuple{Int64}}
3 ┄ %3 = φ (#2 => %2)::NamedTuple{(:x,), Tuple{Int64}}
└── return %3
) => NamedTuple{(:x,), Tuple{Int64}}In code with very very many local variables, this appears to slow things down, but does still get resolved in a type stable manner.
E.g.
julia> function Base.setindex(nt::NamedTuple{names}, x, ::Val{n}) where {names, n}
NamedTuple{(names..., n)}((values(nt)..., x))
end
julia> let N = 12, M = 21
args = [Symbol(:arg, n) for n ∈ 1:N]
vars = [Symbol(:var, m) for m ∈ 1:M]
defargs = map(enumerate(vars)) do (m, var)
if isodd(m)
:(local $var)
else
:($var = $(rand((1, "hi", QuoteNode(:bye), 2.0, 3 => 4 + im))))
end
end
defs = Expr(:block, defargs...)
setindicesargs = map([args; vars]) do x
:(if @isdefined $x
nt = Base.setindex(nt, $x, Val{$(QuoteNode(x))}())
end)
end
thelocals = Expr(:block, :(nt = NamedTuple()), setindicesargs..., :nt)
@eval @time begin
#quote
function f($(args...),)
$defs
$thelocals
end
Core.Compiler.return_type(f, Tuple{$(rand((Int, String, Symbol, ComplexF64), N)...),})
end
end
0.056692 seconds (487.75 k allocations: 23.483 MiB, 92.06% compilation time)
NamedTuple{(:arg1, :arg2, :arg3, :arg4, :arg5, :arg6, :arg7, :arg8, :arg9, :arg10, :arg11, :arg12, :var2, :var4, :var6, :var8, :var10, :var12, :var14, :var16, :var18, :var20), Tuple{Int64, String, Int64, Int64, ComplexF64, ComplexF64, String, String, Symbol, Int64, ComplexF64, String, Float64, Float64, Symbol, Pair{Int64, Complex{Int64}}, Symbol, Pair{Int64, Complex{Int64}}, Symbol, Symbol, String, String}}Metadata
Metadata
Assignees
Labels
No labels