From e69f87db162d19b7aa4c60c6a58a1cf8edd44d4e Mon Sep 17 00:00:00 2001
From: Keno Fischer <keno@juliacomputing.com>
Date: Tue, 28 Nov 2023 21:40:01 +0000
Subject: [PATCH] sroa: Lift restriction that all_same optimization must give
 SSAValue

This restriction has been in there since this code was added in #44557.
Unfortunately, I can't tell from the commit history why this restriction
was added. It's possible that it was trying to avoid putting things into
statement position that were not allowed in the phi block, but we have
cleaned that up since (#50308 and related), so let's see if this restriction
is still required, since I was seeing some suboptimial optimization results
because of this.
---
 base/compiler/ssair/inlining.jl |  2 +-
 base/compiler/ssair/ir.jl       |  6 +++---
 base/compiler/ssair/passes.jl   | 32 ++++++++++++++++++--------------
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/base/compiler/ssair/inlining.jl b/base/compiler/ssair/inlining.jl
index 326a4094e4134..17d990e4b4aee 100644
--- a/base/compiler/ssair/inlining.jl
+++ b/base/compiler/ssair/inlining.jl
@@ -1918,7 +1918,7 @@ function ssa_substitute_op!(insert_node!::Inserter, subst_inst::Instruction, @no
             end
         end
     end
-    isa(val, Union{SSAValue, NewSSAValue}) && return val # avoid infinite loop
+    isa(val, Union{SSAValue, OldSSAValue, NewSSAValue}) && return val # avoid infinite loop
     urs = userefs(val)
     for op in urs
         op[] = ssa_substitute_op!(insert_node!, subst_inst, op[], ssa_substitute)
diff --git a/base/compiler/ssair/ir.jl b/base/compiler/ssair/ir.jl
index dc19c9a09fbba..651adf04f2d8a 100644
--- a/base/compiler/ssair/ir.jl
+++ b/base/compiler/ssair/ir.jl
@@ -464,7 +464,7 @@ struct UndefToken end; const UNDEF_TOKEN = UndefToken()
         isdefined(stmt, :val) || return OOB_TOKEN
         op == 1 || return OOB_TOKEN
         return stmt.val
-    elseif isa(stmt, Union{SSAValue, NewSSAValue, GlobalRef})
+    elseif isa(stmt, Union{SSAValue, OldSSAValue, NewSSAValue, GlobalRef})
         op == 1 || return OOB_TOKEN
         return stmt
     elseif isa(stmt, UpsilonNode)
@@ -520,7 +520,7 @@ end
     elseif isa(stmt, ReturnNode)
         op == 1 || throw(BoundsError())
         stmt = typeof(stmt)(v)
-    elseif isa(stmt, Union{SSAValue, NewSSAValue, GlobalRef})
+    elseif isa(stmt, Union{SSAValue, OldSSAValue, NewSSAValue, GlobalRef})
         op == 1 || throw(BoundsError())
         stmt = v
     elseif isa(stmt, UpsilonNode)
@@ -550,7 +550,7 @@ end
 
 function userefs(@nospecialize(x))
     relevant = (isa(x, Expr) && is_relevant_expr(x)) ||
-        isa(x, GotoIfNot) || isa(x, ReturnNode) || isa(x, SSAValue) || isa(x, NewSSAValue) ||
+        isa(x, GotoIfNot) || isa(x, ReturnNode) || isa(x, SSAValue) || isa(x, OldSSAValue) || isa(x, NewSSAValue) ||
         isa(x, PiNode) || isa(x, PhiNode) || isa(x, PhiCNode) || isa(x, UpsilonNode) || isa(x, EnterNode)
     return UseRefIterator(x, relevant)
 end
diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
index 944e444bf6d06..0a4710068c51b 100644
--- a/base/compiler/ssair/passes.jl
+++ b/base/compiler/ssair/passes.jl
@@ -678,24 +678,28 @@ function perform_lifting!(compact::IncrementalCompact,
         end
     end
 
-    the_leaf_val = isa(the_leaf, LiftedValue) ? the_leaf.val : nothing
-    if !isa(the_leaf_val, SSAValue)
-        all_same = false
-    end
-
-    if all_same
+    if all_same && isa(the_leaf, LiftedValue)
         dominates_all = true
-        if lazydomtree !== nothing
-            domtree = get!(lazydomtree)
-            for item in visited_philikes
-                if !dominates_ssa(compact, domtree, the_leaf_val, item)
-                    dominates_all = false
-                    break
+        the_leaf_val = the_leaf.val
+        if isa(the_leaf_val, Union{SSAValue, OldSSAValue})
+            if lazydomtree === nothing
+                # Must conservatively assume this
+                dominates_all = false
+            else
+                domtree = get!(lazydomtree)
+                for item in visited_philikes
+                    if !dominates_ssa(compact, domtree, the_leaf_val, item)
+                        dominates_all = false
+                        break
+                    end
                 end
             end
-            if dominates_all
-                return the_leaf
+        end
+        if dominates_all
+            if isa(the_leaf, OldSSAValue)
+                the_leaf = simple_walk(compact, the_leaf)
             end
+            return the_leaf
         end
     end