ir: Set refined when narrowing PhiNode in IncrementalCompact (#54061)

This is a bit of a tricky one. What happens here is that
IncrementalCompact finds that one branch of an if/else is dead, so one
of the incoming values of the PhiNode goes away. This may have the
effect of narrowing the type of the PhiNode (if the branches have
different types). In the latest iteration of our compiler, such
situations need to be annotated with IR_FLAG_REFINED or `Refined()`,
otherwise irinterp will skip them on revisit.

In theory, the fix is quite simple: Check whether the type is being
refined and if so, set the flag. However, this code sits inside
IncrementalCompcat, which currently neither performs lattice operations
nor sets any Refined flags by itself. The three possible options here
are:

1. Thread lattice through IncrementalCompact.
2. Use an egal check inside IncrementalCompact rather than the proper
lattice query. This is legal, but may overapproximate the need for
`Refined`, thus causing unnecessary work later.
3. Move the phi node narrowing outside of IncrementalCompact.

This PR takes a hybrid approach of 2 and 3. Approach 1 is undesirable,
because we do not want to recompile all of IncrementalCompact for every
lattice, just for this one query. I almost took approach 2, but it is
unsatisfying, because of the imprecision, which could cuase a lot of
extra work in pathological cases.

The primary downside of approach 3 is that now every pass that performs
IncrementalCompact with CFG manipulation enabled will need to call the
`reprocess_phi_node!` helper. We have two of these in base (adce_pass!
and `compact!(..., true)` although only the former as access to a
lattice.

The compromise I came up with here is to do the egality check, except
leaving the PhiNode in place when it fails rather than setting `Refined`
internally. I think this is ok. There's a bit of an issue with duplicate
work, because the first thing that the lattice check will do is check
for equality, but that only happens in the case where a refinement
actually is required, which is both rarer and profitable.

Finally, I'm pretty sure I've found this issue before, but I cannot find
any writeup on it, so if somebody remembers me writing it up somewhere,
please let me know so I can add appropriate cross references.

---------

Co-authored-by: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>

Author: Keno

base/compiler/ssair/ir.jl

@@ -1601,17 +1601,15 @@ function process_node!(compact::IncrementalCompact, result_idx::Int, inst::Instr
         end
 
         values = process_phinode_values(values, late_fixup, already_inserted_phi_arg, result_idx, ssa_rename, used_ssas, new_new_used_ssas, do_rename_ssa, mark_refined!)
-        # Don't remove the phi node if it is before the definition of its value
-        # because doing so can create forward references. This should only
-        # happen with dead loops, but can cause problems when optimization
-        # passes look at all code, dead or not. This check should be
-        # unnecessary when DCE can remove those dead loops entirely, so this is
-        # just to be safe.
-        before_def = isassigned(values, 1) && (v = values[1]; isa(v, OldSSAValue)) && idx < v.id
-        if length(edges) == 1 && isassigned(values, 1) && !before_def &&
-                length(cfg_transforms_enabled ?
-                    result_bbs[bb_rename_succ[active_bb]].preds :
-                    compact.ir.cfg.blocks[active_bb].preds) == 1
+
+        # Quick egality check for PhiNode that may be replaced with its incoming
+        # value without needing to set the `Refined` flag. We can't do the actual
+        # refinement check, because we do not have access to the lattice here.
+        # Users may call `reprocess_phi_node!` inside the compaction loop to
+        # revisit PhiNodes with the proper lattice refinement check.
+        if may_replace_phi(values, cfg_transforms_enabled ?
+                result_bbs[bb_rename_succ[active_bb]] :
+                compact.ir.cfg.blocks[active_bb], idx) && argextype(values[1], compact) === inst[:type]
             # There's only one predecessor left - just replace it
             v = values[1]
             @assert !isa(v, NewSSAValue)
@@ -1662,6 +1660,38 @@ function process_node!(compact::IncrementalCompact, result_idx::Int, inst::Instr
     return result_idx
 end
 
+function may_replace_phi(values::Vector{Any}, phi_bb::BasicBlock, idx::Int)
+    length(values) == 1 || return false
+    isassigned(values, 1) || return false
+    length(phi_bb.preds) == 1 || return false
+
+    # Don't remove the phi node if it is before the definition of its value
+    # because doing so can create forward references. This should only
+    # happen with dead loops, but can cause problems when optimization
+    # passes look at all code, dead or not. This check should be
+    # unnecessary when DCE can remove those dead loops entirely, so this is
+    # just to be safe.
+    v = values[1]
+    before_def = isa(v, OldSSAValue) && idx < v.id
+    return !before_def
+end
+
+function reprocess_phi_node!(𝕃ₒ::AbstractLattice, compact::IncrementalCompact, phi::PhiNode, old_idx::Int)
+    phi_bb = compact.active_result_bb
+    did_just_finish_bb(compact) && (phi_bb -= 1)
+    may_replace_phi(phi.values, compact.cfg_transform.result_bbs[phi_bb], compact.idx) || return false
+
+    # There's only one predecessor left - just replace it
+    v = phi.values[1]
+    if !⊑(𝕃ₒ, compact[compact.ssa_rename[old_idx]][:type], argextype(v, compact))
+        v = Refined(v)
+    end
+    compact.ssa_rename[old_idx] = v
+
+    delete_inst_here!(compact)
+    return true
+end
+
 function resize!(compact::IncrementalCompact, nnewnodes::Int)
     old_length = length(compact.result)
     resize!(compact.result, nnewnodes)

base/compiler/ssair/passes.jl

@@ -2009,8 +2009,13 @@ function adce_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
     unionphis = Pair{Int,Any}[] # sorted
     compact = IncrementalCompact(ir, true)
     made_changes = false
-    for ((_, idx), stmt) in compact
+    for ((old_idx, idx), stmt) in compact
         if isa(stmt, PhiNode)
+            if reprocess_phi_node!(𝕃ₒ, compact, stmt, old_idx)
+                # Phi node has a single predecessor and was deleted
+                made_changes = true
+                continue
+            end
             push!(all_phis, idx)
             if is_some_union(compact.result[idx][:type])
                 push!(unionphis, Pair{Int,Any}(idx, Union{}))

test/compiler/irpasses.jl

@@ -1820,3 +1820,19 @@ function f53521()
     end
 end
 @test code_typed(f53521)[1][2] === Nothing
+
+# Test that adce_pass! sets Refined on PhiNode values
+let code = Any[
+    # Basic Block 1
+    GotoIfNot(false, 3)
+    # Basic Block 2
+    nothing
+    # Basic Block 3
+    PhiNode(Int32[1, 2], Any[1.0, 1])
+    ReturnNode(Core.SSAValue(3))
+]
+    ir = make_ircode(code; ssavaluetypes=Any[Any, Nothing, Union{Int64, Float64}, Any])
+    (ir, made_changes) = Core.Compiler.adce_pass!(ir)
+    @test made_changes
+    @test (ir[Core.SSAValue(length(ir.stmts))][:flag] & Core.Compiler.IR_FLAG_REFINED) != 0
+end