Skip to content

Commit

Permalink
optimizer: Add early finalize calls
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpsamaroo
jpsamaroo committed Mar 8, 2022
1 parent 6340ff0 commit 75f83ed
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 9 deletions.
125 changes: 117 additions & 8 deletions base/compiler/optimize.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ function stmt_effect_free end # imported by EscapeAnalysis
function alloc_array_ndims end # imported by EscapeAnalysis
include("compiler/ssair/driver.jl")
using .EscapeAnalysis
import .EscapeAnalysis: EscapeState, ArgEscapeCache, is_ipo_profitable
import .EscapeAnalysis: EscapeState, EscapeInfo, ArgEscapeCache, is_ipo_profitable

"""
cache_escapes!(caller::InferenceResult, estate::EscapeState)
Expand All @@ -111,6 +111,110 @@ function ipo_escape_cache(mi_cache::MICache) where MICache
end
end
null_escape_cache(linfo::Union{InferenceResult,MethodInstance}) = nothing
EARLY_FINALIZE_VERBOSE = RefValue{Bool}(false)
"""
early_finalize!(ir::IRCode, estate::EscapeState, domtree::DomTree) -> IRCode
Analyzes `ir` for heap allocations which escape only via `FinalizerEscape`
(thus having a `finalizer` call associated), and inserts a call to
`finalize(obj)` just before any return where the allocation doesn't escape.
"""
function early_finalize!(ir::IRCode, estate::EscapeState, domtree::DomTree)
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int), "early_finalize!: Analyzing with nargs %d\n", estate.nargs)
# Find all allocations that escape only through a finalizer
isalloc(ir::IRCode, pc::Int) = isexpr(ir.stmts[pc][:inst], :new)
to_finalize = Pair{Int,EscapeInfo}[]
if length(estate.escapes) != estate.nargs+length(ir.stmts)
#=EARLY_FINALIZE_VERBOSE[] &&=# ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int, Int), "early_finalize!: Mismatched stmts (%d) with escape infos (%d)\n", length(ir.stmts), length(estate.escapes))
end
for idx in (estate.nargs+1):length(estate.escapes)
info = estate.escapes[idx]
pc = idx - estate.nargs
if pc > length(ir.stmts)
# FIXME
continue
end
if has_finalizer_escape(info) && isalloc(ir, pc)
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int), "early_finalize!: Found finalizing alloc at %d\n", pc)
push!(to_finalize, pc => info)
end
end

# Find all locations to insert calls to `finalize`
# Currently this is just all returns and unreachables
return_locs = Int[]
for bb in ir.cfg.blocks
inst = ir.stmts.inst[bb.stmts.stop]
if inst isa ReturnNode
if isdefined(inst, :val)
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int), "early_finalize!: Found return at %d\n", bb.stmts.stop)
push!(return_locs, bb.stmts.stop)
else
# Unreachable, try to find `Base.throw` one insn before
block = ir.cfg.blocks[block_for_inst(ir, bb.stmts.stop)]
throw_pc = bb.stmts.stop - 1
if throw_pc in block.stmts
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int), "early_finalize!: Found throw at %d\n", throw_pc)
push!(return_locs, throw_pc)
end
end
end
end

# Insert `finalize` calls in order
for idx in return_locs
non_escaping = Pair{Int,EscapeInfo}[]
for (alloc_idx, alloc_info) in to_finalize
# Check if this allocation is otherwise unescaping at this return
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int, Int, Bool, Bool, Bool),
"early_finalize!: For alloc at %d, return at %d: ReturnEscape(%d), ArgEscape(%d), ThrownEscape(%d)\n",
alloc_idx, idx, has_return_escape(alloc_info, idx), has_arg_escape(alloc_info), has_thrown_escape(alloc_info, idx))
if has_return_escape(alloc_info, idx) ||
has_arg_escape(alloc_info) ||
has_thrown_escape(alloc_info, idx)
continue
end
# Construct all blocks for which the allocation is live
allblocks = Int[]
fdu = FieldDefUse()
for pc in alloc_info.Liveness
pc < 1 && continue
if !(1 <= pc <= length(ir.stmts))
ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int, Int, Int),
"early_finalize!: pc (%d) out-of-bounds (1:%d), nargs %d\n",
pc, length(ir.stmts), estate.nargs)
continue
end
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int, Int, Int), "early_finalize!: For alloc at %d, return at %d, value is live at %d\n", alloc_idx, idx, pc)
push!(fdu.defs, pc)
push!(allblocks, block_for_inst(ir, pc))
end
# Add this allocation if there is any valid def-use pair
# N.B. not necessary: `push!(fdu.uses, idx)`
use = idx
defs = find_def_for_use(ir, domtree, allblocks, fdu, use)
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int, Int, Bool),
"early_finalize!: For alloc at %d, return at %d: Defs(%d)\n",
alloc_idx, idx, defs !== nothing ? defs[1]: -1)
if defs !== nothing && defs[1] == alloc_idx
push!(non_escaping, alloc_idx=>alloc_info)
end
end
if length(non_escaping) > 0
ct_expr = Expr(:foreigncall, :(:jl_get_current_task), Ref{Task}, svec(), 0, :(:ccall))
ct_insn = NewInstruction(ct_expr, Ref{Task})
ct = insert_node!(ir, idx, ct_insn)
for (alloc_idx, alloc_info) in non_escaping
EARLY_FINALIZE_VERBOSE[] && ccall(:jl_safe_printf, Cvoid, (Ptr{UInt8}, Int, Int), "early_finalize!: For alloc at %d, return at %d, inserting finalize\n", alloc_idx, idx)
fin_expr = Expr(:foreigncall, :(:jl_finalize_th), Nothing, svec(Any, Any), 0, :(:ccall), ct, SSAValue(alloc_idx))
fin_insn = NewInstruction(fin_expr, Nothing)
fin = insert_node!(ir, idx, fin_insn)
end
end
end

ir
end

mutable struct OptimizationState
linfo::MethodInstance
Expand Down Expand Up @@ -534,7 +638,8 @@ function optimize(interp::AbstractInterpreter, opt::OptimizationState,
@timeit "optimizer" ir = run_passes(opt.src, opt, caller)
return finish(interp, opt, params, ir, caller)
end

const IR_SAVE = [false]
const IR_CACHE = Any[]
function run_passes(ci::CodeInfo, sv::OptimizationState, caller::InferenceResult)
@timeit "convert" ir = convert_to_ircode(ci, sv)
@timeit "slot2reg" ir = slot2reg(ir, ci, sv)
Expand All @@ -552,19 +657,23 @@ function run_passes(ci::CodeInfo, sv::OptimizationState, caller::InferenceResult
# @timeit "verify 2" verify_ir(ir)
@timeit "compact 2" ir = compact!(ir)
@timeit "SROA" ir, memory_opt = linear_pass!(ir)
if memory_opt
@timeit "memory_opt_pass!" begin
@timeit "Local EA" estate = analyze_escapes(ir,
nargs, #=call_resolved=#true, null_escape_cache)
@timeit "memory_opt_pass!" ir = memory_opt_pass!(ir, estate)
end
@timeit "pre-EA compact" ir = compact!(ir)
@timeit "Local EA" estate = analyze_escapes(ir, nargs, #=call_resolved=#true, null_escape_cache)
if false #memory_opt
@timeit "memory_opt_pass!" ir = memory_opt_pass!(ir, estate)
end
@timeit "EA domtree" domtree = construct_domtree(ir.cfg.blocks)
@timeit "early_finalize_pass!" ir = early_finalize!(ir, estate, domtree)
push!(IR_CACHE, ir)
@timeit "ADCE" ir = adce_pass!(ir)
@timeit "type lift" ir = type_lift_pass!(ir)
@timeit "compact 3" ir = compact!(ir)
if JLOptions().debug_level == 2
@timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))
end
if !IR_SAVE[1]
pop!(IR_CACHE)
end
return ir
end

Expand Down
8 changes: 7 additions & 1 deletion base/compiler/ssair/passes.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,13 @@ function linear_pass!(ir::IRCode)
elseif isexpr(def, :new)
typ = unwrap_unionall(widenconst(argextype(SSAValue(defidx), compact)))
if typ isa DataType
ismutabletype(typ) && continue # mutable SROA is performed later
if ismutabletype(typ)
if stmt.args[1] isa QuoteNode && stmt.args[1].value == :jl_gc_add_finalizer_th
# Track this for `early_finalize!`
anymutability = true
end
continue # mutable SROA is performed later
end
record_immutable_preserve!(new_preserves, def, compact)
push!(preserved, preserved_arg.id)
end
Expand Down

0 comments on commit 75f83ed

Please sign in to comment.