Skip to content

Commit 2cfd57d

Browse files
aviateskaviatesk
authored
optimizer: minor improvements on SROA pass (#43348)
- avoid closures - factor out duplicated code
1 parent db1d2f5 commit 2cfd57d

File tree

2 files changed

+63
-68
lines changed

2 files changed

+63
-68
lines changed

‎base/compiler/ssair/inlining.jl‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -370,9 +370,8 @@ function ir_inline_item!(compact::IncrementalCompact, idx::Int, argexprs::Vector
370370
isa(val, SSAValue) && (compact.used_ssas[val.id] += 1)
371371
return_value = SSAValue(idx′)
372372
inline_compact[idx′] = val
373-
inline_compact.result[idx′][:type] = (isa(val, Argument) || isa(val, Expr)) ?
374-
compact_exprtype(compact, val) :
375-
compact_exprtype(inline_compact, val)
373+
inline_compact.result[idx′][:type] =
374+
compact_exprtype(isa(val, Argument) || isa(val, Expr) ? compact : inline_compact, val)
376375
break
377376
end
378377
inline_compact[idx′] = stmt′
@@ -400,9 +399,8 @@ function ir_inline_item!(compact::IncrementalCompact, idx::Int, argexprs::Vector
400399
push!(pn.edges, inline_compact.active_result_bb-1)
401400
if isa(val, GlobalRef) || isa(val, Expr)
402401
stmt′ = val
403-
inline_compact.result[idx′][:type] = (isa(val, Argument) || isa(val, Expr)) ?
404-
compact_exprtype(compact, val) :
405-
compact_exprtype(inline_compact, val)
402+
inline_compact.result[idx′][:type] =
403+
compact_exprtype(isa(val, Expr) ? compact : inline_compact, val)
406404
insert_node_here!(inline_compact, NewInstruction(GotoNode(post_bb_id),
407405
Any, compact.result[idx′][:line]),
408406
true)
@@ -411,7 +409,6 @@ function ir_inline_item!(compact::IncrementalCompact, idx::Int, argexprs::Vector
411409
push!(pn.values, val)
412410
stmt′ = GotoNode(post_bb_id)
413411
end
414-
415412
end
416413
elseif isa(stmt′, GotoNode)
417414
stmt′ = GotoNode(stmt′.label + bb_offset)

‎base/compiler/ssair/passes.jl‎

Lines changed: 59 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -323,50 +323,25 @@ function is_getfield_captures(@nospecialize(def), compact::IncrementalCompact)
323323
return oc ⊑ Core.OpaqueClosure
324324
end
325325

326+
struct LiftedValue
327+
x
328+
LiftedValue(@nospecialize x) = new(x)
329+
end
330+
const LiftedLeaves = IdDict{Any, Union{Nothing,LiftedValue}}
331+
326332
# try to compute lifted values that can replace `getfield(x, field)` call
327333
# where `x` is an immutable struct that are defined at any of `leaves`
328334
function lift_leaves(compact::IncrementalCompact,
329335
@nospecialize(result_t), field::Int, leaves::Vector{Any})
330336
# For every leaf, the lifted value
331-
lifted_leaves = IdDict{Any, Union{Nothing,LiftedValue}}()
337+
lifted_leaves = LiftedLeaves()
332338
maybe_undef = false
333339
for leaf in leaves
334-
leaf_key = leaf
340+
cache_key = leaf
335341
if isa(leaf, AnySSAValue)
336-
function lift_arg(ref::Core.Compiler.UseRef)
337-
lifted = ref[]
338-
if is_old(compact, leaf) && isa(lifted, SSAValue)
339-
lifted = OldSSAValue(lifted.id)
340-
end
341-
if isa(lifted, GlobalRef) || isa(lifted, Expr)
342-
lifted = insert_node!(compact, leaf, effect_free(NewInstruction(lifted, compact_exprtype(compact, lifted))))
343-
ref[] = lifted
344-
(isa(leaf, SSAValue) && (leaf.id < compact.result_idx)) && push!(compact.late_fixup, leaf.id)
345-
end
346-
lifted_leaves[leaf_key] = LiftedValue(lifted)
347-
nothing
348-
end
349-
function walk_leaf(@nospecialize(leaf))
350-
if isa(leaf, OldSSAValue) && already_inserted(compact, leaf)
351-
leaf = compact.ssa_rename[leaf.id]
352-
if isa(leaf, AnySSAValue)
353-
leaf = simple_walk(compact, leaf)
354-
end
355-
if isa(leaf, AnySSAValue)
356-
def = compact[leaf]
357-
else
358-
def = leaf
359-
end
360-
elseif isa(leaf, AnySSAValue)
361-
def = compact[leaf]
362-
else
363-
def = leaf
364-
end
365-
return Pair{Any, Any}(def, leaf)
366-
end
367-
(def, leaf) = walk_leaf(leaf)
368-
if is_tuple_call(compact, def) && 1 <= field < length(def.args)
369-
lift_arg(UseRef(def, 1 + field))
342+
(def, leaf) = walk_to_def(compact, leaf)
343+
if is_tuple_call(compact, def) && 1 ≤ field < length(def.args)
344+
lift_arg!(compact, leaf, cache_key, def, 1+field, lifted_leaves)
370345
continue
371346
elseif isexpr(def, :new)
372347
typ = widenconst(types(compact)[leaf])
@@ -375,7 +350,7 @@ function lift_leaves(compact::IncrementalCompact,
375350
end
376351
(isa(typ, DataType) && !isabstracttype(typ)) || return nothing
377352
@assert !ismutabletype(typ)
378-
if length(def.args) < 1 + field
353+
if length(def.args) < 1+field
379354
if field > fieldcount(typ)
380355
return nothing
381356
end
@@ -384,7 +359,7 @@ function lift_leaves(compact::IncrementalCompact,
384359
# On this branch, this will be a guaranteed UndefRefError.
385360
# We use the regular undef mechanic to lift this to a boolean slot
386361
maybe_undef = true
387-
lifted_leaves[leaf_key] = nothing
362+
lifted_leaves[cache_key] = nothing
388363
continue
389364
end
390365
return nothing
@@ -398,26 +373,17 @@ function lift_leaves(compact::IncrementalCompact,
398373
end
399374
compact[leaf] = def
400375
end
401-
lifted = def.args[1+field]
402-
if is_old(compact, leaf) && isa(lifted, SSAValue)
403-
lifted = OldSSAValue(lifted.id)
404-
end
405-
if isa(lifted, GlobalRef) || isa(lifted, Expr)
406-
lifted = insert_node!(compact, leaf, effect_free(NewInstruction(lifted, compact_exprtype(compact, lifted))))
407-
def.args[1+field] = lifted
408-
(isa(leaf, SSAValue) && (leaf.id < compact.result_idx)) && push!(compact.late_fixup, leaf.id)
409-
end
410-
lifted_leaves[leaf_key] = LiftedValue(lifted)
376+
lift_arg!(compact, leaf, cache_key, def, 1+field, lifted_leaves)
411377
continue
412378
elseif is_getfield_captures(def, compact)
413379
# Walk to new_opaque_closure
414380
ocleaf = def.args[2]
415381
if isa(ocleaf, AnySSAValue)
416382
ocleaf = simple_walk(compact, ocleaf)
417383
end
418-
ocdef, _ = walk_leaf(ocleaf)
419-
if isexpr(ocdef, :new_opaque_closure) && isa(field, Int) && 1 <= field <= length(ocdef.args)-5
420-
lift_arg(UseRef(ocdef, 5 + field))
384+
ocdef, _ = walk_to_def(compact, ocleaf)
385+
if isexpr(ocdef, :new_opaque_closure) && isa(field, Int) && 1 ≤ field ≤ length(ocdef.args)-5
386+
lift_arg!(compact, leaf, cache_key, ocdef, 5+field, lifted_leaves)
421387
continue
422388
end
423389
return nothing
@@ -445,18 +411,55 @@ function lift_leaves(compact::IncrementalCompact,
445411
else
446412
return nothing
447413
end
448-
elseif isa(leaf, Union{Argument, Expr})
414+
elseif isa(leaf, Argument) || isa(leaf, Expr)
449415
return nothing
450416
end
451417
ismutable(leaf) && return nothing
452418
isdefined(leaf, field) || return nothing
453419
val = getfield(leaf, field)
454420
is_inlineable_constant(val) || return nothing
455-
lifted_leaves[leaf_key] = LiftedValue(quoted(val))
421+
lifted_leaves[cache_key] = LiftedValue(quoted(val))
456422
end
457423
return lifted_leaves, maybe_undef
458424
end
459425

426+
function lift_arg!(
427+
compact::IncrementalCompact, @nospecialize(leaf), @nospecialize(cache_key),
428+
stmt::Expr, argidx::Int, lifted_leaves::LiftedLeaves)
429+
lifted = stmt.args[argidx]
430+
if is_old(compact, leaf) && isa(lifted, SSAValue)
431+
lifted = OldSSAValue(lifted.id)
432+
end
433+
if isa(lifted, GlobalRef) || isa(lifted, Expr)
434+
lifted = insert_node!(compact, leaf, effect_free(NewInstruction(lifted, compact_exprtype(compact, lifted))))
435+
stmt.args[argidx] = lifted
436+
if isa(leaf, SSAValue) && leaf.id < compact.result_idx
437+
push!(compact.late_fixup, leaf.id)
438+
end
439+
end
440+
lifted_leaves[cache_key] = LiftedValue(lifted)
441+
nothing
442+
end
443+
444+
function walk_to_def(compact::IncrementalCompact, @nospecialize(leaf))
445+
if isa(leaf, OldSSAValue) && already_inserted(compact, leaf)
446+
leaf = compact.ssa_rename[leaf.id]
447+
if isa(leaf, AnySSAValue)
448+
leaf = simple_walk(compact, leaf)
449+
end
450+
if isa(leaf, AnySSAValue)
451+
def = compact[leaf]
452+
else
453+
def = leaf
454+
end
455+
elseif isa(leaf, AnySSAValue)
456+
def = compact[leaf]
457+
else
458+
def = leaf
459+
end
460+
return Pair{Any, Any}(def, leaf)
461+
end
462+
460463
make_MaybeUndef(@nospecialize(typ)) = isa(typ, MaybeUndef) ? typ : MaybeUndef(typ)
461464

462465
"""
@@ -505,7 +508,7 @@ function lift_comparison!(compact::IncrementalCompact,
505508
r = egal_tfunc(compact_exprtype(compact, leaf), cmp)
506509
if isa(r, Const)
507510
if lifted_leaves === nothing
508-
lifted_leaves = IdDict{Any, Union{Nothing,LiftedValue}}()
511+
lifted_leaves = LiftedLeaves()
509512
end
510513
lifted_leaves[leaf] = LiftedValue(r.val)
511514
else
@@ -515,7 +518,7 @@ function lift_comparison!(compact::IncrementalCompact,
515518

516519
lifted_val = perform_lifting!(compact,
517520
visited_phinodes, cmp, lifting_cache, Bool,
518-
lifted_leaves::IdDict{Any, Union{Nothing,LiftedValue}}, val)::LiftedValue
521+
lifted_leaves::LiftedLeaves, val)::LiftedValue
519522

520523
compact[idx] = lifted_val.x
521524
end
@@ -532,15 +535,10 @@ function is_old(compact, @nospecialize(old_node_ssa))
532535
!already_inserted(compact, old_node_ssa)
533536
end
534537

535-
struct LiftedValue
536-
x
537-
LiftedValue(@nospecialize x) = new(x)
538-
end
539-
540538
function perform_lifting!(compact::IncrementalCompact,
541539
visited_phinodes::Vector{AnySSAValue}, @nospecialize(cache_key),
542540
lifting_cache::IdDict{Pair{AnySSAValue, Any}, AnySSAValue},
543-
@nospecialize(result_t), lifted_leaves::IdDict{Any, Union{Nothing,LiftedValue}}, @nospecialize(stmt_val))
541+
@nospecialize(result_t), lifted_leaves::LiftedLeaves, @nospecialize(stmt_val))
544542
reverse_mapping = IdDict{AnySSAValue, Int}(ssa => id for (id, ssa) in enumerate(visited_phinodes))
545543

546544
# Insert PhiNodes

0 commit comments

Comments
 (0)