From 89a613b5978ecf06566f5778253ca18c3e3fc2e7 Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Mon, 28 Mar 2022 12:30:01 +0900
Subject: [PATCH] optimizer: eliminate more `isdefined` checks (#44755)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Follows up #44708 -- in that PR I missed the most obvious optimization
opportunity, i.e. we can safely eliminate `isdefined` checks when all
fields are defined at allocation site.
This change allows us to eliminate capturing closure constructions when
the body and callsite of capture closure is available within a optimized
frame, e.g.:
```julia
function abmult(r::Int, x0)
    if r < 0
        r = -r
    end
    f = x -> x * r
    return @inline f(x0)
end
```
```diff
diff --git a/_master.jl b/_pr.jl
index ea06d865b75..c38f221090f 100644
--- a/_master.jl
+++ b/_pr.jl
@@ -1,24 +1,19 @@
 julia> @code_typed abmult(-3, 3)
 CodeInfo(
-1 ── %1  = Core.Box::Type{Core.Box}
-│    %2  = %new(%1, r@_2)::Core.Box
-│    %3  = Core.isdefined(%2, :contents)::Bool
-└───       goto #3 if not %3
+1 ──       goto #3 if not true
 2 ──       goto #4
 3 ──       $(Expr(:throw_undef_if_not, :r, false))::Any
-4 ┄─ %7  = (r@_2 < 0)::Any
-└───       goto #9 if not %7
-5 ── %9  = Core.isdefined(%2, :contents)::Bool
-└───       goto #7 if not %9
+4 ┄─ %4  = (r@_2 < 0)::Any
+└───       goto #9 if not %4
+5 ──       goto #7 if not true
 6 ──       goto #8
 7 ──       $(Expr(:throw_undef_if_not, :r, false))::Any
-8 ┄─ %13 = -r@_2::Any
-9 ┄─ %14 = φ (#4 => r@_2, #8 => %13)::Any
-│    %15 = Core.isdefined(%2, :contents)::Bool
-└───       goto #11 if not %15
+8 ┄─ %9  = -r@_2::Any
+9 ┄─ %10 = φ (#4 => r@_2, #8 => %9)::Any
+└───       goto #11 if not true
 10 ─       goto #12
 11 ─       $(Expr(:throw_undef_if_not, :r, false))::Any
-12 ┄ %19 = (x0 * %14)::Any
+12 ┄ %14 = (x0 * %10)::Any
 └───       goto #13
-13 ─       return %19
+13 ─       return %14
 ) => Any
```
---
 base/compiler/ssair/passes.jl | 22 ++++++++++------------
 test/compiler/irpasses.jl     | 25 +++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
index bf2e6d3aae17f..9a0468e8e36ea 100644
--- a/base/compiler/ssair/passes.jl
+++ b/base/compiler/ssair/passes.jl
@@ -372,10 +372,7 @@ function lift_leaves(compact::IncrementalCompact,
                 lift_arg!(compact, leaf, cache_key, def, 1+field, lifted_leaves)
                 continue
             elseif isexpr(def, :new)
-                typ = widenconst(types(compact)[leaf])
-                if isa(typ, UnionAll)
-                    typ = unwrap_unionall(typ)
-                end
+                typ = unwrap_unionall(widenconst(types(compact)[leaf]))
                 (isa(typ, DataType) && !isabstracttype(typ)) || return nothing
                 @assert !ismutabletype(typ)
                 if length(def.args) < 1+field
@@ -786,10 +783,7 @@ function sroa_pass!(ir::IRCode)
                         push!(preserved, preserved_arg.id)
                         continue
                     elseif isexpr(def, :new)
-                        typ = widenconst(argextype(SSAValue(defidx), compact))
-                        if isa(typ, UnionAll)
-                            typ = unwrap_unionall(typ)
-                        end
+                        typ = unwrap_unionall(widenconst(argextype(SSAValue(defidx), compact)))
                         if typ isa DataType && !ismutabletype(typ)
                             record_immutable_preserve!(new_preserves, def, compact)
                             push!(preserved, preserved_arg.id)
@@ -948,10 +942,7 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse
         defexpr = ir[SSAValue(idx)][:inst]
         isexpr(defexpr, :new) || continue
         newidx = idx
-        typ = ir.stmts[newidx][:type]
-        if isa(typ, UnionAll)
-            typ = unwrap_unionall(typ)
-        end
+        typ = unwrap_unionall(ir.stmts[newidx][:type])
         # Could still end up here if we tried to setfield! on an immutable, which would
         # error at runtime, but is not illegal to have in the IR.
         ismutabletype(typ) || continue
@@ -1023,6 +1014,13 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse
                     end
                     has_safe_def(ir, get_domtree(), allblocks, du, newidx, use.idx) || @goto skip
                 end
+            else # always have some definition at the allocation site
+                for i = 1:length(du.uses)
+                    use = du.uses[i]
+                    if use.kind === :isdefined
+                        ir[SSAValue(use.idx)][:inst] = true
+                    end
+                end
             end
         end
         # Everything accounted for. Go field by field and perform idf:
diff --git a/test/compiler/irpasses.jl b/test/compiler/irpasses.jl
index fedc622b91148..6c77891bede5a 100644
--- a/test/compiler/irpasses.jl
+++ b/test/compiler/irpasses.jl
@@ -493,6 +493,31 @@ let src = code_typed1(isdefined_elim)
 end
 @test isdefined_elim() == Any[]
 
+function abmult(r::Int, x0)
+    if r < 0
+        r = -r
+    end
+    f = x -> x * r
+    return @inline f(x0)
+end
+let src = code_typed1(abmult, (Int,Int))
+    @test is_scalar_replaced(src)
+end
+@test abmult(-3, 3) == 9
+
+function abmult2(r0::Int, x0)
+    r::Int = r0
+    if r < 0
+        r = -r
+    end
+    f = x -> x * r
+    return f(x0)
+end
+let src = code_typed1(abmult2, (Int,Int))
+    @test is_scalar_replaced(src)
+end
+@test abmult2(-3, 3) == 9
+
 # comparison lifting
 # ==================