skip inferring calls that lead to throw (#35982)

Author: JeffBezanson

base/compiler/abstractinterpretation.jl

@@ -36,6 +36,9 @@ end
 
 function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f), argtypes::Vector{Any}, @nospecialize(atype), sv::InferenceState,
                                   max_methods::Int = InferenceParams(interp).MAX_METHODS)
+    if sv.currpc in sv.throw_blocks
+        return Any
+    end
     mt = ccall(:jl_method_table_for, Any, (Any,), atype)
     if mt === nothing
         add_remark!(interp, sv, "Could not identify method table for call")

base/compiler/inferencestate.jl

@@ -31,6 +31,7 @@ mutable struct InferenceState
     n_handlers::Int
     # ssavalue sparsity and restart info
     ssavalue_uses::Vector{BitSet}
+    throw_blocks::BitSet
 
     cycle_backedges::Vector{Tuple{InferenceState, LineNum}} # call-graph backedges connecting from callee to caller
     callers_in_cycle::Vector{InferenceState}
@@ -80,6 +81,7 @@ mutable struct InferenceState
         s_types[1] = s_argtypes
 
         ssavalue_uses = find_ssavalue_uses(code, nssavalues)
+        throw_blocks = find_throw_blocks(code)
 
         # exception handlers
         cur_hand = nothing
@@ -106,7 +108,7 @@ mutable struct InferenceState
             nargs, s_types, s_edges,
             Union{}, W, 1, n,
             cur_hand, handler_at, n_handlers,
-            ssavalue_uses,
+            ssavalue_uses, throw_blocks,
             Vector{Tuple{InferenceState,LineNum}}(), # cycle_backedges
             Vector{InferenceState}(), # callers_in_cycle
             #=parent=#nothing,

base/compiler/optimize.jl

@@ -285,7 +285,7 @@ plus_saturate(x::Int, y::Int) = max(x, y, x+y)
 # known return type
 isknowntype(@nospecialize T) = (T === Union{}) || isa(T, Const) || isconcretetype(widenconst(T))
 
-function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}, slottypes::Vector{Any}, params::OptimizationParams)
+function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}, slottypes::Vector{Any}, params::OptimizationParams, error_path::Bool = false)
     head = ex.head
     if is_meta_expr_head(head)
         return 0
@@ -320,7 +320,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
                 return 0
             elseif (f === Main.Core.arrayref || f === Main.Core.const_arrayref) && length(ex.args) >= 3
                 atyp = argextype(ex.args[3], src, sptypes, slottypes)
-                return isknowntype(atyp) ? 4 : params.inline_nonleaf_penalty
+                return isknowntype(atyp) ? 4 : error_path ? params.inline_error_path_cost : params.inline_nonleaf_penalty
             end
             fidx = find_tfunc(f)
             if fidx === nothing
@@ -330,7 +330,11 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
             end
             return T_FFUNC_COST[fidx]
         end
-        return params.inline_nonleaf_penalty
+        extyp = line == -1 ? Any : src.ssavaluetypes[line]
+        if extyp === Union{}
+            return 0
+        end
+        return error_path ? params.inline_error_path_cost : params.inline_nonleaf_penalty
     elseif head === :foreigncall || head === :invoke
         # Calls whose "return type" is Union{} do not actually return:
         # they are errors. Since these are not part of the typical
@@ -347,7 +351,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
         end
         a = ex.args[2]
         if a isa Expr
-            cost = plus_saturate(cost, statement_cost(a, -1, src, sptypes, slottypes, params))
+            cost = plus_saturate(cost, statement_cost(a, -1, src, sptypes, slottypes, params, error_path))
         end
         return cost
     elseif head === :copyast
@@ -365,10 +369,11 @@ end
 function inline_worthy(body::Array{Any,1}, src::CodeInfo, sptypes::Vector{Any}, slottypes::Vector{Any},
                        params::OptimizationParams, cost_threshold::Integer=params.inline_cost_threshold)
     bodycost::Int = 0
+    throw_blocks = find_throw_blocks(body)
     for line = 1:length(body)
         stmt = body[line]
         if stmt isa Expr
-            thiscost = statement_cost(stmt, line, src, sptypes, slottypes, params)::Int
+            thiscost = statement_cost(stmt, line, src, sptypes, slottypes, params, line in throw_blocks)::Int
         elseif stmt isa GotoNode
             # loops are generally always expensive
             # but assume that forward jumps are already counted for from

base/compiler/ssair/inlining.jl

@@ -997,7 +997,9 @@ end
 function assemble_inline_todo!(ir::IRCode, sv::OptimizationState)
     # todo = (inline_idx, (isva, isinvoke, na), method, spvals, inline_linetable, inline_ir, lie)
     todo = Any[]
+    skip = find_throw_blocks(ir.stmts.inst, RefValue(ir))
     for idx in 1:length(ir.stmts)
+        idx in skip && continue
         r = process_simple!(ir, idx, sv.params, sv.world)
         r === nothing && continue
 

base/compiler/types.jl

@@ -45,6 +45,7 @@ struct OptimizationParams
     inline_cost_threshold::Int  # number of CPU cycles beyond which it's not worth inlining
     inline_nonleaf_penalty::Int # penalty for dynamic dispatch
     inline_tupleret_bonus::Int  # extra willingness for non-isbits tuple return types
+    inline_error_path_cost::Int # cost of (un-optimized) calls in blocks that throw
 
     # Duplicating for now because optimizer inlining requires it.
     # Keno assures me this will be removed in the near future
@@ -57,6 +58,7 @@ struct OptimizationParams
             inline_cost_threshold::Int = 100,
             inline_nonleaf_penalty::Int = 1000,
             inline_tupleret_bonus::Int = 400,
+            inline_error_path_cost::Int = 20,
             max_methods::Int = 3,
             tuple_splat::Int = 32,
             union_splitting::Int = 4,
@@ -66,6 +68,7 @@ struct OptimizationParams
             inline_cost_threshold,
             inline_nonleaf_penalty,
             inline_tupleret_bonus,
+            inline_error_path_cost,
             max_methods,
             tuple_splat,
             union_splitting,

base/compiler/utilities.jl

@@ -219,6 +219,72 @@ function find_ssavalue_uses(e::Expr, uses::Vector{BitSet}, line::Int)
     end
 end
 
+function is_throw_call(e::Expr)
+    if e.head === :call
+        f = e.args[1]
+        if isa(f, GlobalRef)
+            ff = abstract_eval_global(f.mod, f.name)
+            if isa(ff, Const) && ff.val === Core.throw
+                return true
+            end
+        end
+    end
+    return false
+end
+
+function find_throw_blocks(code::Vector{Any}, ir = RefValue{IRCode}())
+    stmts = BitSet()
+    n = length(code)
+    try_depth = 0
+    for i in n:-1:1
+        s = code[i]
+        if isa(s, Expr)
+            if s.head === :enter
+                try_depth -= 1
+            elseif s.head === :leave
+                try_depth += (s.args[1]::Int)
+            elseif s.head === :gotoifnot
+                tgt = s.args[2]::Int
+                if i+1 in stmts && tgt in stmts
+                    push!(stmts, i)
+                end
+            elseif s.head === :return
+            elseif is_throw_call(s) || s.head === :unreachable
+                if try_depth == 0
+                    push!(stmts, i)
+                end
+            elseif i+1 in stmts
+                push!(stmts, i)
+            end
+        elseif isa(s, ReturnNode)
+            if try_depth == 0 && !isdefined(s, :val)
+                push!(stmts, i)
+            end
+        elseif isa(s, GotoNode)
+            tgt = s.label
+            if isassigned(ir)
+                tgt = first(ir[].cfg.blocks[tgt].stmts)
+            end
+            if tgt in stmts
+                push!(stmts, i)
+            end
+        elseif isa(s, GotoIfNot)
+            if i+1 in stmts
+                tgt = s.dest::Int
+                if isassigned(ir)
+                    tgt = first(ir[].cfg.blocks[tgt].stmts)
+                end
+                if tgt in stmts
+                    push!(stmts, i)
+                end
+            end
+        elseif i+1 in stmts
+            push!(stmts, i)
+        end
+    end
+    return stmts
+end
+
 # using a function to ensure we can infer this
 @inline slot_id(s) = isa(s, SlotNumber) ? (s::SlotNumber).id :
     isa(s, Argument) ? (s::Argument).n : (s::TypedSlot).id

base/ntuple.jl

@@ -14,7 +14,8 @@ julia> ntuple(i -> 2*i, 4)
 (2, 4, 6, 8)
 ```
 """
-function ntuple(f::F, n::Integer) where F
+@inline function ntuple(f::F, n::Integer) where F
+    # marked inline since this benefits from constant propagation of `n`
     t = n == 0  ? () :
         n == 1  ? (f(1),) :
         n == 2  ? (f(1), f(2)) :

test/stacktraces.jl

@@ -84,8 +84,8 @@ end
 
 module inlined_test
 using Test
-@inline g(x) = (y = throw("a"); y) # the inliner does not insert the proper markers when inlining a single expression
-@inline h(x) = (y = g(x); y)       # this test could be extended to check for that if we switch to linear representation
+@inline g(x) = (x == 3 && throw("a"); x)
+@inline h(x) = (x == 3 && g(x); x)
 f(x) = (y = h(x); y)
 trace = (try; f(3); catch; stacktrace(catch_backtrace()); end)[1:3]
 can_inline = Bool(Base.JLOptions().can_inline)