optimizer: lift more comparisons (#43227)

This commit implements more comparison liftings.
Especially, this change enables the compiler to lift `isa`/`isdefined`
checks (i.e. replace a comparison call with ϕ-node by CFG
union-splitting).

For example, the code snippet below will run 500x faster:
```julia
function compute(n)
    s = 0
    itr = 1:n
    st = iterate(itr)
    while isdefined(st, 2) # mimic our iteration protocol with
`isdefined`
        v, st = st
        s += v
        st = iterate(itr, st)
    end
    s
end
```

Although it seems like the codegen for `isa` is fairly optimized already
and so I could not find any performance benefit for `isa`-lifting
(`code_llvm` emits mostly equivalent code), but I hope it's more ideal
if we can do the equivalent optimization on Julia level so that we can
just consult to `code_typed` for performance optimization.

Author: aviatesk

base/compiler/ssair/passes.jl

@@ -471,19 +471,21 @@ end
 make_MaybeUndef(@nospecialize(typ)) = isa(typ, MaybeUndef) ? typ : MaybeUndef(typ)
 
 """
-    lift_comparison!(compact::IncrementalCompact, idx::Int, stmt::Expr)
+    lift_comparison!(cmp, compact::IncrementalCompact, idx::Int, stmt::Expr)
 
-Replaces `φ(x, y)::Union{X,Y} === constant` by `φ(x === constant, y === constant)`,
-where `x === constant` and `y === constant` can be replaced with constant `Bool`eans.
-It helps codegen avoid generating expensive code for `===` with `Union` types.
+Replaces `cmp(φ(x, y)::Union{X,Y}, constant)` by `φ(cmp(x, constant), cmp(y, constant))`,
+where `cmp(x, constant)` and `cmp(y, constant)` can be replaced with constant `Bool`eans.
+It helps codegen avoid generating expensive code for `cmp` with `Union` types.
 In particular, this is supposed to improve the performance of the iteration protocol:
 ```julia
 while x !== nothing
     x = iterate(...)::Union{Nothing,Tuple{Any,Any}}
 end
 ```
 """
-function lift_comparison!(compact::IncrementalCompact,
+function lift_comparison! end
+
+function lift_comparison!(::typeof(===), compact::IncrementalCompact,
     idx::Int, stmt::Expr, lifting_cache::IdDict{Pair{AnySSAValue, Any}, AnySSAValue})
     args = stmt.args
     length(args) == 3 || return
@@ -493,37 +495,59 @@ function lift_comparison!(compact::IncrementalCompact,
     vr = argextype(rhs, compact)
     if isa(vl, Const)
         isa(vr, Const) && return
-        cmp = vl
-        typeconstraint = widenconst(vr)
         val = rhs
+        target = lhs
     elseif isa(vr, Const)
-        cmp = vr
-        typeconstraint = widenconst(vl)
         val = lhs
+        target = rhs
     else
         return
     end
 
-    valtyp = widenconst(argextype(val, compact))
-    isa(valtyp, Union) || return # bail out if there won't be a good chance for lifting
+    lift_comparison_leaves!(egal_tfunc, compact, val, target, lifting_cache, idx)
+end
 
-    leaves, visited_phinodes = collect_leaves(compact, val, valtyp)
+function lift_comparison!(::typeof(isa), compact::IncrementalCompact,
+    idx::Int, stmt::Expr, lifting_cache::IdDict{Pair{AnySSAValue, Any}, AnySSAValue})
+    args = stmt.args
+    length(args) == 3 || return
+    lift_comparison_leaves!(isa_tfunc, compact, args[2], args[3], lifting_cache, idx)
+end
+
+function lift_comparison!(::typeof(isdefined), compact::IncrementalCompact,
+    idx::Int, stmt::Expr, lifting_cache::IdDict{Pair{AnySSAValue, Any}, AnySSAValue})
+    args = stmt.args
+    length(args) == 3 || return
+    lift_comparison_leaves!(isdefined_tfunc, compact, args[2], args[3], lifting_cache, idx)
+end
+
+function lift_comparison_leaves!(@specialize(tfunc),
+    compact::IncrementalCompact, @nospecialize(val), @nospecialize(target),
+    lifting_cache::IdDict{Pair{AnySSAValue, Any}, AnySSAValue}, idx::Int)
+    typeconstraint = widenconst(argextype(val, compact))
+    if isa(val, Union{OldSSAValue, SSAValue})
+        val, typeconstraint = simple_walk_constraint(compact, val, typeconstraint)
+    end
+    isa(typeconstraint, Union) || return # bail out if there won't be a good chance for lifting
+    leaves, visited_phinodes = collect_leaves(compact, val, typeconstraint)
     length(leaves) ≤ 1 && return # bail out if we don't have multiple leaves
 
-    # Let's check if we evaluate the comparison for each one of the leaves
+    # check if we can evaluate the comparison for each one of the leaves
+    cmp = argextype(target, compact)
     lifted_leaves = nothing
     for leaf in leaves
-        r = egal_tfunc(argextype(leaf, compact), cmp)
-        if isa(r, Const)
+        result = tfunc(argextype(leaf, compact), cmp)
+        if isa(result, Const)
             if lifted_leaves === nothing
                 lifted_leaves = LiftedLeaves()
             end
-            lifted_leaves[leaf] = LiftedValue(r.val)
+            lifted_leaves[leaf] = LiftedValue(result.val)
         else
-            return # TODO In some cases it might be profitable to hoist the === here
+            return # TODO In some cases it might be profitable to hoist the comparison here
         end
     end
 
+    # perform lifting
     lifted_val = perform_lifting!(compact,
         visited_phinodes, cmp, lifting_cache, Bool,
         lifted_leaves::LiftedLeaves, val)::LiftedValue
@@ -715,10 +739,14 @@ function sroa_pass!(ir::IRCode)
             canonicalize_typeassert!(compact, idx, stmt)
             continue
         elseif is_known_call(stmt, (===), compact)
-            lift_comparison!(compact, idx, stmt, lifting_cache)
+            lift_comparison!(===, compact, idx, stmt, lifting_cache)
+            continue
+        elseif is_known_call(stmt, isa, compact)
+            lift_comparison!(isa, compact, idx, stmt, lifting_cache)
+            continue
+        elseif is_known_call(stmt, isdefined, compact)
+            lift_comparison!(isdefined, compact, idx, stmt, lifting_cache)
             continue
-        # elseif is_known_call(stmt, isa, compact)
-            # TODO do a similar optimization as `lift_comparison!` for `===`
         else
             continue
         end

test/compiler/irpasses.jl

@@ -4,7 +4,34 @@ using Test
 using Base.Meta
 using Core: PhiNode, SSAValue, GotoNode, PiNode, QuoteNode, ReturnNode, GotoIfNot
 
-# Tests for domsort
+# utilities
+# =========
+
+import Core.Compiler: argextype, singleton_type
+
+argextype(@nospecialize args...) = argextype(args..., Any[])
+code_typed1(args...; kwargs...) = first(only(code_typed(args...; kwargs...)))::Core.CodeInfo
+get_code(args...; kwargs...) = code_typed1(args...; kwargs...).code
+
+# check if `x` is a statement with a given `head`
+isnew(@nospecialize x) = Meta.isexpr(x, :new)
+
+# check if `x` is a dynamic call of a given function
+iscall(y) = @nospecialize(x) -> iscall(y, x)
+function iscall((src, f)::Tuple{Core.CodeInfo,Base.Callable}, @nospecialize(x))
+    return iscall(x) do @nospecialize x
+        singleton_type(argextype(x, src)) === f
+    end
+end
+iscall(pred::Base.Callable, @nospecialize(x)) = Meta.isexpr(x, :call) && pred(x.args[1])
+
+# check if `x` is a statically-resolved call of a function whose name is `sym`
+isinvoke(y) = @nospecialize(x) -> isinvoke(y, x)
+isinvoke(sym::Symbol, @nospecialize(x)) = isinvoke(mi->mi.def.name===sym, x)
+isinvoke(pred::Function, @nospecialize(x)) = Meta.isexpr(x, :invoke) && pred(x.args[1]::Core.MethodInstance)
+
+# domsort
+# =======
 
 ## Test that domsort doesn't mangle single-argument phis (#29262)
 let m = Meta.@lower 1 + 1
@@ -67,25 +94,8 @@ let m = Meta.@lower 1 + 1
     Core.Compiler.verify_ir(ir)
 end
 
-# Tests for SROA
-
-import Core.Compiler: argextype, singleton_type
-const EMPTY_SPTYPES = Any[]
-
-code_typed1(args...; kwargs...) = first(only(code_typed(args...; kwargs...)))::Core.CodeInfo
-get_code(args...; kwargs...) = code_typed1(args...; kwargs...).code
-
-# check if `x` is a statement with a given `head`
-isnew(@nospecialize x) = Meta.isexpr(x, :new)
-
-# check if `x` is a dynamic call of a given function
-iscall(y) = @nospecialize(x) -> iscall(y, x)
-function iscall((src, f)::Tuple{Core.CodeInfo,Function}, @nospecialize(x))
-    return iscall(x) do @nospecialize x
-        singleton_type(argextype(x, src, EMPTY_SPTYPES)) === f
-    end
-end
-iscall(pred::Function, @nospecialize(x)) = Meta.isexpr(x, :call) && pred(x.args[1])
+# SROA
+# ====
 
 struct ImmutableXYZ; x; y; z; end
 mutable struct MutableXYZ; x; y; z; end
@@ -277,6 +287,38 @@ let src = code_typed1((Any,Any,Any)) do x, y, z
     @test_broken !any(isnew, src.code)
 end
 
+let # should work with constant globals
+    # immutable case
+    # --------------
+    src = @eval Module() begin
+        const REF_FLD = :x
+        struct ImmutableRef{T}
+            x::T
+        end
+
+        code_typed((Int,)) do x
+            r = ImmutableRef{Int}(x) # should be eliminated
+            x = getfield(r, REF_FLD) # should be eliminated
+            return sin(x)
+        end |> only |> first
+    end
+    @test count(iscall((src, getfield)), src.code) == 0
+    @test count(isnew, src.code) == 0
+
+    # mutable case
+    # ------------
+    src = @eval Module() begin
+        const REF_FLD = :x
+        code_typed() do
+            r = Ref{Int}(42) # should be eliminated
+            x = getfield(r, REF_FLD) # should be eliminated
+            return sin(x)
+        end |> only |> first
+    end
+    @test count(iscall((src, getfield)), src.code) == 0
+    @test count(isnew, src.code) == 0
+end
+
 # should work nicely with inlining to optimize away a complicated case
 # adapted from http://wiki.luajit.org/Allocation-Sinking-Optimization#implementation%5B
 struct Point
@@ -296,6 +338,75 @@ let src = code_typed1(compute_points)
     @test !any(isnew, src.code)
 end
 
+# comparison lifting
+# ==================
+
+let # lifting `===`
+    src = code_typed1((Bool,Int,)) do c, x
+        y = c ? x : nothing
+        y === nothing # => ϕ(false, true)
+    end
+    @test count(iscall((src, ===)), src.code) == 0
+
+    # should optimize away the iteration protocol
+    src = code_typed1((Int,)) do n
+        s = 0
+        for i in 1:n
+            s += i
+        end
+        s
+    end
+    @test !any(src.code) do @nospecialize x
+        iscall((src, ===), x) && argextype(x.args[2], src) isa Union
+    end
+end
+
+let # lifting `isa`
+    src = code_typed1((Bool,Int,)) do c, x
+        y = c ? x : nothing
+        isa(y, Int) # => ϕ(true, false)
+    end
+    @test count(iscall((src, isa)), src.code) == 0
+
+    src = code_typed1((Int,)) do n
+        s = 0
+        itr = 1:n
+        st = iterate(itr)
+        while !isa(st, Nothing)
+            i, st = itr
+            s += i
+            st = iterate(itr, st)
+        end
+        s
+    end
+    @test !any(src.code) do @nospecialize x
+        iscall((src, isa), x) && argextype(x.args[2], src) isa Union
+    end
+end
+
+let # lifting `isdefined`
+    src = code_typed1((Bool,Some{Int},)) do c, x
+        y = c ? x : nothing
+        isdefined(y, 1) # => ϕ(true, false)
+    end
+    @test count(iscall((src, isdefined)), src.code) == 0
+
+    src = code_typed1((Int,)) do n
+        s = 0
+        itr = 1:n
+        st = iterate(itr)
+        while isdefined(st, 2)
+            i, st = itr
+            s += i
+            st = iterate(itr, st)
+        end
+        s
+    end
+    @test !any(src.code) do @nospecialize x
+        iscall((src, isdefined), x) && argextype(x.args[2], src) isa Union
+    end
+end
+
 mutable struct Foo30594; x::Float64; end
 Base.copy(x::Foo30594) = Foo30594(x.x)
 function add!(p::Foo30594, off::Foo30594)
@@ -611,48 +722,6 @@ exc39508 = ErrorException("expected")
 end
 @test test39508() === exc39508
 
-let # `sroa_pass!` should work with constant globals
-    # immutable pass
-    src = @eval Module() begin
-        const REF_FLD = :x
-        struct ImmutableRef{T}
-            x::T
-        end
-
-        code_typed((Int,)) do x
-            r = ImmutableRef{Int}(x) # should be eliminated
-            x = getfield(r, REF_FLD) # should be eliminated
-            return sin(x)
-        end |> only |> first
-    end
-    @test !any(src.code) do @nospecialize(stmt)
-        Meta.isexpr(stmt, :call) || return false
-        ft = Core.Compiler.argextype(stmt.args[1], src, EMPTY_SPTYPES)
-        return Core.Compiler.widenconst(ft) == typeof(getfield)
-    end
-    @test !any(src.code) do @nospecialize(stmt)
-        return Meta.isexpr(stmt, :new)
-    end
-
-    # mutable pass
-    src = @eval Module() begin
-        const REF_FLD = :x
-        code_typed() do
-            r = Ref{Int}(42) # should be eliminated
-            x = getfield(r, REF_FLD) # should be eliminated
-            return sin(x)
-        end |> only |> first
-    end
-    @test !any(src.code) do @nospecialize(stmt)
-        Meta.isexpr(stmt, :call) || return false
-        ft = Core.Compiler.argextype(stmt.args[1], src, EMPTY_SPTYPES)
-        return Core.Compiler.widenconst(ft) == typeof(getfield)
-    end
-    @test !any(src.code) do @nospecialize(stmt)
-        return Meta.isexpr(stmt, :new)
-    end
-end
-
 let
     # `typeassert` elimination after SROA
     # NOTE we can remove this optimization once inference is able to reason about memory-effects
@@ -666,11 +735,7 @@ let
         end |> only |> first
     end
     # eliminate `typeassert(x2.x, Foo)`
-    @test all(src.code) do @nospecialize stmt
-        Meta.isexpr(stmt, :call) || return true
-        ft = Core.Compiler.argextype(stmt.args[1], src, EMPTY_SPTYPES)
-        return Core.Compiler.widenconst(ft) !== typeof(typeassert)
-    end
+    @test count(iscall((src, typeassert)), src.code) == 0
 end
 
 let