[NewOptimizer] Add back getfield elim pass

This adds back a (fixed version of) the NewOptimizer getfield elim
pass. This passes all getfield elim (and `@allocated`) tests in base,
with the exception of two tests that check for getfield elim on mutables.
That part is planned in the design of this pass, but will be part of a
follow up PR.

Author: Keno

base/compiler/ssair/driver.jl

@@ -129,9 +129,10 @@ function run_passes(ci::CodeInfo, nargs::Int, linetable::Vector{LineInfoNode})
             IRCode(code, lines, cfg, argtypes, mod, meta)
         end
     ir = construct_ssa!(ci, ir, domtree, defuse_insts, nargs)
-    domtree = construct_domtree(ir.cfg)
     ir = compact!(ir)
     verify_ir(ir)
+    ir = getfield_elim_pass!(ir)
+    ir = compact!(ir)
     ir = type_lift_pass!(ir)
     ir = compact!(ir)
     verify_ir(ir)

base/compiler/ssair/ir.jl

@@ -355,6 +355,10 @@ function getindex(compact::IncrementalCompact, idx)
     end
 end
 
+function getindex(view::TypesView, v::OldSSAValue)
+    return view.ir.ir.types[v.id]
+end
+
 function setindex!(compact::IncrementalCompact, v, idx)
     if idx < compact.result_idx
         # Kill count for current uses
@@ -372,8 +376,8 @@ end
 
 function getindex(view::TypesView, idx)
     isa(idx, SSAValue) && (idx = idx.id)
-    if isa(view.ir, IncrementalCompact) && idx < view.compact.result_idx
-        return view.compact.result_types[idx]
+    if isa(view.ir, IncrementalCompact) && idx < view.ir.result_idx
+        return view.ir.result_types[idx]
     else
         ir = isa(view.ir, IncrementalCompact) ? view.ir.ir : view.ir
         if idx <= length(ir.types)
@@ -500,7 +504,8 @@ function next(compact::IncrementalCompact, (idx, active_bb, old_result_idx)::Tup
 end
 
 function maybe_erase_unused!(extra_worklist, compact, idx)
-   if stmt_effect_free(compact.result[idx], compact.ir, compact.ir.mod)
+    effect_free = stmt_effect_free(compact.result[idx], compact.ir, compact.ir.mod)
+    if effect_free
         for ops in userefs(compact.result[idx])
             val = ops[]
             if isa(val, SSAValue)

base/compiler/ssair/passes.jl

@@ -1,3 +1,110 @@
+function compact_exprtype(compact, value)
+    if isa(value, Union{SSAValue, OldSSAValue})
+        return types(compact)[value]
+    elseif isa(value, Argument)
+        return compact.ir.argtypes[value.n]
+    end
+    exprtype(value, compact.ir, compact.ir.mod)
+end
+
+function getfield_elim_pass!(ir::IRCode)
+    compact = IncrementalCompact(ir)
+    insertions = Vector{Any}()
+    for (idx, stmt) in compact
+        isa(stmt, Expr) || continue
+        is_known_call(stmt, getfield, ir, ir.mod) || continue
+        isa(stmt.args[2], SSAValue) || continue
+        field = stmt.args[3]
+        isa(field, QuoteNode) && (field = field.value)
+        isa(field, Union{Int, Symbol}) || continue
+        orig_defidx = defidx = stmt.args[2].id
+        def = compact[defidx]
+        typeconstraint = types(compact)[defidx]
+        phi_locs = Tuple{Int, Int}[]
+        while true
+            if isa(def, PiNode)
+                typeconstraint = typeintersect(typeconstraint, def.typ)
+                if isa(def.val, SSAValue)
+                    defidx = def.val.id
+                    def = compact[defidx]
+                else
+                    def = def.val
+                end
+                continue
+            elseif isa(def, PhiNode)
+                possible_predecessors = collect(Iterators.filter(1:length(def.edges)) do n
+                    isassigned(def.values, n) || return false
+                    value = def.values[n]
+                    edge_typ = compact_exprtype(compact, value)
+                    return edge_typ ⊑ typeconstraint
+                end)
+                # For now, only look at unique predecessors
+                if length(possible_predecessors) == 1
+                    n = possible_predecessors[1]
+                    pred = def.edges[n]
+                    val = def.values[n]
+                    if isa(val, SSAValue)
+                        push!(phi_locs, (pred, defidx))
+                        defidx = val.id
+                        def = compact[defidx]
+                    elseif def == val
+                        # This shouldn't really ever happen, but
+                        # patterns like this can occur in dead code,
+                        # so bail out.
+                        break
+                    else
+                        def = val
+                    end
+                    continue
+                end
+            end
+            break
+        end
+        if isa(def, Expr) && is_known_call(def, tuple, ir, ir.mod) && isa(field, Int) && 1 <= field < length(def.args)
+            forwarded = def.args[1+field]
+        elseif isexpr(def, :new)
+            typ = def.typ
+            if isa(typ, UnionAll)
+                typ = unwrap_unionall(typ)
+            end
+            isa(typ, DataType) || continue
+            !typ.mutable || continue
+            if isa(field, Symbol)
+                field = fieldindex(typ, field, false)
+                field == 0 && continue
+            elseif isa(field, Integer)
+                (1 <= field <= fieldcount(typ)) || continue
+            end
+            forwarded = def.args[1+field]
+        else
+            continue
+        end
+        if !isempty(phi_locs) && isa(forwarded, SSAValue)
+            # TODO: We have have to use BB ids for phi_locs
+            # to avoid index invalidation.
+            push!(insertions, (idx, phi_locs))
+        end
+        compact[idx] = forwarded
+    end
+    ir = finish(compact)
+    for (idx, phi_locs) in insertions
+        # For non-dominating load-store forward, we may have to insert extra phi nodes
+        # TODO: Can use the domtree to eliminate unnecessary phis, but ok for now
+        forwarded = ir.stmts[idx]
+        if isa(forwarded, SSAValue)
+            forwarded_typ = ir.types[forwarded.id]
+            for (pred, pos) in reverse!(phi_locs)
+                node = PhiNode()
+                push!(node.edges, pred)
+                push!(node.values, forwarded)
+                forwarded = insert_node!(ir, pos, forwarded_typ, node)
+            end
+        end
+        ir.stmts[idx] = forwarded
+    end
+    ir
+end
+
 function type_lift_pass!(ir::IRCode)
     type_ctx_uses = Vector{Vector{Int}}[]
     has_non_type_ctx_uses = IdSet{Int}()
@@ -81,4 +188,4 @@ function type_lift_pass!(ir::IRCode)
         end
     end
     ir
-end
+end

base/compiler/utilities.jl

@@ -181,6 +181,8 @@ function exprtype(@nospecialize(x), src, mod::Module)
         return (x::TypedSlot).typ
     elseif isa(x, SSAValue)
         return abstract_eval_ssavalue(x::SSAValue, src)
+    elseif isa(x, Argument)
+        return src.argtypes[x.n]
     elseif isa(x, Symbol)
         return abstract_eval_global(mod, x::Symbol)
     elseif isa(x, QuoteNode)