From a07cfa3cfc67591f0387b050602b22d135b0784d Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <aviatesk@gmail.com>
Date: Mon, 17 Jan 2022 17:45:11 +0900
Subject: [PATCH] run EA for SROA and check its accuracy

---
 base/compiler/optimize.jl                     | 9 ++-------
 base/compiler/ssair/EscapeAnalysis/EAUtils.jl | 6 ++----
 base/compiler/ssair/passes.jl                 | 8 +++++---
 test/compiler/irpasses.jl                     | 2 +-
 4 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/base/compiler/optimize.jl b/base/compiler/optimize.jl
index ba4ca7d5bf685..e921e5fbed267 100644
--- a/base/compiler/optimize.jl
+++ b/base/compiler/optimize.jl
@@ -517,16 +517,11 @@ function run_passes(ci::CodeInfo, sv::OptimizationState)
     @timeit "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     # @timeit "verify 2" verify_ir(ir)
     @timeit "compact 2" ir = compact!(ir)
-    @timeit "SROA"      ir = sroa_pass!(ir)
+    nargs = let def = sv.linfo.def; isa(def, Method) ? Int(def.nargs) : 0; end
+    @timeit "SROA"      ir = sroa_pass!(ir, nargs)
     @timeit "ADCE"      ir = adce_pass!(ir)
     @timeit "type lift" ir = type_lift_pass!(ir)
     @timeit "compact 3" ir = compact!(ir)
-    # TODO validate EA by comparing it to the existing SROA
-    nargs = let def = sv.linfo.def
-        isa(def, Method) ? Int(def.nargs) : 0
-    end
-    @timeit "EA" estate = analyze_escapes(ir, nargs)
-    cache_escapes!(sv.linfo, estate, ir)
     if JLOptions().debug_level == 2
         @timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))
     end
diff --git a/base/compiler/ssair/EscapeAnalysis/EAUtils.jl b/base/compiler/ssair/EscapeAnalysis/EAUtils.jl
index 78ebf5eff206f..fb4499d01452a 100644
--- a/base/compiler/ssair/EscapeAnalysis/EAUtils.jl
+++ b/base/compiler/ssair/EscapeAnalysis/EAUtils.jl
@@ -252,9 +252,7 @@ function run_passes_with_ea(interp::EscapeAnalyzer, ci::CodeInfo, sv::Optimizati
     @timeit "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     # @timeit "verify 2" verify_ir(ir)
     @timeit "compact 2" ir = compact!(ir)
-    nargs = let def = sv.linfo.def
-        isa(def, Method) ? Int(def.nargs) : 0
-    end
+    nargs = let def = sv.linfo.def; isa(def, Method) ? Int(def.nargs) : 0; end
     local state
     try
         @timeit "collect escape information" state = analyze_escapes(ir, nargs)
@@ -270,7 +268,7 @@ function run_passes_with_ea(interp::EscapeAnalyzer, ci::CodeInfo, sv::Optimizati
     interp.ir = cacheir
     interp.state = state
     interp.linfo = sv.linfo
-    @timeit "SROA"      ir = sroa_pass!(ir)
+    @timeit "SROA"      ir = sroa_pass!(ir, nargs)
     @timeit "ADCE"      ir = adce_pass!(ir)
     @timeit "type lift" ir = type_lift_pass!(ir)
     @timeit "compact 3" ir = compact!(ir)
diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
index 1df68da4e2b3d..e7be739426c31 100644
--- a/base/compiler/ssair/passes.jl
+++ b/base/compiler/ssair/passes.jl
@@ -641,7 +641,7 @@ its argument).
 In a case when all usages are fully eliminated, `struct` allocation may also be erased as
 a result of succeeding dead code elimination.
 """
-function sroa_pass!(ir::IRCode)
+function sroa_pass!(ir::IRCode, nargs::Int)
     compact = IncrementalCompact(ir)
     defuses = nothing # will be initialized once we encounter mutability in order to reduce dynamic allocations
     lifting_cache = IdDict{Pair{AnySSAValue, Any}, AnySSAValue}()
@@ -813,7 +813,7 @@ function sroa_pass!(ir::IRCode)
         used_ssas = copy(compact.used_ssas)
         simple_dce!(compact, (x::SSAValue) -> used_ssas[x.id] -= 1)
         ir = complete(compact)
-        sroa_mutables!(ir, defuses, used_ssas)
+        sroa_mutables!(ir, defuses, used_ssas, nargs)
         return ir
     else
         simple_dce!(compact)
@@ -821,9 +821,10 @@ function sroa_pass!(ir::IRCode)
     end
 end
 
-function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse}}, used_ssas::Vector{Int})
+function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse}}, used_ssas::Vector{Int}, nargs::Int)
     # initialization of domtree is delayed to avoid the expensive computation in many cases
     local domtree = nothing
+    estate = analyze_escapes(ir, nargs)
     for (idx, (intermediaries, defuse)) in defuses
         intermediaries = collect(intermediaries)
         # Check if there are any uses we did not account for. If so, the variable
@@ -899,6 +900,7 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int, Tuple{SPCSet, SSADefUse
                 end
             end
         end
+        is_sroa_eligible(estate[SSAValue(idx)]) || println("[EA] bad EA: ", ir.argtypes[1:nargs], " at ", idx)
         # Everything accounted for. Go field by field and perform idf:
         # Compute domtree now, needed below, now that we have finished compacting the IR.
         # This needs to be after we iterate through the IR with `IncrementalCompact`
diff --git a/test/compiler/irpasses.jl b/test/compiler/irpasses.jl
index f355a42d7b08c..cbbf4375541e1 100644
--- a/test/compiler/irpasses.jl
+++ b/test/compiler/irpasses.jl
@@ -407,7 +407,7 @@ let m = Meta.@lower 1 + 1
     src.ssaflags = fill(Int32(0), nstmts)
     ir = Core.Compiler.inflate_ir(src, Any[], Any[Any, Any])
     @test Core.Compiler.verify_ir(ir) === nothing
-    ir = @test_nowarn Core.Compiler.sroa_pass!(ir)
+    ir = @test_nowarn Core.Compiler.sroa_pass!(ir, 0)
     @test Core.Compiler.verify_ir(ir) === nothing
 end