From 0435e02b7468c69a8ec48d882b5890b247988166 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Thu, 5 Mar 2020 04:29:30 -0600
Subject: [PATCH] RFC: support external compiler passes

---
 base/compiler/optimize.jl | 24 ++++++++++++++++++++++++
 test/compiler/irpasses.jl | 17 +++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/base/compiler/optimize.jl b/base/compiler/optimize.jl
index 56b602c9356ae..0fbe6ddcadd46 100644
--- a/base/compiler/optimize.jl
+++ b/base/compiler/optimize.jl
@@ -166,6 +166,30 @@ end
 function optimize(opt::OptimizationState, @nospecialize(result))
     def = opt.linfo.def
     nargs = Int(opt.nargs) - 1
+    # Execute external passes
+    modified = true
+    while modified
+        modified = false
+        for i = 1:length(opt.src.code)
+            stmt = opt.src.code[i]
+            if isexpr(stmt, :meta)
+                for j = 1:length(stmt.args)
+                    a = stmt.args[j]
+                    if isexpr(a, :external_pass)
+                        f = a.args[1]
+                        if f === :start
+                            f = a.args[2]
+                        end
+                        if isa(f, Function)
+                            Core._apply_latest(f, (opt, i))
+                            modified = true
+                            break
+                        end
+                    end
+                end
+            end
+        end
+    end
     @timeit "optimizer" ir = run_passes(opt.src, nargs, opt)
     force_noinline = _any(@nospecialize(x) -> isexpr(x, :meta) && x.args[1] === :noinline, ir.meta)
 
diff --git a/test/compiler/irpasses.jl b/test/compiler/irpasses.jl
index 042d91d2ff700..2f34c533be89c 100644
--- a/test/compiler/irpasses.jl
+++ b/test/compiler/irpasses.jl
@@ -295,3 +295,20 @@ end
 const _some_coeffs = (1,[2],3,4)
 splat_from_globalref(x) = (x, _some_coeffs...,)
 @test splat_from_globalref(0) == (0, 1, [2], 3, 4)
+
+# External passes
+const _opt = Ref{Any}(nothing)
+avx_pass(opt) = (_opt[] = opt; opt)
+macro avx(ex)
+    esc(Expr(:block, Expr(:meta, Expr(:external_pass, :start, avx_pass)), ex, Expr(:meta, Expr(:external_pass, :stop))))
+end
+function myselfdot(a)
+    s = zero(eltype(a))
+    @avx for item in a
+        s += item^2
+    end
+    return s
+end
+a = [0.1, 0.2, 0.3]
+@test myselfdot(a) ≈ 0.14
+@test isa(_opt[], Core.Compiler.OptimizationState)