Support NewPM in some places (#2710)

Author: vchuravy

src/compiler.jl

@@ -1263,9 +1263,12 @@ function nested_codegen!(
     edges = edges::Vector{Any}
     push!(edges, funcspec)
 
-    LLVM.ModulePassManager() do pm
-        API.AddPreserveNVVMPass!(pm, true) #=Begin=#
-        LLVM.run!(pm, otherMod)
+    LLVM.@dispose pb=LLVM.NewPMPassBuilder() begin
+        registerEnzymeAndPassPipeline!(pb)
+        LLVM.add!(pb, LLVM.NewPMModulePassManager()) do mpm
+            LLVM.add!(mpm, PreserveNVVMPass())
+        end
+        LLVM.run!(pb, mod)
     end
 
     if DumpPreNestedCheck[]
@@ -2752,10 +2755,7 @@ function enzyme!(
     for f in collect(functions(mod))
         API.EnzymeFixupBatchedJuliaCallingConvention(f)
     end
-    ModulePassManager() do pm
-        dce!(pm)
-        LLVM.run!(pm, mod)
-    end
+    run!(DCEPass(), mod)
     fix_decayaddr!(mod)
     adjointf = adjointf == nothing ? nothing : functions(mod)[adjointfname]
     augmented_primalf =
@@ -4502,9 +4502,12 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT
         permit_inlining!(f)
     end
 
-    LLVM.ModulePassManager() do pm
-        API.AddPreserveNVVMPass!(pm, true) #=Begin=#
-        LLVM.run!(pm, mod)
+    LLVM.@dispose pb=LLVM.NewPMPassBuilder() begin
+        registerEnzymeAndPassPipeline!(pb)
+        LLVM.add!(pb, LLVM.NewPMModulePassManager()) do mpm
+            LLVM.add!(mpm, PreserveNVVMPass())
+        end
+        LLVM.run!(pb, mod)
     end
 
     primalf = meta.entry
@@ -5164,10 +5167,7 @@ end
                 push!(toremove, name(f))
             end
         end
-        ModulePassManager() do pm
-            always_inliner!(pm)
-            LLVM.run!(pm, mod)
-        end
+        run!(AlwaysInlinerPass(), mod)
         for fname in toremove
             if haskey(functions(mod), fname)
                 f = functions(mod)[fname]
@@ -5186,10 +5186,14 @@ end
         augmented_primalf = nothing
     end
 
-    LLVM.ModulePassManager() do pm
-        API.AddPreserveNVVMPass!(pm, false) #=Begin=#
-        LLVM.run!(pm, mod)
+    LLVM.@dispose pb=LLVM.NewPMPassBuilder() begin
+        registerEnzymeAndPassPipeline!(pb)
+        LLVM.add!(pb, LLVM.NewPMModulePassManager()) do mpm
+            LLVM.add!(mpm, PreserveNVVMEndPass())
+        end
+        LLVM.run!(pb, mod)
     end
+
     if !(primal_target isa GPUCompiler.NativeCompilerTarget)
         mark_gpu_intrinsics!(primal_target, mod)
     end

src/compiler/optimize.jl

@@ -1,73 +1,13 @@
-struct PipelineConfig
-    Speedup::Cint
-    Size::Cint
-    lower_intrinsics::Cint
-    dump_native::Cint
-    external_use::Cint
-    llvm_only::Cint
-    always_inline::Cint
-    enable_early_simplifications::Cint
-    enable_early_optimizations::Cint
-    enable_scalar_optimizations::Cint
-    enable_loop_optimizations::Cint
-    enable_vector_pipeline::Cint
-    remove_ni::Cint
-    cleanup::Cint
+function registerEnzymeAndPassPipeline!(pb::NewPMPassBuilder)
+    enzyme_callback = cglobal((:registerEnzymeAndPassPipeline, API.libEnzyme))
+    LLVM.API.LLVMPassBuilderExtensionsPushRegistrationCallbacks(pb.exts, enzyme_callback)
 end
 
-const RunAttributor = Ref(true)
-
-function pipeline_options(;
-    lower_intrinsics::Bool = true,
-    dump_native::Bool = false,
-    external_use::Bool = false,
-    llvm_only::Bool = false,
-    always_inline::Bool = true,
-    enable_early_simplifications::Bool = true,
-    enable_early_optimizations::Bool = true,
-    enable_scalar_optimizations::Bool = true,
-    enable_loop_optimizations::Bool = true,
-    enable_vector_pipeline::Bool = true,
-    remove_ni::Bool = true,
-    cleanup::Bool = true,
-    Size::Cint = 0,
-    Speedup::Cint = 3,
-)
-    return PipelineConfig(
-        Speedup,
-        Size,
-        lower_intrinsics,
-        dump_native,
-        external_use,
-        llvm_only,
-        always_inline,
-        enable_early_simplifications,
-        enable_early_optimizations,
-        enable_scalar_optimizations,
-        enable_loop_optimizations,
-        enable_vector_pipeline,
-        remove_ni,
-        cleanup,
-    )
-end
+LLVM.@function_pass "jl-inst-simplify" JLInstSimplifyPass
+LLVM.@module_pass "preserve-nvvm" PreserveNVVMPass
+LLVM.@module_pass "preserve-nvvm-end" PreserveNVVMEndPass
 
-function run_jl_pipeline(pm::ModulePassManager, tm::LLVM.TargetMachine; kwargs...)
-    config = Ref(pipeline_options(; kwargs...))
-    function jl_pipeline(m)
-        @dispose pb = NewPMPassBuilder() begin
-            add!(pb, NewPMModulePassManager()) do mpm
-                @ccall jl_build_newpm_pipeline(
-                    mpm.ref::Ptr{Cvoid},
-                    pb.ref::Ptr{Cvoid},
-                    config::Ptr{PipelineConfig},
-                )::Cvoid
-            end
-            LLVM.run!(mpm, m, tm)
-        end
-        return true
-    end
-    add!(pm, ModulePass("JLPipeline", jl_pipeline))
-end
+const RunAttributor = Ref(true)
 
 @static if VERSION < v"1.11.0-DEV.428"
 else
@@ -215,22 +155,7 @@ function loop_optimizations_tm!(pm::LLVM.ModulePassManager, tm::LLVM.TargetMachi
             loop_unswitch!(pm)
         end
     else
-        run_jl_pipeline(
-            pm,
-            tm;
-            lower_intrinsics = false,
-            dump_native = false,
-            external_use = false,
-            llvm_only = false,
-            always_inline = false,
-            enable_early_simplifications = false,
-            enable_early_optimizations = false,
-            enable_scalar_optimizations = false,
-            enable_loop_optimizations = true,
-            enable_vector_pipeline = false,
-            remove_ni = false,
-            cleanup = false,
-        )
+        @assert false
     end
 end
 
@@ -253,36 +178,7 @@ function more_loop_optimizations_tm!(pm::LLVM.ModulePassManager, tm::LLVM.Target
         loop_deletion!(pm)
         loop_unroll!(pm) # TODO: in Julia createSimpleLoopUnroll
     else
-        # LowerSIMDLoopPass
-        # LoopRotatePass [opt >= 2]
-        # LICMPass
-        # JuliaLICMPass
-        # SimpleLoopUnswitchPass
-        # LICMPass
-        # JuliaLICMPass
-        # IRCEPass
-        # LoopInstSimplifyPass
-        #   - in ours this is instcombine with jlinstsimplify
-        # LoopIdiomRecognizePass
-        # IndVarSimplifyPass
-        # LoopDeletionPass
-        # LoopFullUnrollPass
-        run_jl_pipeline(
-            pm,
-            tm;
-            lower_intrinsics = false,
-            dump_native = false,
-            external_use = false,
-            llvm_only = false,
-            always_inline = false,
-            enable_early_simplifications = false,
-            enable_early_optimizations = false,
-            enable_scalar_optimizations = false,
-            enable_loop_optimizations = true,
-            enable_vector_pipeline = false,
-            remove_ni = false,
-            cleanup = false,
-        )
+        @assert false
     end
 end
 

src/llvm/transforms.jl

@@ -2401,10 +2401,8 @@ function removeDeadArgs!(mod::LLVM.Module, tm::LLVM.TargetMachine)
     # and including 12 (but fixed 13+), Attributor will incorrectly change functions that
     # call code with undef to become unreachable, even when there exist other valid
     # callsites. See: https://godbolt.org/z/9Y3Gv6q5M
-    ModulePassManager() do pm
-        global_dce!(pm)
-        LLVM.run!(pm, mod)
-    end
+    run!(GlobalDCEPass(), mod)
+
     # Prevent dead-arg-elimination of functions which we may require args for in the derivative
     funcT = LLVM.FunctionType(LLVM.VoidType(), LLVMType[], vararg = true)
     if LLVM.version().major <= 15