Add support for specifying the AA pipeline.

deps/LLVMExtra/include/LLVMExtra.h

@@ -254,6 +254,8 @@ void LLVMPassBuilderExtensionsRegisterFunctionPass(LLVMPassBuilderExtensionsRef
                                                    const char *PassName,
                                                    LLVMJuliaFunctionPassCallback Callback,
                                                    void *Thunk);
+void LLVMPassBuilderExtensionsSetAAPipeline(LLVMPassBuilderExtensionsRef Extensions,
+                                            const char *AAPipeline);
 LLVMErrorRef LLVMRunJuliaPasses(LLVMModuleRef M, const char *Passes,
                                 LLVMTargetMachineRef TM, LLVMPassBuilderOptionsRef Options,
                                 LLVMPassBuilderExtensionsRef Extensions);

deps/LLVMExtra/lib/NewPM.cpp

@@ -1,5 +1,6 @@
 #include "LLVMExtra.h"
 
+#include <llvm/Analysis/AliasAnalysis.h>
 #include <llvm/IR/Verifier.h>
 #include <llvm/Passes/PassBuilder.h>
 #include <llvm/Passes/StandardInstrumentations.h>
@@ -18,10 +19,6 @@ namespace llvm {
 // Keep this in sync with PassBuilderBindings.cpp!
 class LLVMPassBuilderOptions {
 public:
-  explicit LLVMPassBuilderOptions(bool DebugLogging = false, bool VerifyEach = false,
-                                  PipelineTuningOptions PTO = PipelineTuningOptions())
-      : DebugLogging(DebugLogging), VerifyEach(VerifyEach), PTO(PTO) {}
-
   bool DebugLogging;
   bool VerifyEach;
   PipelineTuningOptions PTO;
@@ -30,7 +27,13 @@ DEFINE_SIMPLE_CONVERSION_FUNCTIONS(LLVMPassBuilderOptions, LLVMPassBuilderOption
 
 class LLVMPassBuilderExtensions {
 public:
+  // A callback to register additional pipeline parsing callbacks with the pass builder.
+  // This is used to support Julia's passes.
   void (*RegistrationCallback)(void *);
+
+  // A list of callbacks that each register a single custom module or function pass.
+  // These callbacks are generated here in C++, and match against a pass name.
+  // This is used to enable custom LLVM passes implemented in Julia.
   SmallVector<std::function<bool(StringRef, ModulePassManager &,
                                  ArrayRef<PassBuilder::PipelineElement>)>,
               2>
@@ -39,6 +42,9 @@ class LLVMPassBuilderExtensions {
                                  ArrayRef<PassBuilder::PipelineElement>)>,
               2>
       FunctionPipelineParsingCallbacks;
+
+  // A pipeline describing the alias analysis passes to run.
+  const char *AAPipeline;
 };
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(LLVMPassBuilderExtensions, LLVMPassBuilderExtensionsRef)
 } // namespace llvm
@@ -120,6 +126,15 @@ void LLVMPassBuilderExtensionsRegisterFunctionPass(LLVMPassBuilderExtensionsRef
   return;
 }
 
+// Alias analysis pipeline
+
+void LLVMPassBuilderExtensionsSetAAPipeline(LLVMPassBuilderExtensionsRef Extensions,
+                                            const char *AAPipeline) {
+  LLVMPassBuilderExtensions *PassExts = unwrap(Extensions);
+  PassExts->AAPipeline = AAPipeline;
+  return;
+}
+
 
 // Vendored API entrypoint
 
@@ -154,6 +169,15 @@ LLVMErrorRef LLVMRunJuliaPasses(LLVMModuleRef M, const char *Passes,
   CGSCCAnalysisManager CGAM;
   ModuleAnalysisManager MAM;
   PB.registerLoopAnalyses(LAM);
+  if (PassExts->AAPipeline) {
+    // If we have a custom AA pipeline, we need to register it _before_ calling
+    // registerFunctionAnalyses, or the default alias analysis pipeline is used.
+    AAManager AA;
+    if (auto Err = PB.parseAAPipeline(AA, PassExts->AAPipeline)) {
+      return wrap(std::move(Err));
+    }
+    FAM.registerPass([&] { return std::move(AA); });
+  }
   PB.registerFunctionAnalyses(FAM);
   PB.registerCGSCCAnalyses(CGAM);
   PB.registerModuleAnalyses(MAM);

lib/15/libLLVM_extra.jl

@@ -377,6 +377,10 @@ function LLVMPassBuilderExtensionsRegisterFunctionPass(Options, PassName, Callba
     ccall((:LLVMPassBuilderExtensionsRegisterFunctionPass, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring, LLVMJuliaFunctionPassCallback, Ptr{Cvoid}), Options, PassName, Callback, Thunk)
 end
 
+function LLVMPassBuilderExtensionsSetAAPipeline(Extensions, AAPipeline)
+    ccall((:LLVMPassBuilderExtensionsSetAAPipeline, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring), Extensions, AAPipeline)
+end
+
 function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
     ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
 end

lib/16/libLLVM_extra.jl

@@ -377,6 +377,10 @@ function LLVMPassBuilderExtensionsRegisterFunctionPass(Options, PassName, Callba
     ccall((:LLVMPassBuilderExtensionsRegisterFunctionPass, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring, LLVMJuliaFunctionPassCallback, Ptr{Cvoid}), Options, PassName, Callback, Thunk)
 end
 
+function LLVMPassBuilderExtensionsSetAAPipeline(Extensions, AAPipeline)
+    ccall((:LLVMPassBuilderExtensionsSetAAPipeline, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring), Extensions, AAPipeline)
+end
+
 function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
     ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
 end

lib/17/libLLVM_extra.jl

@@ -337,6 +337,10 @@ function LLVMPassBuilderExtensionsRegisterFunctionPass(Options, PassName, Callba
     ccall((:LLVMPassBuilderExtensionsRegisterFunctionPass, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring, LLVMJuliaFunctionPassCallback, Ptr{Cvoid}), Options, PassName, Callback, Thunk)
 end
 
+function LLVMPassBuilderExtensionsSetAAPipeline(Extensions, AAPipeline)
+    ccall((:LLVMPassBuilderExtensionsSetAAPipeline, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring), Extensions, AAPipeline)
+end
+
 function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
     ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
 end

lib/18/libLLVM_extra.jl

@@ -271,6 +271,10 @@ function LLVMPassBuilderExtensionsRegisterFunctionPass(Options, PassName, Callba
     ccall((:LLVMPassBuilderExtensionsRegisterFunctionPass, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring, LLVMJuliaFunctionPassCallback, Ptr{Cvoid}), Options, PassName, Callback, Thunk)
 end
 
+function LLVMPassBuilderExtensionsSetAAPipeline(Extensions, AAPipeline)
+    ccall((:LLVMPassBuilderExtensionsSetAAPipeline, libLLVMExtra), Cvoid, (LLVMPassBuilderExtensionsRef, Cstring), Extensions, AAPipeline)
+end
+
 function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
     ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
 end

src/newpm.jl

@@ -3,7 +3,8 @@
 
 ## pass managers
 
-export NewPMModulePassManager, NewPMCGSCCPassManager, NewPMFunctionPassManager, NewPMLoopPassManager
+export NewPMModulePassManager, NewPMCGSCCPassManager, NewPMFunctionPassManager,
+       NewPMLoopPassManager, NewPMAAManager
 
 abstract type AbstractPassManager end
 
@@ -174,6 +175,7 @@ mutable struct NewPMPassBuilder <: AbstractPassManager
     opts::API.LLVMPassBuilderOptionsRef
     exts::API.LLVMPassBuilderExtensionsRef
     passes::Vector{String}
+    aa_passes::Vector{String}
     custom_passes::Vector{NewPMCustomPass}
 end
 
@@ -185,7 +187,7 @@ Base.unsafe_convert(::Type{API.LLVMPassBuilderOptionsRef}, pb::NewPMPassBuilder)
 function NewPMPassBuilder(; kwargs...)
     opts = API.LLVMCreatePassBuilderOptions()
     exts = API.LLVMCreatePassBuilderExtensions()
-    obj = mark_alloc(NewPMPassBuilder(opts, exts, [], []))
+    obj = mark_alloc(NewPMPassBuilder(opts, exts, [], [], []))
 
     for (name, value) in pairs(kwargs)
         if name == :verify_each
@@ -247,6 +249,7 @@ run!
 
 function run!(pb::NewPMPassBuilder, mod::Module, tm::Union{Nothing,TargetMachine}=nothing)
     isempty(pb.passes) && return
+    pipeline = join(pb.passes, ",")
 
     # XXX: The Base API is too restricted, not supporting custom passes
     #      or Julia's pass registration callback
@@ -273,7 +276,13 @@ function run!(pb::NewPMPassBuilder, mod::Module, tm::Union{Nothing,TargetMachine
         julia_callback = cglobal(:jl_register_passbuilder_callbacks)
         API.LLVMPassBuilderExtensionsSetRegistrationCallback(pb.exts, julia_callback)
 
-        @check API.LLVMRunJuliaPasses(mod, string(pb), something(tm, C_NULL),
+        # register AA pipeline
+        if !isempty(pb.aa_passes)
+            aa_pipeline = join(pb.aa_passes, ",")
+            API.LLVMPassBuilderExtensionsSetAAPipeline(pb.exts, aa_pipeline)
+        end
+
+        @check API.LLVMRunJuliaPasses(mod, pipeline, something(tm, C_NULL),
                                       pb.opts, pb.exts)
     end
 end
@@ -863,3 +872,28 @@ Base.string(options::LICMPassOptions) =
 @loop_pass "licm" LICMPass LICMPassOptions
 
 @loop_pass "lnicm" LNICMPass LICMPassOptions
+
+
+## alias analyses
+
+struct NewPMAAManager <: AbstractPassManager
+    passes::Vector{String}
+
+    NewPMAAManager() = new([])
+end
+
+Base.string(pb::NewPMAAManager) = join(pb.passes, ",")
+
+add!(pb::NewPMPassBuilder, aa::NewPMAAManager) = push!(pb.aa_passes, string(aa))
+add!(pm::NewPMAAManager, aa::NewPMAAManager) =
+    error("Alias analyses can only be added to the top-level pass builder")
+
+macro aa_pass(pass_name, class_name, params=nothing)
+    define_pass(pass_name, class_name, params)
+end
+
+@aa_pass "basic-aa" BasicAA
+@aa_pass "objc-arc-aa" ObjCARCAA
+@aa_pass "scev-aa" SCEVAA
+@aa_pass "scoped-noalias-aa" ScopedNoAliasAA
+@aa_pass "tbaa" TypeBasedAA

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
+IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/newpm_tests.jl

@@ -2,6 +2,7 @@
 @testitem "newpm" begin
 
 using LLVM.Interop
+using IOCapture
 
 function test_module()
     mod = LLVM.Module("test")
@@ -19,7 +20,7 @@ function test_module()
 end
 
 @testset "pass builder" begin
-    @dispose ctx=Context() pb=NewPMPassBuilder() begin
+    @dispose ctx=Context() begin
         # single pass
         @dispose mod=test_module() begin
             # by string
@@ -74,7 +75,7 @@ end
 end
 
 @testset "pass manager" begin
-    @dispose ctx=Context() pb=NewPMPassBuilder() begin
+    @dispose ctx=Context() begin
         # pass manager interface
         @dispose pb=NewPMPassBuilder() mod=test_module() begin
             add!(pb, "no-op-module")
@@ -364,4 +365,44 @@ end
     end
 end
 
+@testset "alias analyses" begin
+    # default pipeline
+    @dispose ctx=Context() pb=NewPMPassBuilder() begin
+        @dispose pb=NewPMPassBuilder(;debug_logging=true) mod=test_module() begin
+            add!(pb, "aa-eval")
+
+            io = IOCapture.capture() do
+                run!(pb, mod)
+            end
+
+            @test contains(io.output, "Running analysis: BasicAA")
+            @test contains(io.output, "Running analysis: TypeBasedAA")
+            @test contains(io.output, "Running analysis: ScopedNoAliasAA")
+        end
+    end
+
+    # custom pipeline
+    @dispose ctx=Context() pb=NewPMPassBuilder() begin
+        @dispose pb=NewPMPassBuilder(;debug_logging=true) mod=test_module() begin
+            add!(pb, NewPMAAManager()) do aa
+                # by string
+                add!(aa, "basic-aa")
+
+                # by object
+                add!(aa, SCEVAA())
+            end
+            add!(pb, "aa-eval")
+
+            io = IOCapture.capture() do
+                run!(pb, mod)
+            end
+
+            @test contains(io.output, "Running analysis: BasicAA")
+            @test contains(io.output, "Running analysis: SCEVAA")
+            @test !contains(io.output, "Running analysis: TypeBasedAA")
+            @test !contains(io.output, "Running analysis: ScopedNoAliasAA")
+        end
+    end
+end
+
 end