Make EnterNode save/restore dynamic scope #52309

Keno · 2023-11-26T16:10:25Z

As discussed in #51352, this gives EnterNode the ability to set (and restore on leave or catch edge) jl_current_task->scope. Manual modifications of the task field after the task has started are considered undefined behavior. In addition, we gain a new intrinsic to access current_task->scope and both inference and the optimizer will forward scopes from EnterNodes to this intrinsic (non-interprocedurally). Together with #51993 this is sufficient to fully optimize ScopedValues (non-interprocedurally at least).

vtjnash · 2023-11-26T21:32:54Z

src/codegen.cpp

 std::map<int, jl_varinfo_t> phic_slots;
+ std::map<int, std::pair<Value*, Value*> > scope_restore;


I thought we were preferring DenseMap since it has better performance and [] behavior, or maybe that does still need .at(), but at least has that operator?

I thought std::map was better for small maps, but I don't really care either way.

For small maps we should probably use SmallDenseMap

std::map has similar characteristics to std::set: it uses a single allocation per pair inserted into the map, it offers log(n) lookup with an extremely large constant factor, imposes a space penalty of 3 pointers per pair in the map, etc.

Looks like the gcc version that we have on CI has a false positive warning in SmallDenseMap. I really don't think it's worth micro-optimizing the datastructure here, so I'm just gonna switch this back to std::map

src/codegen.cpp

src/interpreter.c

src/julia-syntax.scm

base/compiler/tfuncs.jl

vtjnash · 2023-11-26T21:55:44Z

base/compiler/tfuncs.jl

@@ -2488,6 +2488,12 @@ function builtin_effects(𝕃::AbstractLattice, @nospecialize(f::Builtin), argty
 return Effects(EFFECTS_TOTAL;
 consistent = (isa(setting, Const) && setting.val === :conditional) ? ALWAYS_TRUE : ALWAYS_FALSE,
 nothrow = compilerbarrier_nothrow(setting, nothing))
+ elseif f === Core.current_scope
+ length(argtypes) == 0 || return Effects(EFFECTS_THROWS; consistent=ALWAYS_FALSE)


argtypes == Any[Vararg] is valid here too

And if length(argtypes) >= 1 && !isvarargtype(argtypes[1]), then isn't it just EFFECTS_THROWS?

vtjnash · 2023-11-26T22:01:14Z

base/scopedvalues.jl

- ct.scope = $(Scope)(current_scope, $(exprs...))
- $(Expr(:tryfinally, esc(ex), :(ct.scope = current_scope)))
- end
+ Expr(:tryfinally, esc(ex), :(), :($(Scope)($(Core.current_scope)()::Union{Nothing, Scope}, $(exprs...))))


Suggested change

Expr(:tryfinally, esc(ex), :(), :($(Scope)($(Core.current_scope)()::Union{Nothing, Scope}, $(exprs...))))

Expr(:tryfinally, esc(ex), :(), :(Scope(Core.current_scope()::Union{Nothing, Scope}, $(exprs...))))

It appears that for this PR to work as intended, this call of Scope(...) needs to be constant-folded? And am I right in assuming that #51993 will make this happen? I'm considering if a standalone test case can be crafted for the basic functionality this PR implements.

No, constant-folding is not required. The inference support isn't really needed either for the base version of this, but I anticipate needing it downstream.

vtjnash · 2023-11-26T22:04:45Z

base/compiler/abstractinterpretation.jl

@@ -3259,6 +3259,19 @@ function typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
 elseif isa(stmt, EnterNode)
 ssavaluetypes[currpc] = Any
 add_curr_ssaflag!(frame, IR_FLAG_NOTHROW)
+ if isdefined(stmt, :scope)


Does inference / SROA need to understand that this is now a clobbering statement for the current_task scope field, since the scope field is still publicly visible? Or should we hide the field so that users can't break the invariants expected and proven of it?

I don't think either of them can optimize these fields on the current task. I would be open to hiding the scope field further, but I didn't think it was that necessary.

base/compiler/ssair/passes.jl

aviatesk · 2023-11-27T12:01:59Z

base/compiler/tfuncs.jl

@@ -2488,6 +2488,12 @@ function builtin_effects(𝕃::AbstractLattice, @nospecialize(f::Builtin), argty
 return Effects(EFFECTS_TOTAL;
 consistent = (isa(setting, Const) && setting.val === :conditional) ? ALWAYS_TRUE : ALWAYS_FALSE,
 nothrow = compilerbarrier_nothrow(setting, nothing))
+ elseif f === Core.current_scope
+ length(argtypes) == 0 || return Effects(EFFECTS_THROWS; consistent=ALWAYS_FALSE)


And if length(argtypes) >= 1 && !isvarargtype(argtypes[1]), then isn't it just EFFECTS_THROWS?

aviatesk · 2023-11-27T12:14:17Z

base/boot.jl

+EnterNode(old::EnterNode, new_dest::Int) = isdefined(old, :scope) ?
+ EnterNode(new_dest, old.scope) : EnterNode(new_dest)


Suggested change

EnterNode(old::EnterNode, new_dest::Int) = isdefined(old, :scope) ?

EnterNode(new_dest, old.scope) : EnterNode(new_dest)

function EnterNode(node::EnterNode;

dest::Int=node.dest,

scope)

@nospecialize scope

if !(@isdefined scope)

isdefined(node, :scope) || return EnterNode(dest)

scope = node.scope

end

return EnterNode(dest, scope)

end

Considering the possibility that EnterNode could get more complicated in the future (I'm not sure if it is likely though), defining it as Effects(effects::Effects; overrides...)-like constructor might be a better move?

aviatesk · 2023-11-27T12:35:14Z

base/compiler/tfuncs.jl

+function current_scope_tfunc(interp::AbstractInterpreter, sv::InferenceState)
+ pc = sv.currpc
+ while true
+ handleridx = sv.handler_at[pc][2]


Suggested change

handleridx = sv.handler_at[pc][2]

handleridx = sv.handler_at[pc][1]

Shouldn't this look at the active try region rather than the catch region?

base/compiler/tfuncs.jl

aviatesk · 2023-11-27T12:48:55Z

base/scopedvalues.jl

- ct.scope = $(Scope)(current_scope, $(exprs...))
- $(Expr(:tryfinally, esc(ex), :(ct.scope = current_scope)))
- end
+ Expr(:tryfinally, esc(ex), :(), :($(Scope)($(Core.current_scope)()::Union{Nothing, Scope}, $(exprs...))))


It appears that for this PR to work as intended, this call of Scope(...) needs to be constant-folded? And am I right in assuming that #51993 will make this happen? I'm considering if a standalone test case can be crafted for the basic functionality this PR implements.

aviatesk · 2023-11-27T13:00:53Z

Manual modifications of the task field after the task has started are considered undefined behavior.

Maybe we should make it the only way to set a scope, and raise an explicit error if user manually modifies the scope field of a task?

Keno · 2023-11-27T15:26:20Z

Maybe we should make it the only way to set a scope, and raise an explicit error if user manually modifies the scope field of a task?

I do think we need to retain the ability to set the scope before a task is started. It's not used here, but I suspect we may want it in the future. However, we could certainly have a setproperty! that asserts that the task has not been started yet.

aviatesk · 2023-11-28T11:44:10Z

base/compiler/tfuncs.jl

+ pchandler = sv.handlers[handleridx]
+ # Remember that we looked at this handler, so we get re-scheduled
+ # if the scope information changes
+ isdefined(pchandler, :scope_uses) || (pchandler.scope_uses = Int[])


Can't we set scope_uses within compute_trycatch?

We cannot, because we don't necessarily know all the callsites that will infer to a call to current_scope. We could make current_scope an expr to make it obvious, but I don't think it's required.

As discussed in #51352, this gives `EnterNode` the ability to set (and restore on leave or catch edge) jl_current_task->scope. Manual modifications of the task field after the task has started are considered undefined behavior. In addition, we gain a new intrinsic to access current_task->scope and both inference and the optimizer will forward scopes from EnterNodes to this intrinsic (non-interprocedurally). Together with #51993 this is sufficient to fully optimize ScopedValues (non-interprocedurally at least).

Keno requested review from vchuravy and aviatesk November 26, 2023 16:10

Keno force-pushed the kf/enterscope branch from 171c4bf to a352fc7 Compare November 26, 2023 16:29